import hashlib
import os

from .. import _tl
from .._crypto import factorization


def check_prime_and_good_check(prime: int, g: int):
    good_prime_bits_count = 2048
    if prime < 0 or prime.bit_length() != good_prime_bits_count:
        raise ValueError(
            "bad prime count {}, expected {}".format(
                prime.bit_length(), good_prime_bits_count
            )
        )

    # TODO This is awfully slow
    if factorization.Factorization.factorize(prime)[0] != 1:
        raise ValueError('given "prime" is not prime')

    if g == 2:
        if prime % 8 != 7:
            raise ValueError("bad g {}, mod8 {}".format(g, prime % 8))
    elif g == 3:
        if prime % 3 != 2:
            raise ValueError("bad g {}, mod3 {}".format(g, prime % 3))
    elif g == 4:
        pass
    elif g == 5:
        if prime % 5 not in (1, 4):
            raise ValueError("bad g {}, mod5 {}".format(g, prime % 5))
    elif g == 6:
        if prime % 24 not in (19, 23):
            raise ValueError("bad g {}, mod24 {}".format(g, prime % 24))
    elif g == 7:
        if prime % 7 not in (3, 5, 6):
            raise ValueError("bad g {}, mod7 {}".format(g, prime % 7))
    else:
        raise ValueError("bad g {}".format(g))

    prime_sub1_div2 = (prime - 1) // 2
    if factorization.Factorization.factorize(prime_sub1_div2)[0] != 1:
        raise ValueError("(prime - 1) // 2 is not prime")

    # Else it's good


def check_prime_and_good(prime_bytes: bytes, g: int):
    good_prime = bytes(
        (
            0xC7,
            0x1C,
            0xAE,
            0xB9,
            0xC6,
            0xB1,
            0xC9,
            0x04,
            0x8E,
            0x6C,
            0x52,
            0x2F,
            0x70,
            0xF1,
            0x3F,
            0x73,
            0x98,
            0x0D,
            0x40,
            0x23,
            0x8E,
            0x3E,
            0x21,
            0xC1,
            0x49,
            0x34,
            0xD0,
            0x37,
            0x56,
            0x3D,
            0x93,
            0x0F,
            0x48,
            0x19,
            0x8A,
            0x0A,
            0xA7,
            0xC1,
            0x40,
            0x58,
            0x22,
            0x94,
            0x93,
            0xD2,
            0x25,
            0x30,
            0xF4,
            0xDB,
            0xFA,
            0x33,
            0x6F,
            0x6E,
            0x0A,
            0xC9,
            0x25,
            0x13,
            0x95,
            0x43,
            0xAE,
            0xD4,
            0x4C,
            0xCE,
            0x7C,
            0x37,
            0x20,
            0xFD,
            0x51,
            0xF6,
            0x94,
            0x58,
            0x70,
            0x5A,
            0xC6,
            0x8C,
            0xD4,
            0xFE,
            0x6B,
            0x6B,
            0x13,
            0xAB,
            0xDC,
            0x97,
            0x46,
            0x51,
            0x29,
            0x69,
            0x32,
            0x84,
            0x54,
            0xF1,
            0x8F,
            0xAF,
            0x8C,
            0x59,
            0x5F,
            0x64,
            0x24,
            0x77,
            0xFE,
            0x96,
            0xBB,
            0x2A,
            0x94,
            0x1D,
            0x5B,
            0xCD,
            0x1D,
            0x4A,
            0xC8,
            0xCC,
            0x49,
            0x88,
            0x07,
            0x08,
            0xFA,
            0x9B,
            0x37,
            0x8E,
            0x3C,
            0x4F,
            0x3A,
            0x90,
            0x60,
            0xBE,
            0xE6,
            0x7C,
            0xF9,
            0xA4,
            0xA4,
            0xA6,
            0x95,
            0x81,
            0x10,
            0x51,
            0x90,
            0x7E,
            0x16,
            0x27,
            0x53,
            0xB5,
            0x6B,
            0x0F,
            0x6B,
            0x41,
            0x0D,
            0xBA,
            0x74,
            0xD8,
            0xA8,
            0x4B,
            0x2A,
            0x14,
            0xB3,
            0x14,
            0x4E,
            0x0E,
            0xF1,
            0x28,
            0x47,
            0x54,
            0xFD,
            0x17,
            0xED,
            0x95,
            0x0D,
            0x59,
            0x65,
            0xB4,
            0xB9,
            0xDD,
            0x46,
            0x58,
            0x2D,
            0xB1,
            0x17,
            0x8D,
            0x16,
            0x9C,
            0x6B,
            0xC4,
            0x65,
            0xB0,
            0xD6,
            0xFF,
            0x9C,
            0xA3,
            0x92,
            0x8F,
            0xEF,
            0x5B,
            0x9A,
            0xE4,
            0xE4,
            0x18,
            0xFC,
            0x15,
            0xE8,
            0x3E,
            0xBE,
            0xA0,
            0xF8,
            0x7F,
            0xA9,
            0xFF,
            0x5E,
            0xED,
            0x70,
            0x05,
            0x0D,
            0xED,
            0x28,
            0x49,
            0xF4,
            0x7B,
            0xF9,
            0x59,
            0xD9,
            0x56,
            0x85,
            0x0C,
            0xE9,
            0x29,
            0x85,
            0x1F,
            0x0D,
            0x81,
            0x15,
            0xF6,
            0x35,
            0xB1,
            0x05,
            0xEE,
            0x2E,
            0x4E,
            0x15,
            0xD0,
            0x4B,
            0x24,
            0x54,
            0xBF,
            0x6F,
            0x4F,
            0xAD,
            0xF0,
            0x34,
            0xB1,
            0x04,
            0x03,
            0x11,
            0x9C,
            0xD8,
            0xE3,
            0xB9,
            0x2F,
            0xCC,
            0x5B,
        )
    )

    if good_prime == prime_bytes:
        if g in (3, 4, 5, 7):
            return  # It's good

    check_prime_and_good_check(int.from_bytes(prime_bytes, "big"), g)


