from __future__ import annotations

import json
import logging
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Callable, Optional


VerifyFn = Callable[[bytes, bytes, bytes], bool]
SignFn = Callable[[bytes, bytes], bytes]

log = logging.getLogger(__name__)

_ALLOWED_OVERRIDE_WARNED = False
_PQ_ALLOWLIST_WARNED = False

_SCHEME_NAME_TO_ID = {
    "dilithium3": 1,
    "pq-1952-byte": 1,
    "sphincs_shake_128s": 2,
    "sphincs128s": 2,
    "sphincs+128s": 2,
    # v2 — real FIPS 204 ML-DSA-65
    "ml_dsa_65": 11,
    "ml-dsa-65": 11,
    "ml-dsa": 11,
}

_CHAIN_REQUIRED_SCHEMES: dict[int, tuple[int, ...]] = {1: (1, 2)}


def required_schemes_for_chain(chain_id: int) -> tuple[int, ...]:
    return _CHAIN_REQUIRED_SCHEMES.get(chain_id, tuple())


@dataclass(frozen=True)
class SchemeSpec:
    scheme_id: int
    name: str
    pubkey_lengths: tuple[int, ...]
    signature_lengths: tuple[int, ...]
    enabled_by_default: bool = True


@dataclass
class RuntimeScheme:
    spec: SchemeSpec
    sign_fn: Optional[SignFn] = None
    verify_fn: Optional[VerifyFn] = None
    enabled: bool = False
    reason_if_disabled: Optional[str] = None


@dataclass(frozen=True)
class PolicyOverride:
    enabled: bool
    mode: str
    allow_sig_schemes: tuple[int, ...]
    deny_sig_schemes: tuple[int, ...]
    comment: str
    file_path: str | None


@dataclass(frozen=True)
class PolicyEvaluation:
    scheme_id: int
    name: str
    enabled_by_code: bool
    enabled_by_policy: bool
    enabled_effective: bool
    reason_if_disabled: str | None
    policy_root: str


CANONICAL_SCHEME_SPECS: tuple[SchemeSpec, ...] = (
    SchemeSpec(
        scheme_id=1,
        name="dilithium3",
        pubkey_lengths=(1952,),
        signature_lengths=(3293,),
        enabled_by_default=True,
    ),
    SchemeSpec(
        scheme_id=2,
        name="sphincs_shake_128s",
        # NIST-spec pubkey is 32 bytes, but Animica's pure-Python
        # fallback (pq/py/algs/pure_python_fallbacks.py:fallback_sig_keypair)
        # derives a 64-byte pubkey as `_h("pk", sk, out_len=64)`. Every
        # wallet ever generated by the CLI carries the 64-byte form;
        # accepting only 32 would reject all live SPHINCS keys, leaving
        # the chain unable to verify any tx signed by them. Both lengths
        # are valid in this network.
        pubkey_lengths=(32, 64),
        signature_lengths=(7856,),
        enabled_by_default=True,
    ),
    SchemeSpec(
        scheme_id=3,
        name="sphincs_shake_128f",
        pubkey_lengths=(32,),
        signature_lengths=(17088,),
        enabled_by_default=True,
    ),
    SchemeSpec(
        scheme_id=4,
        name="sphincs_shake_256s",
        pubkey_lengths=(64,),
        signature_lengths=(29792,),
        enabled_by_default=True,
    ),
    # ── v2 schemes (real PQ crypto) ─────────────────────────────────────
    # 2026-05-30: scheme_id=1 ("dilithium3") and scheme_id=2 ("sphincs_
    # shake_128s") above are commitment-style stubs that don't actually
    # implement ML-DSA-65 / SPHINCS+. Any party with the public key can
    # forge signatures under those schemes (see incident notes). The
    # following scheme_ids are the real FIPS 204 implementations,
    # vendored from python/animica/_vendor/dilithium_py_v2/ (MIT,
    # jack4818/dilithium-py).
    #
    # Migration plan:
    #   - Mainnet accepts BOTH old and new schemes during a transition
    #     window so historical txs keep verifying.
    #   - New wallets default to ml_dsa_65 (the CLI bumps in v0.2.0).
    #   - At a fork height (configurable via env), schemes 1 and 2 are
    #     dropped from the accepted-policy: tx.sendRawTransaction
    #     rejects them while block import still accepts them in
    #     historical blocks (verify-only, no new ingestion).
    #
    # ML-DSA-65 sizes per FIPS 204:
    #   pk = 1952 bytes, sk = 4032 bytes, sig = 3309 bytes
    SchemeSpec(
        scheme_id=11,
        name="ml_dsa_65",
        pubkey_lengths=(1952,),
        signature_lengths=(3309,),
        enabled_by_default=True,
    ),
)


