"""Core ADS abstractions: Auth, Mode, ProofStream, auth/unauth primitives.

Ports ADS4All's Core.hs (GADT Auth) and Monad.hs (ADS monad) to Python
using frozen dataclasses and context managers.
"""

from __future__ import annotations

import hmac
from collections import deque
from contextlib import contextmanager
from dataclasses import dataclass
from enum import Enum, auto
from typing import Generic, TypeVar

from .errors import ADSError, ProofError, VerificationError
from .hash import shallow_hash

T = TypeVar("T")


class Mode(Enum):
    """Operating mode: Prover has values, Verifier has only hashes."""

    PROVER = auto()
    VERIFIER = auto()


@dataclass(frozen=True, slots=True)
class Auth(Generic[T]):
    """Authenticated wrapper. Prover holds hash+value, Verifier holds hash only.

    Ports Haskell's GADT: AuthP Hash a | AuthV Hash
    """

    _hash: bytes
    _value: T | None = None

    @staticmethod
    def prover(h: bytes, value: T) -> Auth[T]:
        """Create a prover-mode Auth (hash + value)."""
        return Auth(_hash=h, _value=value)

    @staticmethod
    def verifier(h: bytes) -> Auth[T]:
        """Create a verifier-mode Auth (hash only)."""
        return Auth(_hash=h, _value=None)

    def digest(self) -> bytes:
        """Return the hash."""
        return self._hash

    def unwrap(self) -> T:
        """Extract the value. Raises ADSError if in verifier mode."""
        if self._value is None:
            raise ADSError("Cannot unwrap Auth in verifier mode")
        return self._value

    @property
    def is_prover(self) -> bool:
        """True if this Auth carries a value (prover mode)."""
        return self._value is not None


class ProofStream:
    """Manages proof hash accumulation (prover) and consumption (verifier).

    Ports Haskell's ADS monad (ExceptT String (StateT ProofObject IO)).
    """

    __slots__ = ("mode", "_stream", "_finalized")

    def __init__(self, mode: Mode, proof: list[bytes] | None = None) -> None:
        self.mode = mode
        self._stream: deque[bytes] = deque(proof or [])
        self._finalized = False

    def append(self, h: bytes) -> None:
        """Prover appends a hash to the proof stream."""
        if self._finalized:
            raise ADSError("Cannot append to finalized proof stream")
        if self.mode != Mode.PROVER:
            raise ADSError("Only prover can append to proof stream")
        assert len(h) == 32, f"Hash must be 32 bytes, got {len(h)}"
        self._stream.append(h)

    def consume(self) -> bytes:
        """Verifier consumes (pops) the next hash from the proof stream."""
        if self._finalized:
            raise ADSError("Cannot consume from finalized proof stream")
        if self.mode != Mode.VERIFIER:
            raise ADSError("Only verifier can consume from proof stream")
        if not self._stream:
            raise ProofError("Proof stream exhausted")
        return self._stream.popleft()

    def get_proof(self) -> list[bytes]:
        """Return the accumulated proof (prover mode)."""
        return list(self._stream)

    @property
    def is_complete(self) -> bool:
        """True if all proof elements have been consumed."""
        return len(self._stream) == 0

    def _finalize(self) -> None:
        """Mark as finalized -- no further append/consume allowed."""
        self._finalized = True

    def __len__(self) -> int:
        return len(self._stream)


def auth(value: T, ctx: ProofStream) -> Auth[T]:
    """Wrap a value in an Auth container.

    Prover: stores hash + value.
    Verifier: stores hash only.
    """
    h = shallow_hash(value)
    if ctx.mode == Mode.PROVER:
        return Auth.prover(h, value)
    else:
        return Auth.verifier(h)


def unauth(authenticated: Auth[T], ctx: ProofStream) -> T | None:
    """Unwrap an authenticated value, managing the proof stream.

    Prover: extracts value, appends hash to proof stream.
    Verifier: consumes hash from proof stream, verifies match.
    """
    if ctx.mode == Mode.PROVER:
        ctx.append(authenticated.digest())
        return authenticated.unwrap()
    else:
        proof_hash = ctx.consume()
        if not hmac.compare_digest(proof_hash, authenticated.digest()):
            raise VerificationError(
                f"Hash mismatch: expected {authenticated.digest().hex()[:16]}..., "
                f"got {proof_hash.hex()[:16]}..."
            )
        return None


@contextmanager
def prover_context():
    """Run authenticated operations in prover mode.

    Usage:
        with prover_context() as ctx:
            auth_tree = auth(tree, ctx)
        proof = ctx.get_proof()
    """
    ctx = ProofStream(Mode.PROVER)
    yield ctx
    ctx._finalize()


@contextmanager
def verifier_context(proof: list[bytes]):
    """Run authenticated operations in verifier mode.

    Checks that all proof elements are consumed on exit.
    """
    ctx = ProofStream(Mode.VERIFIER, proof)
    yield ctx
    if not ctx.is_complete:
        remaining = len(ctx)
        ctx._finalize()
        raise VerificationError(
            f"Proof verification failed: {remaining} unused proof elements"
        )
    ctx._finalize()


def run_prover(operation, *args, **kwargs):
    """Run an operation in prover mode, returning (result, proof)."""
    ctx = ProofStream(Mode.PROVER)
    result = operation(*args, ctx=ctx, **kwargs)
    ctx._finalize()
    return result, ctx.get_proof()


def run_verifier(proof: list[bytes], operation, *args, **kwargs):
    """Run an operation in verifier mode with a proof.

    Raises VerificationError if proof is invalid or incomplete.
    """
    ctx = ProofStream(Mode.VERIFIER, proof)
    result = operation(*args, ctx=ctx, **kwargs)
    if not ctx.is_complete:
        remaining = len(ctx)
        ctx._finalize()
        raise VerificationError(
            f"Proof verification failed: {remaining} unused proof elements"
        )
    ctx._finalize()
    return result