def is_good_large(number: int, p: int) -> bool:
    return number > 0 and p - number > 0


SIZE_FOR_HASH = 256


def num_bytes_for_hash(number: bytes) -> bytes:
    return bytes(SIZE_FOR_HASH - len(number)) + number


def big_num_for_hash(g: int) -> bytes:
    return g.to_bytes(SIZE_FOR_HASH, "big")


def sha256(*p: bytes) -> bytes:
    hash = hashlib.sha256()
    for q in p:
        hash.update(q)
    return hash.digest()


def is_good_mod_exp_first(modexp, prime) -> bool:
    diff = prime - modexp
    min_diff_bits_count = 2048 - 64
    max_mod_exp_size = 256
    if (
        diff < 0
        or diff.bit_length() < min_diff_bits_count
        or modexp.bit_length() < min_diff_bits_count
        or (modexp.bit_length() + 7) // 8 > max_mod_exp_size
    ):
        return False
    return True


def xor(a: bytes, b: bytes) -> bytes:
    return bytes(x ^ y for x, y in zip(a, b))


def pbkdf2sha512(password: bytes, salt: bytes, iterations: int):
    return hashlib.pbkdf2_hmac("sha512", password, salt, iterations)


def compute_hash(
    algo: _tl.PasswordKdfAlgoSHA256SHA256PBKDF2HMACSHA512iter100000SHA256ModPow,
    password: str,
):
    hash1 = sha256(algo.salt1, password.encode("utf-8"), algo.salt1)
    hash2 = sha256(algo.salt2, hash1, algo.salt2)
    hash3 = pbkdf2sha512(hash2, algo.salt1, 100000)
    return sha256(algo.salt2, hash3, algo.salt2)


def compute_digest(
    algo: _tl.PasswordKdfAlgoSHA256SHA256PBKDF2HMACSHA512iter100000SHA256ModPow,
    password: str,
):
    try:
        check_prime_and_good(algo.p, algo.g)
    except ValueError:
        raise ValueError("bad p/g in password")

    value = pow(
        algo.g,
        int.from_bytes(compute_hash(algo, password), "big"),
        int.from_bytes(algo.p, "big"),
    )

    return big_num_for_hash(value)


