"""The quantumpipe protocol state machine, sans I/O. session()
initiates ("alice"), session(msg1) responds ("bob"); results are bytes
to send or signal objects; once established, session(Message) encrypts a
frame and session(frame) decrypts one back. See docs/design.md."""

from __future__ import annotations

import logging
import struct
from collections.abc import MutableMapping
from dataclasses import dataclass

from .chab3 import Receiver, Sender, random
from .exceptions import ProtocolError, _Burn
from .handshake import Registration, Resumption, open_hello
from .kdf import kdf

_log = logging.getLogger(__name__)


# one-line, secret-free rendering of call inputs/results for the log
def _describe(obj) -> str:
    if obj is None:
        return "-"
    if isinstance(obj, (bytes, bytearray)):
        return f"frame[{len(obj)}]"
    if isinstance(obj, Message):
        return f"Message(kind={obj.kind}, {len(obj.data)}B)"
    if isinstance(obj, Verification):
        return f"Verification({obj}, answer={obj.answer})"
    return type(obj).__name__


# --------------------------------------------------------------------------
# credentials
#
# generate_registration_secret() creates the one-time registration secret,
# given to both parties out of band and stored as a bare 32-byte store
# entry. The first successful registration replaces it with the composite.

_CREDENTIAL = struct.Struct("32s32s1184s32s4s?")  # idroot|shared|peer_m|peer_x|sas|vf


def generate_registration_secret() -> bytes:
    """Fresh 256-bit registration secret, given to both parties once."""
    return random(32)


def credential_id(cred: bytes) -> bytes:
    """A stable 16-byte identifier for a stored credential (recomputable
    from the record); rotates when the credential does."""
    if len(cred) == _CREDENTIAL.size:
        cred = cred[:-1]  # exclude the verified flag
    return kdf("cred/id", cred, length=16)


def _credential_shared(cred: bytes) -> bytes:
    # the bare secret itself, or the composite's rotated field
    return cred if len(cred) == 32 else _CREDENTIAL.unpack(cred)[1]


def _same_credential(a: bytes, b: bytes) -> bool:
    # same credential, ignoring the verified flag byte on composites
    if len(a) != len(b):
        return False
    if len(a) == _CREDENTIAL.size:
        return a[:-1] == b[:-1]
    return a == b


# --------------------------------------------------------------------------
# signals and traffic


@dataclass
class Verification:
    """Out-of-band comparison request (registration only): fill .answer
    and pass back -- True = verified, None = proceed unverified, False =
    reject (the credential is deleted)."""

    sas: bytes
    answer: bool | None = None

    def __str__(self) -> str:
        return self.sas.hex()


@dataclass(frozen=True)
class Established:
    """Handshake complete: sas = the credential's comparison digest,
    verified = its flag at establishment."""

    sas: bytes
    verified: bool


@dataclass(frozen=True)
class Message:
    """One traffic frame: an opaque user kind plus bytes. Kind 0 is
    reserved for in-band protocol messages."""

    data: bytes
    kind: int


# dispatch in/out types
_In = bytes | Message | Verification | None
_Out = bytes | Message | Verification | Established


# --------------------------------------------------------------------------
# the state machine


