"""Key material for the three PoC identities.

The PoC uses three ES256 (P-256) keys — the common algorithm for both
ACT and ECT per draft-nennemann-act-01 §5 and draft-nennemann-wimse-ect-01 §5.

Identities:
  user        — issues the ACT mandate (iss in Phase 1)
  agent       — subject of the mandate, signs Phase 2 record, signs ECT on
                every MCP tool call (sub in ACT, iss in ECT)
  mcp-server  — audience / verifier (aud in both ACT and ECT)

Keys are written to ``keys/`` as PEM files on first run; subsequent runs
load them. This mimics pre-shared-key deployment per ACT §5.2 Tier 1.
"""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric.ec import (
    EllipticCurvePrivateKey,
    EllipticCurvePublicKey,
)

from act.crypto import KeyRegistry, generate_p256_keypair


IDENTITIES = ("user", "agent", "mcp-server")


@dataclass
class Identity:
    name: str
    kid: str
    private_key: EllipticCurvePrivateKey
    public_key: EllipticCurvePublicKey


def _pem_paths(keys_dir: Path, name: str) -> tuple[Path, Path]:
    return keys_dir / f"{name}.priv.pem", keys_dir / f"{name}.pub.pem"


def _load_or_generate(keys_dir: Path, name: str) -> Identity:
    priv_path, pub_path = _pem_paths(keys_dir, name)
    if priv_path.exists() and pub_path.exists():
        priv_bytes = priv_path.read_bytes()
        priv = serialization.load_pem_private_key(priv_bytes, password=None)
        assert isinstance(priv, EllipticCurvePrivateKey), (
            f"{name}.priv.pem is not a P-256 private key"
        )
        pub = priv.public_key()
    else:
        priv, pub = generate_p256_keypair()
        keys_dir.mkdir(parents=True, exist_ok=True)
        priv_path.write_bytes(
            priv.private_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PrivateFormat.PKCS8,
                encryption_algorithm=serialization.NoEncryption(),
            )
        )
        pub_path.write_bytes(
            pub.public_bytes(
                encoding=serialization.Encoding.PEM,
                format=serialization.PublicFormat.SubjectPublicKeyInfo,
            )
        )
    kid = f"kid:{name}:v1"
    return Identity(name=name, kid=kid, private_key=priv, public_key=pub)


def load_identities(keys_dir: str | Path = "keys") -> dict[str, Identity]:
    """Load all three PoC identities, generating key material if missing."""
    keys_dir = Path(keys_dir)
    return {name: _load_or_generate(keys_dir, name) for name in IDENTITIES}


def build_key_registry(identities: dict[str, Identity]) -> KeyRegistry:
    """Assemble an ACT KeyRegistry with every identity's public key."""
    reg = KeyRegistry()
    for ident in identities.values():
        reg.register(ident.kid, ident.public_key)
    return reg


def build_ect_key_resolver(identities: dict[str, Identity]):
    """Return an ECT KeyResolver callable that maps kid → public key."""
    kid_to_pub: dict[str, EllipticCurvePublicKey] = {
        ident.kid: ident.public_key for ident in identities.values()
    }

    def _resolve(kid: str) -> EllipticCurvePublicKey | None:
        return kid_to_pub.get(kid)

    return _resolve