# https://github.com/telegramdesktop/tdesktop/blob/18b74b90451a7db2379a9d753c9cbaf8734b4d5d/Telegram/SourceFiles/core/core_cloud_password.cpp
def compute_check(request: _tl.account.Password, password: str):
    algo = request.current_algo
    if not isinstance(
        algo, _tl.PasswordKdfAlgoSHA256SHA256PBKDF2HMACSHA512iter100000SHA256ModPow
    ):
        raise ValueError(
            "unsupported password algorithm {}".format(algo.__class__.__name__)
        )

    pw_hash = compute_hash(algo, password)

    p = int.from_bytes(algo.p, "big")
    g = algo.g
    B = int.from_bytes(request.srp_B, "big")
    try:
        check_prime_and_good(algo.p, g)
    except ValueError:
        raise ValueError("bad p/g in password")

    if not is_good_large(B, p):
        raise ValueError("bad b in check")

    x = int.from_bytes(pw_hash, "big")
    p_for_hash = num_bytes_for_hash(algo.p)
    g_for_hash = big_num_for_hash(g)
    b_for_hash = num_bytes_for_hash(request.srp_B)
    g_x = pow(g, x, p)
    k = int.from_bytes(sha256(p_for_hash, g_for_hash), "big")
    kg_x = (k * g_x) % p

    def generate_and_check_random():
        random_size = 256
        while True:
            random = os.urandom(random_size)
            a = int.from_bytes(random, "big")
            A = pow(g, a, p)
            if is_good_mod_exp_first(A, p):
                a_for_hash = big_num_for_hash(A)
                u = int.from_bytes(sha256(a_for_hash, b_for_hash), "big")
                if u > 0:
                    return (a, a_for_hash, u)

    a, a_for_hash, u = generate_and_check_random()
    g_b = (B - kg_x) % p
    if not is_good_mod_exp_first(g_b, p):
        raise ValueError("bad g_b")

    ux = u * x
    a_ux = a + ux
    S = pow(g_b, a_ux, p)
    K = sha256(big_num_for_hash(S))
    M1 = sha256(
        xor(sha256(p_for_hash), sha256(g_for_hash)),
        sha256(algo.salt1),
        sha256(algo.salt2),
        a_for_hash,
        b_for_hash,
        K,
    )

    return _tl.InputCheckPasswordSRP(request.srp_id, bytes(a_for_hash), bytes(M1))

import hashlib
import os

from .. import _tl
from .._crypto import factorization


def check_prime_and_good_check(prime: int, g: int):
    good_prime_bits_count = 2048
    if prime < 0 or prime.bit_length() != good_prime_bits_count:
        raise ValueError(
            "bad prime count {}, expected {}".format(
                prime.bit_length(), good_prime_bits_count
            )
        )

    # TODO This is awfully slow
    if factorization.Factorization.factorize(prime)[0] != 1:
        raise ValueError('given "prime" is not prime')

    if g == 2:
        if prime % 8 != 7:
            raise ValueError("bad g {}, mod8 {}".format(g, prime % 8))
    elif g == 3:
        if prime % 3 != 2:
            raise ValueError("bad g {}, mod3 {}".format(g, prime % 3))
    elif g == 4:
        pass
    elif g == 5:
        if prime % 5 not in (1, 4):
            raise ValueError("bad g {}, mod5 {}".format(g, prime % 5))
    elif g == 6:
        if prime % 24 not in (19, 23):
            raise ValueError("bad g {}, mod24 {}".format(g, prime % 24))
    elif g == 7:
        if prime % 7 not in (3, 5, 6):
            raise ValueError("bad g {}, mod7 {}".format(g, prime % 7))
    else:
        raise ValueError("bad g {}".format(g))

    prime_sub1_div2 = (prime - 1) // 2
    if factorization.Factorization.factorize(prime_sub1_div2)[0] != 1:
        raise ValueError("(prime - 1) // 2 is not prime")

    # Else it's good


