"""Exact integer arithmetic: primes, divisibility, modular arithmetic, combinatorics, bases."""

from __future__ import annotations

from typing import Any, Union

import sympy as sp
from sympy.ntheory.modular import crt

from mcp_calculator.ops.common import MathError, check_list, clip, to_str
from mcp_calculator.ops.parsing import parse
from mcp_calculator.ops.spec import NumberTheoryOperation

IntLike = Union[int, str]

# operation -> (min values, max values, value names for error messages)
_ARITY: dict[str, tuple[int, int, str]] = {
    "gcd": (2, 10_000, "[a, b, ...]"),
    "lcm": (2, 10_000, "[a, b, ...]"),
    "factorize": (1, 1, "[n]"),
    "is_prime": (1, 1, "[n]"),
    "next_prime": (1, 1, "[n]"),
    "previous_prime": (1, 1, "[n]"),
    "nth_prime": (1, 1, "[k]"),
    "primes_in_range": (2, 2, "[low, high]"),
    "divisors": (1, 1, "[n]"),
    "totient": (1, 1, "[n]"),
    "mod": (2, 2, "[a, m]"),
    "mod_pow": (3, 3, "[base, exponent, modulus]"),
    "mod_inverse": (2, 2, "[a, m]"),
    "crt": (2, 10_000, "[r1, r2, ...] plus 'moduli'"),
    "factorial": (1, 1, "[n]"),
    "binomial": (2, 2, "[n, k]"),
    "permutations": (2, 2, "[n, k]"),
    "fibonacci": (1, 1, "[n]"),
    "continued_fraction": (1, 1, "[x] (rational or expression)"),
    "base_convert": (1, 1, "[value] plus 'from_base' / 'to_base'"),
}

_MAX_FACTORIAL = 20_000
_MAX_PRIMES = 10_000
_DIGITS = "0123456789abcdefghijklmnopqrstuvwxyz"


def number_theory(
    operation: NumberTheoryOperation,
    values: list[IntLike],
    moduli: list[IntLike] | None = None,
    from_base: int = 10,
    to_base: int = 10,
) -> dict[str, Any]:
    if operation not in _ARITY:
        raise MathError(f"unknown operation {operation!r}")
    lo, hi, names = _ARITY[operation]
    check_list(values, "values", min_items=lo, max_items=hi)
    if len(values) > hi or len(values) < lo:
        raise MathError(f"'{operation}' expects values = {names}")

    if operation == "base_convert":
        return _base_convert(str(values[0]), from_base, to_base)
    if operation == "continued_fraction":
        x = parse(str(values[0]))
        if not x.is_real:
            raise MathError("continued fractions need a real number")
        max_terms = 50 if x.is_Rational else 20
        terms = _continued_fraction(x, max_terms)
        out: dict[str, Any] = {"result": terms}
        if len(terms) == max_terms:
            out["note"] = f"first {max_terms} terms"
        return out

    n = [_int(v) for v in values]
    match operation:
        case "gcd":
            result: Any = sp.gcd_list(n)
        case "lcm":
            result = sp.lcm_list(n)
        case "factorize":
            _check_magnitude(n[0], 10**60, "factorize")
            factors = sp.factorint(n[0])
            pretty = " * ".join(f"{p}^{e}" if e > 1 else str(p) for p, e in factors.items()) or str(n[0])
            return {"result": {str(p): e for p, e in factors.items()}, "product_form": pretty}
        case "is_prime":
            result = bool(sp.isprime(n[0]))
        case "next_prime":
            result = sp.nextprime(n[0])
        case "previous_prime":
            if n[0] <= 2:
                raise MathError("there is no prime below 2")
            result = sp.prevprime(n[0])
        case "nth_prime":
            if not 1 <= n[0] <= 10**7:
                raise MathError("k must be between 1 and 10^7")
            result = sp.prime(n[0])
        case "primes_in_range":
            low, high = n
            primes = list(sp.primerange(low, high + 1)) if high - low <= 10**7 else None
            if primes is None:
                raise MathError("range too large (max width 10^7)")
            out = {"result": primes[:_MAX_PRIMES], "count": len(primes)}
            if len(primes) > _MAX_PRIMES:
                out["note"] = f"only the first {_MAX_PRIMES} are listed"
            return out
        case "divisors":
            _check_magnitude(n[0], 10**40, "divisors")
            divs = sp.divisors(n[0])
            return {"result": divs[:_MAX_PRIMES], "count": len(divs), "sum": sum(divs)}
        case "totient":
            _check_magnitude(n[0], 10**60, "totient")
            result = sp.totient(n[0])
        case "mod":
            if n[1] == 0:
                raise MathError("modulus must not be 0")
            result = n[0] % n[1]
        case "mod_pow":
            if n[2] == 0:
                raise MathError("modulus must not be 0")
            try:
                result = pow(n[0], n[1], n[2])
            except ValueError:
                raise MathError("a negative exponent needs base and modulus to be coprime") from None
        case "mod_inverse":
            try:
                result = sp.mod_inverse(n[0], n[1])
            except ValueError:
                raise MathError(f"{n[0]} has no inverse modulo {n[1]} (not coprime)") from None
        case "crt":
            if not moduli or len(moduli) != len(n):
                raise MathError("'crt' needs 'moduli' with one modulus per residue in 'values'")
            m = [_int(v) for v in moduli]
            solution = crt(m, n)
            if solution is None:
                return {"result": None, "note": "no solution: the congruences are inconsistent"}
            return {"result": int(solution[0]), "modulus": int(solution[1]), "meaning": f"x ≡ {solution[0]} (mod {solution[1]})"}
        case "factorial":
            if not 0 <= n[0] <= _MAX_FACTORIAL:
                raise MathError(f"n must be between 0 and {_MAX_FACTORIAL}")
            result = sp.factorial(n[0])
        case "binomial":
            result = sp.binomial(n[0], n[1])
        case "permutations":
            if n[1] > n[0] or n[1] < 0:
                result = 0
            else:
                result = sp.ff(n[0], n[1])
        case "fibonacci":
            if not 0 <= n[0] <= 100_000:
                raise MathError("n must be between 0 and 100000")
            result = sp.fibonacci(n[0])
        case _:
            raise MathError(f"unknown operation {operation!r}")

    return _int_out(result)


