from __future__ import annotations
from dataclasses import dataclass
import logging
from typing import *
import gmpy2
from . import Challenge, Serde
_logger = logging.getLogger(__name__)

@dataclass
class Problem(Serde):
    difficulty: int
    num: int
    num_bits: int

@dataclass
class Solution(Serde):
    status: str
    p: Optional[int]
    q: Optional[int]

@dataclass
class Verif(Serde):
    n: int
    p: int
    q: int

def validate_breaking_rsa_solution(solution: Solution, verif: Verif, prob: Problem | None=None, require_success_status: bool=True) -> tuple[bool, str | None]:
    if require_success_status and getattr(solution, 'status', None) != 'success':
        return (False, f"Solution status is '{getattr(solution, 'status', None)}', not 'success'")
    if solution.p is None or solution.q is None:
        return (False, 'Solution has missing factors (p or q is null)')
    try:
        sol_p = int(solution.p)
        sol_q = int(solution.q)
    except (TypeError, ValueError):
        return (False, f'Solution p/q are not valid integers: p={solution.p}, q={solution.q}')
    sol_p, sol_q = (min(sol_p, sol_q), max(sol_p, sol_q))
    p_check, q_check = (min(verif.p, verif.q), max(verif.p, verif.q))
    n = verif.n
    if sol_p * sol_q != n:
        return (False, f'p * q != n: {sol_p} * {sol_q} != {n}')
    if sol_p != p_check or sol_q != q_check:
        return (False, f"Factors don't match expected: got ({sol_p}, {sol_q}), expected ({p_check}, {q_check})")
    if prob is not None:
        if prob.num != n or sol_p * sol_q != prob.num:
            return (False, f'n mismatch between problem ({prob.num}) and verif ({n})')
    return (True, None)

def _gen_prime(num_bits: int, rng: gmpy2.random_state) -> gmpy2.mpz:
    n = gmpy2.mpz_urandomb(rng, num_bits)
    n |= 1 << num_bits - 1 | 1
    return gmpy2.next_prime(n)

@dataclass
class BreakingRSA(Challenge[Problem, Solution, Verif]):
    difficulty: int
    num_bits: int

    def generate(self, seed: int) -> tuple[Problem, Verif]:
        rng = gmpy2.random_state(seed)
        bits_p = self.num_bits // 2
        bits_q = self.num_bits - bits_p
        fermat_thresh = 2 ** max(bits_p - 100, 1)
        while True:
            p = _gen_prime(bits_p, rng)
            q = _gen_prime(bits_q, rng)
            n = p * q
            if n.bit_length() == self.num_bits and abs(p - q) > fermat_thresh:
                break
        n = int(n)
        p = int(p)
        q = int(q)
        _logger.info(f'Generated {self.num_bits}-bit semiprime ({len(str(n))} digits)')
        problem = Problem(self.difficulty, n, self.num_bits)
        verif = Verif(n, p, q)
        return (problem, verif)

    def verify(self, prob: Problem, sol: Solution, verif: Verif) -> bool:
        success, reason = validate_breaking_rsa_solution(sol, verif, prob, require_success_status=False)
        if success:
            _logger.info('Verification SUCCESS')
        elif sol.p is None or sol.q is None:
            _logger.info(f'Verification FAILURE: missing factors, got p={sol.p} and' + f' q={sol.q}, with solution status {sol.status}')
        else:
            _logger.info(f'Verification FAILURE: {reason or 'factors do not match'}')
        return success