def check_prime_and_good(prime_bytes: bytes, g: int):
    good_prime = bytes(
        (
            0xC7,
            0x1C,
            0xAE,
            0xB9,
            0xC6,
            0xB1,
            0xC9,
            0x04,
            0x8E,
            0x6C,
            0x52,
            0x2F,
            0x70,
            0xF1,
            0x3F,
            0x73,
            0x98,
            0x0D,
            0x40,
            0x23,
            0x8E,
            0x3E,
            0x21,
            0xC1,
            0x49,
            0x34,
            0xD0,
            0x37,
            0x56,
            0x3D,
            0x93,
            0x0F,
            0x48,
            0x19,
            0x8A,
            0x0A,
            0xA7,
            0xC1,
            0x40,
            0x58,
            0x22,
            0x94,
            0x93,
            0xD2,
            0x25,
            0x30,
            0xF4,
            0xDB,
            0xFA,
            0x33,
            0x6F,
            0x6E,
            0x0A,
            0xC9,
            0x25,
            0x13,
            0x95,
            0x43,
            0xAE,
            0xD4,
            0x4C,
            0xCE,
            0x7C,
            0x37,
            0x20,
            0xFD,
            0x51,
            0xF6,
            0x94,
            0x58,
            0x70,
            0x5A,
            0xC6,
            0x8C,
            0xD4,
            0xFE,
            0x6B,
            0x6B,
            0x13,
            0xAB,
            0xDC,
            0x97,
            0x46,
            0x51,
            0x29,
            0x69,
            0x32,
            0x84,
            0x54,
            0xF1,
            0x8F,
            0xAF,
            0x8C,
            0x59,
            0x5F,
            0x64,
            0x24,
            0x77,
            0xFE,
            0x96,
            0xBB,
            0x2A,
            0x94,
            0x1D,
            0x5B,
            0xCD,
            0x1D,
            0x4A,
            0xC8,
            0xCC,
            0x49,
            0x88,
            0x07,
            0x08,
            0xFA,
            0x9B,
            0x37,
            0x8E,
            0x3C,
            0x4F,
            0x3A,
            0x90,
            0x60,
            0xBE,
            0xE6,
            0x7C,
            0xF9,
            0xA4,
            0xA4,
            0xA6,
            0x95,
            0x81,
            0x10,
            0x51,
            0x90,
            0x7E,
            0x16,
            0x27,
            0x53,
            0xB5,
            0x6B,
            0x0F,
            0x6B,
            0x41,
            0x0D,
            0xBA,
            0x74,
            0xD8,
            0xA8,
            0x4B,
            0x2A,
            0x14,
            0xB3,
            0x14,
            0x4E,
            0x0E,
            0xF1,
            0x28,
            0x47,
            0x54,
            0xFD,
            0x17,
            0xED,
            0x95,
            0x0D,
            0x59,
            0x65,
            0xB4,
            0xB9,
            0xDD,
            0x46,
            0x58,
            0x2D,
            0xB1,
            0x17,
            0x8D,
            0x16,
            0x9C,
            0x6B,
            0xC4,
            0x65,
            0xB0,
            0xD6,
            0xFF,
            0x9C,
            0xA3,
            0x92,
            0x8F,
            0xEF,
            0x5B,
            0x9A,
            0xE4,
            0xE4,
            0x18,
            0xFC,
            0x15,
            0xE8,
            0x3E,
            0xBE,
            0xA0,
            0xF8,
            0x7F,
            0xA9,
            0xFF,
            0x5E,
            0xED,
            0x70,
            0x05,
            0x0D,
            0xED,
            0x28,
            0x49,
            0xF4,
            0x7B,
            0xF9,
            0x59,
            0xD9,
            0x56,
            0x85,
            0x0C,
            0xE9,
            0x29,
            0x85,
            0x1F,
            0x0D,
            0x81,
            0x15,
            0xF6,
            0x35,
            0xB1,
            0x05,
            0xEE,
            0x2E,
            0x4E,
            0x15,
            0xD0,
            0x4B,
            0x24,
            0x54,
            0xBF,
            0x6F,
            0x4F,
            0xAD,
            0xF0,
            0x34,
            0xB1,
            0x04,
            0x03,
            0x11,
            0x9C,
            0xD8,
            0xE3,
            0xB9,
            0x2F,
            0xCC,
            0x5B,
        )
    )

    if good_prime == prime_bytes:
        if g in (3, 4, 5, 7):
            return  # It's good

    check_prime_and_good_check(int.from_bytes(prime_bytes, "big"), g)


def is_good_large(number: int, p: int) -> bool:
    return number > 0 and p - number > 0


SIZE_FOR_HASH = 256


def num_bytes_for_hash(number: bytes) -> bytes:
    return bytes(SIZE_FOR_HASH - len(number)) + number


def big_num_for_hash(g: int) -> bytes:
    return g.to_bytes(SIZE_FOR_HASH, "big")


def sha256(*p: bytes) -> bytes:
    hash = hashlib.sha256()
    for q in p:
        hash.update(q)
    return hash.digest()