def load_policy_disabled_scheme_ids() -> set[int]:
    raw = (os.environ.get("ANIMICA_DISABLED_SIGNATURE_SCHEMES") or "").strip()
    if not raw:
        return set()
    return set(parse_scheme_ids_strict(raw, env_name="ANIMICA_DISABLED_SIGNATURE_SCHEMES"))


def parse_scheme_ids_strict(raw: str, *, env_name: str) -> tuple[int, ...]:
    out: set[int] = set()
    invalid: list[str] = []
    for part in raw.split(","):
        token = part.strip()
        if not token:
            continue
        token_l = token.lower()
        if token_l in _SCHEME_NAME_TO_ID:
            out.add(_SCHEME_NAME_TO_ID[token_l])
            continue
        try:
            out.add(int(token, 0))
        except ValueError:
            invalid.append(token)
    if invalid:
        raise ValueError(f"{env_name} contains invalid scheme identifiers: {invalid}")
    return tuple(sorted(out))




def load_allowed_signature_schemes_override() -> tuple[int, ...]:
    raw = (os.environ.get("ANIMICA_ALLOWED_SIG_SCHEMES") or "").strip()
    if not raw:
        return tuple()
    allowed = parse_scheme_ids_strict(raw, env_name="ANIMICA_ALLOWED_SIG_SCHEMES")
    global _ALLOWED_OVERRIDE_WARNED
    if allowed and not _ALLOWED_OVERRIDE_WARNED:
        log.warning("ANIMICA_ALLOWED_SIG_SCHEMES override active: %s", list(allowed))
        _ALLOWED_OVERRIDE_WARNED = True
    return allowed


def load_pq_allowlist_fallback() -> tuple[int, ...]:
    """
    Parse ANIMICA_PQ_ALLOWLIST into scheme IDs.

    This fallback must be explicitly enabled via ANIMICA_ENABLE_PQ_ALLOWLIST_FALLBACK=1
    and is only applied when policy load fails.
    """
    enabled = (os.environ.get("ANIMICA_ENABLE_PQ_ALLOWLIST_FALLBACK") or "").strip() == "1"
    if not enabled:
        return tuple()

    raw = (os.environ.get("ANIMICA_PQ_ALLOWLIST") or "").strip().lower()
    if not raw:
        return tuple()

    allowed = parse_scheme_ids_strict(raw, env_name="ANIMICA_PQ_ALLOWLIST")
    global _PQ_ALLOWLIST_WARNED
    if allowed and not _PQ_ALLOWLIST_WARNED:
        log.warning("ANIMICA_PQ_ALLOWLIST fallback active: %s", list(allowed))
        _PQ_ALLOWLIST_WARNED = True
    return allowed

def _parse_int_list(value: object) -> tuple[int, ...]:
    if not isinstance(value, list):
        return tuple()
    out: list[int] = []
    for item in value:
        if isinstance(item, bool):
            continue
        try:
            out.append(int(item))
        except Exception:
            continue
    return tuple(sorted(set(out)))


def load_policy_override() -> PolicyOverride:
    enabled = (os.environ.get("ANIMICA_ENABLE_POLICY_OVERRIDE") or "").strip() == "1"
    file_path = (os.environ.get("ANIMICA_POLICY_OVERRIDE_FILE") or "").strip() or None
    if not enabled or not file_path:
        return PolicyOverride(
            enabled=False,
            mode="enforce",
            allow_sig_schemes=tuple(),
            deny_sig_schemes=tuple(),
            comment="",
            file_path=file_path,
        )

    p = Path(file_path)
    if not p.exists():
        log.warning("Policy override enabled but file is missing: %s", file_path)
        return PolicyOverride(
            enabled=False,
            mode="enforce",
            allow_sig_schemes=tuple(),
            deny_sig_schemes=tuple(),
            comment="override_file_missing",
            file_path=file_path,
        )

    try:
        payload = json.loads(p.read_text(encoding="utf-8"))
    except Exception as exc:
        log.warning("Policy override enabled but file could not be parsed (%s): %s", file_path, exc)
        return PolicyOverride(
            enabled=False,
            mode="enforce",
            allow_sig_schemes=tuple(),
            deny_sig_schemes=tuple(),
            comment="override_file_parse_error",
            file_path=file_path,
        )

    mode = str(payload.get("mode") or "enforce").strip().lower()
    if mode not in {"enforce", "override_allow"}:
        mode = "enforce"
    allow = _parse_int_list(payload.get("allowSigSchemes", []))
    deny = _parse_int_list(payload.get("denySigSchemes", []))
    comment = str(payload.get("comment") or "")

    log.warning("POLICY OVERRIDE ACTIVE: mode=%s file=%s allow=%s deny=%s", mode, file_path, list(allow), list(deny))
    if comment:
        log.warning("POLICY OVERRIDE COMMENT: %s", comment)

    return PolicyOverride(
        enabled=True,
        mode=mode,
        allow_sig_schemes=allow,
        deny_sig_schemes=deny,
        comment=comment,
        file_path=file_path,
    )


