"""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, WireError, random
from .exceptions import ProtocolError
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 WireError as e:
            # unauthentic box: proves no key knowledge, nothing burns
            _log.debug("fail state=%s: %r", self.state, e)
            raise ProtocolError(str(e)) from None
        except (ValueError, struct.error) as e:
            # 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: %r", self.state, e)
            raise ProtocolError("credential burned") 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)