def is_good_mod_exp_first(modexp, prime) -> bool:
    diff = prime - modexp
    min_diff_bits_count = 2048 - 64
    max_mod_exp_size = 256
    if (
        diff < 0
        or diff.bit_length() < min_diff_bits_count
        or modexp.bit_length() < min_diff_bits_count
        or (modexp.bit_length() + 7) // 8 > max_mod_exp_size
    ):
        return False
    return True


def xor(a: bytes, b: bytes) -> bytes:
    return bytes(x ^ y for x, y in zip(a, b))


def pbkdf2sha512(password: bytes, salt: bytes, iterations: int):
    return hashlib.pbkdf2_hmac("sha512", password, salt, iterations)


def compute_hash(
    algo: _tl.PasswordKdfAlgoSHA256SHA256PBKDF2HMACSHA512iter100000SHA256ModPow,
    password: str,
):
    hash1 = sha256(algo.salt1, password.encode("utf-8"), algo.salt1)
    hash2 = sha256(algo.salt2, hash1, algo.salt2)
    hash3 = pbkdf2sha512(hash2, algo.salt1, 100000)
    return sha256(algo.salt2, hash3, algo.salt2)


def compute_digest(
    algo: _tl.PasswordKdfAlgoSHA256SHA256PBKDF2HMACSHA512iter100000SHA256ModPow,
    password: str,
):
    try:
        check_prime_and_good(algo.p, algo.g)
    except ValueError:
        raise ValueError("bad p/g in password")

    value = pow(
        algo.g,
        int.from_bytes(compute_hash(algo, password), "big"),
        int.from_bytes(algo.p, "big"),
    )

    return big_num_for_hash(value)


# https://github.com/telegramdesktop/tdesktop/blob/18b74b90451a7db2379a9d753c9cbaf8734b4d5d/Telegram/SourceFiles/core/core_cloud_password.cpp
def compute_check(request: _tl.account.Password, password: str):
    algo = request.current_algo
    if not isinstance(
        algo, _tl.PasswordKdfAlgoSHA256SHA256PBKDF2HMACSHA512iter100000SHA256ModPow
    ):
        raise ValueError(
            "unsupported password algorithm {}".format(algo.__class__.__name__)
        )

    pw_hash = compute_hash(algo, password)

    p = int.from_bytes(algo.p, "big")
    g = algo.g
    B = int.from_bytes(request.srp_B, "big")
    try:
        check_prime_and_good(algo.p, g)
    except ValueError:
        raise ValueError("bad p/g in password")

    if not is_good_large(B, p):
        raise ValueError("bad b in check")

    x = int.from_bytes(pw_hash, "big")
    p_for_hash = num_bytes_for_hash(algo.p)
    g_for_hash = big_num_for_hash(g)
    b_for_hash = num_bytes_for_hash(request.srp_B)
    g_x = pow(g, x, p)
    k = int.from_bytes(sha256(p_for_hash, g_for_hash), "big")
    kg_x = (k * g_x) % p

    def generate_and_check_random():
        random_size = 256
        while True:
            random = os.urandom(random_size)
            a = int.from_bytes(random, "big")
            A = pow(g, a, p)
            if is_good_mod_exp_first(A, p):
                a_for_hash = big_num_for_hash(A)
                u = int.from_bytes(sha256(a_for_hash, b_for_hash), "big")
                if u > 0:
                    return (a, a_for_hash, u)

    a, a_for_hash, u = generate_and_check_random()
    g_b = (B - kg_x) % p
    if not is_good_mod_exp_first(g_b, p):
        raise ValueError("bad g_b")

    ux = u * x
    a_ux = a + ux
    S = pow(g_b, a_ux, p)
    K = sha256(big_num_for_hash(S))
    M1 = sha256(
        xor(sha256(p_for_hash), sha256(g_for_hash)),
        sha256(algo.salt1),
        sha256(algo.salt2),
        a_for_hash,
        b_for_hash,
        K,
    )

    return _tl.InputCheckPasswordSRP(request.srp_id, bytes(a_for_hash), bytes(M1))

# TODO This is awfully slow # Else it's good # It's good # https://github.com/telegramdesktop/tdesktop/blob/18b74b90451a7db2379a9d753c9cbaf8734b4d5d/Telegram/SourceFiles/core/core_cloud_password.cpp