def _int_out(result: Any) -> dict[str, Any]:
    if isinstance(result, bool):
        return {"result": result}
    value = int(result)
    text = str(value)
    out: dict[str, Any] = {"result": clip(text) if len(text) > 15 else value}
    if len(text) > 15:
        out["digits"] = len(text.lstrip("-"))
        out["note"] = "integer given as string to keep every digit"
    return out


def _int(value: IntLike) -> int:
    if isinstance(value, bool):
        raise MathError("booleans are not integers here")
    if isinstance(value, int):
        return value
    if isinstance(value, float):
        if not value.is_integer():
            raise MathError(f"{value} is not an integer")
        return int(value)
    parsed = parse(str(value))
    if not parsed.is_Integer:
        raise MathError(f"{value!r} is not an integer")
    return int(parsed)


def _check_magnitude(n: int, bound: int, operation: str) -> None:
    if abs(n) > bound:
        raise MathError(f"'{operation}' is limited to |n| <= 10^{len(str(bound)) - 1}")


def _continued_fraction(x: sp.Expr, max_terms: int) -> list[int]:
    terms = []
    for term in sp.continued_fraction_iterator(x):
        terms.append(int(term))
        if len(terms) >= max_terms:
            break
    return terms


def _base_convert(value: str, from_base: int, to_base: int) -> dict[str, Any]:
    if not (2 <= from_base <= 36 and 2 <= to_base <= 36):
        raise MathError("bases must be between 2 and 36")
    text = value.strip().lower().replace("_", "")
    for prefix, base in (("0x", 16), ("0b", 2), ("0o", 8)):
        if text.lstrip("-").startswith(prefix) and from_base == base:
            text = text.replace(prefix, "", 1)
    if len(text) > 1000:
        raise MathError("value too long")
    try:
        n = int(text, from_base)
    except ValueError:
        raise MathError(f"{value!r} is not a valid base-{from_base} number") from None
    if n == 0:
        digits = "0"
    else:
        digits, m = "", abs(n)
        while m:
            m, r = divmod(m, to_base)
            digits = _DIGITS[r] + digits
        digits = ("-" if n < 0 else "") + digits
    return {"result": digits, "decimal": to_str(n), "from_base": from_base, "to_base": to_base}