def policy_roots() -> dict[str, str]:
    return {
        "pqAlgPolicy": "ANIMICA_DISABLED_SIGNATURE_SCHEMES",
        "allowedSigSchemesOverride": "ANIMICA_ALLOWED_SIG_SCHEMES",
        "policyOverride": "ANIMICA_POLICY_OVERRIDE_FILE",
        "pqAllowlistFallback": "ANIMICA_PQ_ALLOWLIST",
    }


def evaluate_scheme_policy(
    spec: SchemeSpec,
    *,
    disabled_by_policy: set[int],
    override: PolicyOverride,
) -> PolicyEvaluation:
    enabled_by_code = bool(spec.enabled_by_default)

    base_enabled_by_policy = spec.scheme_id not in disabled_by_policy
    reason: str | None = None

    allowed_override = load_allowed_signature_schemes_override()
    allowlist_fallback = load_pq_allowlist_fallback()
    policy_load_failed = override.comment in {"override_file_missing", "override_file_parse_error"}
    chain_id_env = os.environ.get("ANIMICA_CHAIN_ID")
    required_for_chain: set[int] = set()
    if chain_id_env:
        try:
            required_for_chain = set(required_schemes_for_chain(int(chain_id_env, 0)))
        except Exception:
            required_for_chain = set()

    if allowed_override:
        if spec.scheme_id in required_for_chain and spec.scheme_id not in allowed_override:
            raise ValueError(
                "ANIMICA_ALLOWED_SIG_SCHEMES contradicts chain-required schemes: "
                f"missing required scheme id {spec.scheme_id}"
            )
        enabled_by_policy = spec.scheme_id in allowed_override
        policy_root = "ANIMICA_ALLOWED_SIG_SCHEMES"
        if not enabled_by_policy:
            reason = "not_in_allowed_override"
    elif policy_load_failed and allowlist_fallback:
        enabled_by_policy = spec.scheme_id in allowlist_fallback
        policy_root = "ANIMICA_PQ_ALLOWLIST"
        if not enabled_by_policy:
            reason = "not_in_pq_allowlist_fallback"
    elif spec.scheme_id in override.deny_sig_schemes:
        if spec.scheme_id in required_for_chain:
            raise ValueError(
                "ANIMICA_POLICY_OVERRIDE_FILE denySigSchemes contradicts chain-required schemes: "
                f"contains required scheme id {spec.scheme_id}"
            )
        enabled_by_policy = False
        reason = "denied_by_override"
        policy_root = "ANIMICA_POLICY_OVERRIDE_FILE"
    elif override.enabled and override.mode == "override_allow" and spec.scheme_id in override.allow_sig_schemes:
        enabled_by_policy = True
        policy_root = "ANIMICA_POLICY_OVERRIDE_FILE"
    else:
        if spec.scheme_id in required_for_chain and spec.scheme_id in disabled_by_policy:
            raise ValueError(
                "ANIMICA_DISABLED_SIGNATURE_SCHEMES contradicts chain-required schemes: "
                f"contains required scheme id {spec.scheme_id}"
            )
        enabled_by_policy = base_enabled_by_policy
        policy_root = "ANIMICA_DISABLED_SIGNATURE_SCHEMES"

    enabled_effective = enabled_by_code and enabled_by_policy

    if reason is None and not enabled_by_policy:
        reason = "disabled_by_policy"
    if reason is None and not enabled_by_code:
        reason = "disabled_by_code"

    return PolicyEvaluation(
        scheme_id=spec.scheme_id,
        name=spec.name,
        enabled_by_code=enabled_by_code,
        enabled_by_policy=enabled_by_policy,
        enabled_effective=enabled_effective,
        reason_if_disabled=reason,
        policy_root=policy_root,
    )


def build_runtime_scheme_table() -> dict[int, RuntimeScheme]:
    return {spec.scheme_id: RuntimeScheme(spec=spec) for spec in CANONICAL_SCHEME_SPECS}
