"""Agent key material: ECDSA P-256 key generation, DID derivation, fingerprinting."""

from __future__ import annotations

import hashlib
import json
import os
from dataclasses import dataclass
from typing import Any

from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives.serialization import (
    Encoding,
    NoEncryption,
    PrivateFormat,
    PublicFormat,
)


@dataclass
class AgentKeyMaterial:
    private_jwk: dict[str, Any]
    public_jwk: dict[str, Any]
    did: str
    verification_method: str
    key_fingerprint: str


def _stable_json(value: Any) -> str:
    """JSON-serialize with sorted keys and no whitespace, matching the TS SDK's stableStringify."""
    if isinstance(value, dict):
        entries = sorted(value.items())
        inner = ",".join(f"{json.dumps(k)}:{_stable_json(v)}" for k, v in entries)
        return "{" + inner + "}"
    if isinstance(value, list):
        return "[" + ",".join(_stable_json(item) for item in value) + "]"
    return json.dumps(value)


def _base64url_encode(data: bytes) -> str:
    import base64

    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")


def _base64url_decode(s: str) -> bytes:
    import base64

    padding = 4 - len(s) % 4
    if padding != 4:
        s += "=" * padding
    return base64.urlsafe_b64decode(s)


def _ec_key_to_jwk(private_key: ec.EllipticCurvePrivateKey) -> tuple[dict, dict]:
    """Export an EC private key as (private_jwk, public_jwk) dicts."""
    private_numbers = private_key.private_numbers()
    public_numbers = private_numbers.public_numbers

    byte_len = (public_numbers.curve.key_size + 7) // 8

    public_jwk = {
        "crv": "P-256",
        "kty": "EC",
        "x": _base64url_encode(public_numbers.x.to_bytes(byte_len, "big")),
        "y": _base64url_encode(public_numbers.y.to_bytes(byte_len, "big")),
    }

    private_jwk = {
        **public_jwk,
        "d": _base64url_encode(private_numbers.private_value.to_bytes(byte_len, "big")),
    }

    return private_jwk, public_jwk


def did_jwk(public_jwk: dict[str, Any]) -> str:
    """Derive a did:jwk from a public JWK."""
    encoded = _base64url_encode(_stable_json(public_jwk).encode("utf-8"))
    return f"did:jwk:{encoded}"


def verification_method_for_did(did: str) -> str:
    """Return the #0 verification method for a DID."""
    return f"{did}#0"


def fingerprint_public_jwk(public_jwk: dict[str, Any]) -> str:
    """SHA-256 hex fingerprint of a public JWK (deterministic, sorted)."""
    return hashlib.sha256(_stable_json(public_jwk).encode("utf-8")).hexdigest()


def generate_agent_key_material() -> AgentKeyMaterial:
    """Generate an ECDSA P-256 key pair with DID and fingerprint."""
    private_key = ec.generate_private_key(ec.SECP256R1())
    private_jwk, public_jwk = _ec_key_to_jwk(private_key)

    did = did_jwk(public_jwk)

    return AgentKeyMaterial(
        private_jwk=private_jwk,
        public_jwk=public_jwk,
        did=did,
        verification_method=verification_method_for_did(did),
        key_fingerprint=fingerprint_public_jwk(public_jwk),
    )


def load_agent_key_material() -> AgentKeyMaterial:
    """Load agent key material from env var, file path env var, or ./agent-key.json."""

    def from_export(data: dict[str, Any]) -> AgentKeyMaterial:
        return AgentKeyMaterial(
            private_jwk=data["privateJwk"],
            public_jwk=data["publicJwk"],
            did=data["did"],
            verification_method=data["verificationMethod"],
            key_fingerprint=data["keyFingerprint"],
        )

    env_json = os.environ.get("HALLPASS_AGENT_KEY")
    if env_json:
        return from_export(json.loads(env_json))

    env_file = os.environ.get("HALLPASS_AGENT_KEY_FILE")
    if env_file:
        with open(env_file, encoding="utf-8") as f:
            return from_export(json.load(f))

    default_path = "agent-key.json"
    if os.path.exists(default_path):
        with open(default_path, encoding="utf-8") as f:
            return from_export(json.load(f))

    raise FileNotFoundError(
        "No agent key material found. Set HALLPASS_AGENT_KEY (JSON), "
        "HALLPASS_AGENT_KEY_FILE (path), or place agent-key.json in the working directory."
    )