class Session:
    """One end of the protocol state machine; call it to advance.

    store maps peer ids to credentials; key_id selects the credential to
    initiate with (none -> respond); ad is application session context,
    bound into msg1 (both sides must pass the same value).
    """

    def __init__(self, store: MutableMapping, key_id=None, /, ad: bytes = b""):
        self._store = store
        self._key_id = None
        self._ad = ad
        self._registering: bool | None = None  # unknown until a credential loads
        self._exchange: Registration | Resumption | None = None
        self._initiator: bool | None = None  # unknown until the first call
        self._pending: bytes | None = None  # a tail message held for the gate
        self._verified_answer = False
        self._phase = "hello"
        if key_id is not None:
            self._load_credential(key_id)

    def __call__(self, msg: _In = None, /) -> list[_Out]:
        _log.debug("call state=%s input=%s", self.state, _describe(msg))
        try:
            if isinstance(msg, Message):  # outbound user traffic
                out = self._st_data(msg)
            else:
                match self._phase:
                    case "hello":
                        out = self._st_hello(msg)
                    case "msg2":
                        out = self._st_msg2(msg)
                    case "msg3":
                        out = self._st_msg3(msg)
                    case "msg4":
                        out = self._st_msg4(msg)
                    case "msg5":
                        out = self._st_msg5(msg)
                    case "verify":
                        out = self._st_verify(msg)
                    case "established":  # inbound user traffic
                        out = self._st_data(msg)
                    case _:
                        raise ProtocolError(f"unexpected input in state {self.state}")
        except _Burn as b:
            # invalid data inside a validly opened box: the sender knew
            # the keys and misbehaved -- retire the credential
            self._delete_credential()
            _log.debug("fail state=%s: credential burned: %s", self.state, b)
            raise ProtocolError(f"credential burned: {b}") from None
        except (ValueError, struct.error) as e:
            _log.debug("fail state=%s: %r", self.state, e)
            raise ProtocolError(str(e)) from None
        _log.debug(
            "emit state=%s results=[%s]",
            self.state,
            ", ".join(_describe(r) for r in out),
        )
        return out

    # -- outcome ------------------------------------------------------------

    @property
    def state(self) -> str:
        """mode:phase, e.g. "reg:msg2"; "established"/"failed" at the end."""
        if self._phase in ("established", "failed") or self._registering is None:
            return self._phase
        mode = "reg" if self._registering else "auth"
        return f"{mode}:{self._phase}"

    @property
    def verification(self) -> Verification:
        """The out-of-band comparison value (*:verify, or established
        over a still-unverified credential)."""
        if self._phase == "verify" or (
            self._phase == "established" and not self._credential_verified
        ):
            return Verification(self._sas)
        raise ProtocolError(f"nothing to verify in state {self.state}")

    @property
    def verified(self) -> bool:
        """Whether the session was verified (established only)."""
        if self._phase != "established":
            raise ProtocolError("handshake not complete")
        return self._verified

    @property
    def credential(self) -> bytes:
        """The stored credential for this session's peer."""
        if self._phase != "established":
            raise ProtocolError("handshake not complete")
        return self._store[self._key_id]

    def confirm(self, verified: bool) -> None:
        """Post-establishment decision on a still-unverified credential:
        True flags it verified, False deletes it. The session itself is
        unaffected."""
        if self._phase != "established" or self._credential_verified:
            raise ProtocolError(f"nothing to confirm in state {self.state}")
        if verified:
            current = self._store.get(self._key_id)
            if current is None or not _same_credential(current, self._stored):
                raise ProtocolError("credential changed underneath us")
            self._swap_credential(verified=True)
            self._verified = True
        else:
            # rejected: the credential is deleted outright -- a genuine
            # mismatch means a MITM; re-provision a fresh PSK
            self._delete_credential()

    # -- credential loading --------------------------------------------------

    def _load_credential(self, key_id) -> Registration | Resumption:
        # bad local credential -> caller's problem (KeyError/struct.error)
        cred = self._store[key_id]
        self._key_id = key_id
        self._stored = cred  # CAS reference for every later store write
        if len(cred) == 32:  # bare registration secret
            self._idroot = random(32)  # fresh identity root
            self._shared = cred
            self._registering = True
            self._credential_verified = False
            hs: Registration | Resumption = Registration(cred, self._idroot, self._ad)
        else:
            (
                self._idroot,
                self._shared,
                peer_m,
                peer_x,
                self._sas,
                self._credential_verified,
            ) = _CREDENTIAL.unpack(cred)
            self._peer_pub = (peer_m, peer_x)
            self._registering = False
            hs = Resumption(self._shared, self._idroot, self._peer_pub, self._ad)
        self._exchange = hs
        return hs

    # -- store mutations -------------------------------------------------------
    #
    # Every write is a compare-and-swap against the credential this session
    # loaded (self._stored): with several connections in flight on one
    # store, only the first to reach the atomic key change wins, and a
    # loser's later burn/delete/flag-flip never touches a slot that
    # changed underneath it. The store is synchronously accessed (a
    # plain dict), so these checks are atomic by construction.

    def _swap_credential(self, verified: bool) -> None:
        # atomic credential replacement in the store; the application
        # persists the store itself if it wants durability
        self._credential_verified = verified
        self._stored = _CREDENTIAL.pack(
            self._idroot, self._shared, *self._peer_pub, self._sas, verified
        )
        self._store[self._key_id] = self._stored

    def _delete_credential(self) -> None:
        # delete only if the slot still holds OUR credential: a parallel
        # connection may have won the slot in the meantime
        current = self._store.get(self._key_id)
        if current is not None and _same_credential(current, self._stored):
            del self._store[self._key_id]

    # -- state handlers ------------------------------------------------------

    def _st_hello(self, msg: _In) -> list[_Out]:
        if msg is None:  # initiate ("alice")
            hs = self._exchange
            if hs is None:
                raise ProtocolError("initiation requires a key_id")
            self._initiator = True
            self._phase = "msg2"
            return [hs.startA()]
        if not isinstance(msg, bytes):
            raise ProtocolError(f"expecting msg1 in state {self.state}")
        # respond ("bob"): identify the caller by trying msg1 against
        # every credential in the store -- an AEAD open per candidate is
        # cheap, and a box that opens pins the credential (and its mode).
        # No candidate opening proves no key knowledge: nothing burns.
        for key_id, cred in self._store.items():
            try:
                opened = open_hello(
                    msg,
                    shared=_credential_shared(cred),
                    ad=self._ad,
                    flow="reg" if len(cred) == 32 else "res",
                )
            except (struct.error, ValueError):
                continue  # not this credential (or a malformed store entry)
            break
        else:
            raise ProtocolError("invalid authentication tag")
        self._initiator = False
        self._phase = "msg3"
        return [self._load_credential(key_id).respondB(msg, opened)]

    def _st_msg2(self, msg: _In) -> list[_Out]:
        if not isinstance(msg, bytes):
            raise ProtocolError(f"expecting msg2 in state {self.state}")
        hs = self._exchange
        assert hs is not None  # a credential is loaded by now
        msg3 = hs.respondA(msg)
        if self._registering:  # registration continues to msg4
            self._phase = "msg4"
            return [msg3]
        return [msg3, *self._keyed(hs, initiator=True)]

    def _st_msg3(self, msg: _In) -> list[_Out]:
        if not isinstance(msg, bytes):
            raise ProtocolError(f"expecting msg3 in state {self.state}")
        hs = self._exchange
        assert hs is not None  # a credential is loaded by now
        if self._registering:
            assert isinstance(hs, Registration)
            # msg4 is held back until the comparison answer
            self._pending, self._sas = hs.confirmB(msg)
            self._phase = "verify"
            return [Verification(self._sas)]
        assert isinstance(hs, Resumption)
        hs.completeB(msg)
        return self._keyed(hs, initiator=False)

    def _st_msg4(self, msg: _In) -> list[_Out]:
        # registration only: bob's challenge and confirmation arrive;
        # msg5 is held back until the comparison answer
        if not isinstance(msg, bytes):
            raise ProtocolError(f"expecting msg4 in state {self.state}")
        hs = self._exchange
        assert isinstance(hs, Registration)
        self._pending, self._sas = hs.confirmA(msg)
        self._phase = "verify"
        return [Verification(self._sas)]

    def _st_msg5(self, msg: _In) -> list[_Out]:
        # registration only: the initiator's confirmation finishes the
        # exchange -- the key change happens here
        if not isinstance(msg, bytes):
            raise ProtocolError(f"expecting msg5 in state {self.state}")
        hs = self._exchange
        assert isinstance(hs, Registration)
        hs.finishB(msg)
        self._install_keys(hs, initiator=False)
        self._swap_at_gate(self._verified_answer)
        return [self._establish(verified=self._verified_answer)]

    def _st_verify(self, msg: _In) -> list[_Out]:
        # registration only: the comparison answer gates the tail of the
        # exchange. False deletes the credential and fails; True and
        # None both proceed (verified / unverified).
        if not isinstance(msg, Verification):
            raise ProtocolError(f"expecting the comparison answer in {self.state}")
        if msg.answer is False:
            self._delete_credential()
            self._phase = "failed"
            raise ProtocolError("verification rejected")
        hs = self._exchange
        assert isinstance(hs, Registration)  # the resumption flow never gates
        pending = self._pending
        assert pending is not None  # set by the gate that led here
        if self._initiator:
            # everything is in: key change, traffic channels, established
            self._install_keys(hs, initiator=True)
            self._swap_at_gate(msg.answer is True)
            return [pending, self._establish(verified=msg.answer is True)]
        # responder: release the held msg4; the key change waits for msg5
        self._verified_answer = msg.answer is True
        self._phase = "msg5"
        return [pending]

    def _swap_at_gate(self, verified: bool) -> None:
        # the composite atomically replaces the registration secret,
        # flagged by the answer; the swap is the arbitration point for
        # parallel registrations -- the first swap wins, a later one
        # finds the registration secret gone and fails
        if self._store.get(self._key_id) != self._stored:
            self._phase = "failed"
            raise ProtocolError("registration key already spent")
        self._swap_credential(verified=verified)

    def _keyed(self, hs: Registration | Resumption, initiator: bool) -> list[_Out]:
        # resumption only: install the traffic channels, no comparison
        self._install_keys(hs, initiator)
        return [self._establish(verified=self._credential_verified)]

    def _st_data(self, msg: _In) -> list[_Out]:
        # traffic: Messages encrypt into wire frames, frames decrypt
        # back to Messages; the kind byte is opaque here (0 = protocol)
        if self._phase != "established":
            raise ProtocolError(f"no traffic in state {self.state}")
        if isinstance(msg, Message):
            if msg.kind == 0:
                raise ProtocolError("frame kind 0 is reserved for protocol use")
            return [bytes(self._tx.encrypt(bytes([msg.kind]) + msg.data))]
        if not isinstance(msg, bytes):
            raise ProtocolError(f"expecting traffic in state {self.state}")
        plain = self._rx.decrypt(msg)
        if not plain:
            raise ProtocolError("empty traffic frame")
        kind, data = plain[0], bytes(plain[1:])
        if kind == 0:
            # in-band protocol messages will be handled here
            raise ProtocolError("frame kind 0 is reserved for protocol use")
        return [Message(data, kind)]

    # -- transitions shared by both roles ------------------------------------

    def _install_keys(self, hs: Registration | Resumption, initiator: bool) -> None:
        # pull the results out of the exchange and dispose of its state
        self._exchange = None
        if isinstance(hs, Registration):
            send, recv, self._shared, self._sas, self._peer_pub = hs.finish(initiator)
            # rotated shared in memory only: the store swap waits for the answer
        else:
            send, recv = hs.finish(initiator)
        self._tx, self._rx = Sender(send), Receiver(recv)

    def _establish(self, verified: bool) -> Established:
        self._verified = verified
        self._phase = "established"
        return Established(self._sas, verified)