"""Ed25519 hardware receipts.

A sign-off means: this key attests it observed this readout and this document.
It is not UL, NIST, or NVIDIA confidential-compute attestation.
"""

from __future__ import annotations

import hashlib
import json
import os
from datetime import datetime, timezone
from pathlib import Path

from nacl.exceptions import BadSignatureError
from nacl.signing import SigningKey, VerifyKey

from schema.models import Bundle, HardwareReadout, SignOff, SignerKind

KEY_DIR = Path(os.environ.get("MYCOLO_KEY_DIR", Path.home() / ".config" / "mycolo"))
KEY_PATH = KEY_DIR / "ed25519.seed"


def _sha256(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def hardware_fingerprint(hw: HardwareReadout) -> str:
    blob = json.dumps(
        hw.model_dump(mode="json", exclude={"os_name"}),
        sort_keys=True,
        separators=(",", ":"),
    ).encode()
    return _sha256(blob)


def canonical_bundle(bundle: Bundle) -> bytes:
    payload = bundle.model_dump(mode="json", exclude={"content_sha256", "signoffs"})
    hw = payload.get("hardware")
    if isinstance(hw, dict):
        hw.pop("os_name", None)
    return json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()


def stamp(bundle: Bundle) -> Bundle:
    bundle.hardware_fingerprint = hardware_fingerprint(bundle.hardware)
    bundle.content_sha256 = _sha256(canonical_bundle(bundle))
    return bundle


def keygen(path: Path = KEY_PATH) -> str:
    path.parent.mkdir(parents=True, exist_ok=True)
    if path.exists():
        raise FileExistsError(f"key already exists: {path}")
    sk = SigningKey.generate()
    path.write_bytes(bytes(sk))
    path.chmod(0o600)
    return bytes(sk.verify_key).hex()


def load_signing_key(path: Path = KEY_PATH) -> SigningKey:
    raw = path.read_bytes()
    return SigningKey(raw)


def public_key_hex(path: Path = KEY_PATH) -> str:
    return bytes(load_signing_key(path).verify_key).hex()


def sign_bundle(
    bundle: Bundle,
    *,
    signer_kind: SignerKind,
    signer_label: str | None,
    path: Path = KEY_PATH,
) -> Bundle:
    bundle = stamp(bundle)
    sk = load_signing_key(path)
    msg = (bundle.content_sha256 or "").encode()
    sig = sk.sign(msg).signature
    bundle.signoffs.append(
        SignOff(
            public_key_hex=bytes(sk.verify_key).hex(),
            signature_hex=sig.hex(),
            signed_at=datetime.now(timezone.utc),
            signer_kind=signer_kind,
            signer_label=signer_label,
        )
    )
    return bundle


def verify_signoff(bundle: Bundle, signoff: SignOff) -> bool:
    digest = bundle.content_sha256 or _sha256(canonical_bundle(bundle))
    try:
        vk = VerifyKey(bytes.fromhex(signoff.public_key_hex))
        vk.verify(digest.encode(), bytes.fromhex(signoff.signature_hex))
        return True
    except (BadSignatureError, ValueError):
        return False


def verify_bundle(bundle: Bundle) -> dict:
    recomputed = _sha256(canonical_bundle(bundle))
    hash_ok = bundle.content_sha256 == recomputed
    expected_fp = hardware_fingerprint(bundle.hardware)
    fp_ok = bundle.hardware_fingerprint == expected_fp
    results = []
    for s in bundle.signoffs:
        ok = hash_ok and verify_signoff(bundle, s)
        results.append(
            {
                "ok": ok,
                "signer_kind": s.signer_kind.value,
                "signer_label": s.signer_label,
                "public_key_hex": s.public_key_hex,
                "statement": s.statement,
            }
        )
    any_ok = any(r["ok"] for r in results)
    return {
        "receipt_ok": bool(any_ok and fp_ok and hash_ok),
        "fingerprint_ok": fp_ok,
        "hash_ok": hash_ok,
        "hardware_fingerprint": expected_fp,
        "signoffs": results,
        "disclaimer": "First-party hardware receipt. Not a laboratory or OEM certification.",
    }


def render_receipt(bundle: Bundle, verification: dict | None = None) -> str:
    v = verification or verify_bundle(bundle)
    hw = bundle.hardware
    lines = [
        "MYCOLO HARDWARE RECEIPT v0",
        "First-party attestation — not a lab / OEM certificate.",
        "",
        f"fingerprint: {v['hardware_fingerprint']}",
        f"fingerprint_ok: {v['fingerprint_ok']}",
        f"receipt_ok: {v['receipt_ok']}",
        f"gpu_or_soc: {hw.gpu_or_soc}",
        f"vram_or_unified_mib: {hw.vram_or_unified_mib}",
        f"driver: {hw.driver}",
        f"cpu: {hw.cpu}",
        f"stack: {hw.stack.value}",
        f"content_sha256: {bundle.content_sha256}",
    ]
    if bundle.model:
        lines.append(f"model: {bundle.model.build_id} {bundle.model.quant or ''}".rstrip())
    if bundle.score:
        lines.append(f"score_tok_s: {bundle.score.tok_per_s} (score trust={bundle.score.trust.value})")
    lines.append("")
    if not v["signoffs"]:
        lines.append("signoffs: none")
    for s in v["signoffs"]:
        lines.append(
            f"signoff: {s['signer_kind']}:{s['signer_label'] or '-'} "
            f"ok={s['ok']} key={s['public_key_hex'][:16]}…"
        )
    lines.append("")
    lines.append(v["disclaimer"])
    return "\n".join(lines) + "\n"