"""
Proof Pack v0 Builder.

Builds the 5-file execution kernel:
  - receipt_pack.jsonl     Canonical receipts in deterministic order
  - verify_report.json     Machine-readable verification results
  - verify_transcript.md   Human-readable summary with attestation block
  - pack_manifest.json     Signed root envelope
  - pack_signature.sig     Detached Ed25519 signature

Signing workflow (per spec):
  1. Build unsigned manifest, validate structure
  2. JCS-canonicalize unsigned manifest, sign with Ed25519
  3. Attach signature to manifest -> pack_manifest.json
  4. Write detached pack_signature.sig with same bytes
"""
from __future__ import annotations

import base64
import hashlib
import json
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional

from assay.claim_verifier import ClaimSetResult, ClaimSpec, verify_claims
from assay.integrity import VerifyResult, verify_receipt_pack
from assay.keystore import AssayKeyStore, DEFAULT_SIGNER_ID, get_default_keystore
from assay._receipts.canonicalize import to_jcs_bytes

try:
    from assay import __version__ as _assay_version
except Exception:
    _assay_version = "0.1.0"


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

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


def _generate_pack_id(*, deterministic_seed: Optional[str] = None) -> str:
    """Generate a pack ID.

    When deterministic_seed is provided (e.g. corpus generation), the ID
    is content-addressed and reproducible.  Otherwise it uses a timestamp
    and random suffix for production packs.
    """
    if deterministic_seed is not None:
        tag = hashlib.sha256(deterministic_seed.encode()).hexdigest()[:8]
        return f"pack_deterministic_{tag}"
    ts = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S")
    return f"pack_{ts}_{uuid.uuid4().hex[:8]}"


def _sort_receipts(entries: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    """Sort receipts by (run_id, seq, receipt_id) for deterministic ordering."""
    def sort_key(r: Dict[str, Any]):
        return (
            r.get("_trace_id", r.get("run_id", "")),
            r.get("seq", 0),
            r.get("receipt_id", ""),
        )
    return sorted(entries, key=sort_key)


# ---------------------------------------------------------------------------
# Transcript generator
# ---------------------------------------------------------------------------

def _generate_transcript(
    pack_id: str,
    attestation: Dict[str, Any],
    verify_result: VerifyResult,
    version: str,
    *,
    generated_at: Optional[str] = None,
) -> str:
    """Generate verify_transcript.md from attestation and verification results."""
    errors_section = "None"
    if verify_result.errors:
        lines = []
        for e in verify_result.errors:
            loc = f" (receipt {e.receipt_index})" if e.receipt_index is not None else ""
            lines.append(f"- **{e.code}**{loc}: {e.message}")
        errors_section = "\n".join(lines)

    warnings_section = "None"
    if verify_result.warnings:
        warnings_section = "\n".join(f"- {w}" for w in verify_result.warnings)

    ts_start = attestation.get("timestamp_start", "N/A")
    ts_end = attestation.get("timestamp_end", "N/A")

    return f"""# Proof Pack Verification Transcript

## Attestation

| Field | Value |
|-------|-------|
| Pack ID | `{pack_id}` |
| Run ID | `{attestation.get('run_id', 'N/A')}` |
| Receipt Integrity | **{attestation.get('receipt_integrity', 'N/A')}** |
| Claim Check | {attestation.get('claim_check', 'N/A')} |
| Assurance Level | {attestation.get('assurance_level', 'N/A')} |
| Mode | {attestation.get('mode', 'N/A')} |
| Receipts | {attestation.get('n_receipts', 0)} |
| Head Hash | `{attestation.get('head_hash', 'N/A')}` |
| Policy Hash | `{attestation.get('policy_hash', 'N/A')}` |
| Suite ID | `{attestation.get('suite_id', 'N/A')}` |
| Time Range | {ts_start} to {ts_end} |

## Verification Errors

{errors_section}

## Verification Warnings

{warnings_section}

---
Generated by Assay {version} at {generated_at or datetime.now(timezone.utc).isoformat()}
Verify: `assay verify-pack <pack_dir>`
"""


# ---------------------------------------------------------------------------
# ProofPack builder
# ---------------------------------------------------------------------------

class ProofPack:
    """Builds the 5-file Proof Pack v0 kernel."""

    def __init__(
        self,
        trace_id: Optional[str] = None,
        entries: Optional[List[Dict[str, Any]]] = None,
        *,
        run_id: Optional[str] = None,
        signer_id: str = DEFAULT_SIGNER_ID,
        policy_hash: Optional[str] = None,
        suite_id: str = "manual",
        suite_hash: Optional[str] = None,
        claim_set_id: str = "none",
        claim_set_hash: Optional[str] = None,
        claims: Optional[List[ClaimSpec]] = None,
        mode: str = "shadow",
    ):
        # run_id is canonical; trace_id accepted as alias for backward compat
        resolved_run_id = run_id or trace_id
        if not resolved_run_id:
            raise ValueError("run_id (or trace_id) is required")
        self.run_id = resolved_run_id
        self.entries = entries or []
        self.signer_id = signer_id
        self.mode = mode
        self.claims = claims

        # Default hashes for fields not yet wired
        self.policy_hash = policy_hash or _sha256_hex(b"default-policy-v0")
        self.suite_id = suite_id
        self.suite_hash = suite_hash or _sha256_hex(suite_id.encode())
        self.claim_set_id = claim_set_id

        # When claims are provided, compute real claim_set_hash from specs
        if claims and not claim_set_hash:
            specs = [c.to_dict() for c in claims]
            self.claim_set_hash = _sha256_hex(to_jcs_bytes(specs))
        else:
            self.claim_set_hash = claim_set_hash or _sha256_hex(claim_set_id.encode())

    def build(
        self,
        output_dir: Path,
        keystore: Optional[AssayKeyStore] = None,
        *,
        pack_id: Optional[str] = None,
        deterministic_ts: Optional[str] = None,
    ) -> Path:
        """Build the 5-file Proof Pack to output_dir.

        Args:
            pack_id: Explicit pack ID.  When omitted, one is generated
                     (random in production, content-addressed if
                     deterministic_ts is also set).
            deterministic_ts: Fixed ISO-8601 timestamp for all
                              time-dependent fields.  Used by
                              conformance corpus generation to ensure
                              bit-identical rebuilds.

        Returns output_dir on success.
        """
        ks = keystore or get_default_keystore()
        output_dir = Path(output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)

        if pack_id is None:
            seed = self.run_id if deterministic_ts else None
            pack_id = _generate_pack_id(deterministic_seed=seed)
        sorted_entries = _sort_receipts(self.entries)

        # 1. Write receipt_pack.jsonl (canonical, deterministic order).
        # JSONL invariant: one JCS-canonical JSON object per line, no blank
        # lines.  Non-empty packs end with exactly one trailing newline.
        # Empty packs produce a 0-byte file (no newline).
        # Verifier counts non-empty lines (line.strip()).
        receipt_lines = []
        for entry in sorted_entries:
            canonical = to_jcs_bytes(entry).decode("utf-8")
            receipt_lines.append(canonical)
        receipt_pack_content = "\n".join(receipt_lines) + "\n" if receipt_lines else ""
        receipt_pack_bytes = receipt_pack_content.encode("utf-8")
        (output_dir / "receipt_pack.jsonl").write_bytes(receipt_pack_bytes)

        # 2. Verify receipts (structural integrity)
        verify_result = verify_receipt_pack(sorted_entries)

        # 2b. Verify claims (semantic) if claims provided
        claim_result: Optional[ClaimSetResult] = None
        if self.claims:
            claim_result = verify_claims(
                sorted_entries, self.claims,
                policy_hash=self.policy_hash, suite_hash=self.suite_hash,
            )

        # 3. Build verify_report.json
        report: Dict[str, Any] = {
            "pack_id": pack_id,
            "run_id": self.run_id,
            **verify_result.to_dict(),
            "verified_at": deterministic_ts or datetime.now(timezone.utc).isoformat(),
            "verifier_version": _assay_version,
        }
        if claim_result is not None:
            report["claim_verification"] = claim_result.to_dict()
        report_bytes = json.dumps(report, indent=2).encode("utf-8")
        (output_dir / "verify_report.json").write_bytes(report_bytes)

        # 4. Build attestation object
        ts_start = None
        ts_end = None
        for entry in sorted_entries:
            ts = entry.get("timestamp") or entry.get("_stored_at")
            if ts:
                if ts_start is None:
                    ts_start = str(ts)
                ts_end = str(ts)

        attestation = {
            "pack_format_version": "0.1.0",
            "fingerprint_version": 1,
            "pack_id": pack_id,
            "run_id": self.run_id,
            "suite_id": self.suite_id,
            "suite_hash": self.suite_hash,
            "verifier_version": _assay_version,
            "canon_version": "jcs-rfc8785",
            "canon_impl": "receipts.jcs",
            "canon_impl_version": _assay_version,
            "policy_hash": self.policy_hash,
            "claim_set_id": self.claim_set_id,
            "claim_set_hash": self.claim_set_hash,
            "receipt_integrity": "PASS" if verify_result.passed else "FAIL",
            "claim_check": (
                "PASS" if claim_result.passed else "FAIL"
            ) if claim_result is not None else "N/A",
            "discrepancy_fingerprint": (
                claim_result.discrepancy_fingerprint
            ) if claim_result is not None else None,
            "assurance_level": "L0",
            "proof_tier": "signed-pack",
            "mode": self.mode,
            "head_hash": verify_result.head_hash or _sha256_hex(b"empty"),
            "head_hash_algorithm": "last-receipt-digest-v0",
            "time_authority": "local_clock",
            "n_receipts": verify_result.receipt_count,
            "timestamp_start": ts_start or deterministic_ts or datetime.now(timezone.utc).isoformat(),
            "timestamp_end": ts_end or deterministic_ts or datetime.now(timezone.utc).isoformat(),
        }

        # 5. Build verify_transcript.md
        transcript = _generate_transcript(
            pack_id, attestation, verify_result, _assay_version,
            generated_at=deterministic_ts,
        )
        transcript_bytes = transcript.encode("utf-8")
        (output_dir / "verify_transcript.md").write_bytes(transcript_bytes)

        # 6. Build unsigned manifest
        attestation_bytes = to_jcs_bytes(attestation)
        attestation_sha256 = _sha256_hex(attestation_bytes)

        files_list = [
            {
                "path": "receipt_pack.jsonl",
                "sha256": _sha256_hex(receipt_pack_bytes),
                "bytes": len(receipt_pack_bytes),
            },
            {
                "path": "verify_report.json",
                "sha256": _sha256_hex(report_bytes),
                "bytes": len(report_bytes),
            },
            {
                "path": "verify_transcript.md",
                "sha256": _sha256_hex(transcript_bytes),
                "bytes": len(transcript_bytes),
            },
        ]

        # All 5 expected files in the pack directory.
        # pack_manifest.json and pack_signature.sig cannot be in the
        # hash-covered `files` array: the manifest can't contain its own
        # hash, and pack_signature.sig is produced AFTER the manifest is
        # signed (circular dependency).  Their integrity is instead
        # protected by the Ed25519 signature and the detached-sig parity
        # check in verify_pack_manifest.
        expected_files = [
            "receipt_pack.jsonl",
            "verify_report.json",
            "verify_transcript.md",
            "pack_manifest.json",
            "pack_signature.sig",
        ]

        # 6b. Embed signer's public key fingerprint for offline verification.
        # Ed25519 public keys are 32 bytes; no reason not to include the full
        # key so packs are self-contained evidence objects.
        vk = ks.get_verify_key(self.signer_id)
        pubkey_bytes = vk.encode()

        unsigned_manifest = {
            "pack_id": pack_id,
            "pack_version": "0.1.0",
            "manifest_version": "1.0.0",
            "hash_alg": "sha256",
            "attestation": attestation,
            "attestation_sha256": attestation_sha256,
            "suite_hash": self.suite_hash,
            "claim_set_id": self.claim_set_id,
            "claim_set_hash": self.claim_set_hash,
            "receipt_count_expected": len(sorted_entries),
            "files": files_list,
            "expected_files": expected_files,
            "signer_id": self.signer_id,
            "signer_pubkey": base64.b64encode(pubkey_bytes).decode("ascii"),
            "signer_pubkey_sha256": _sha256_hex(pubkey_bytes),
            "signature_alg": "ed25519",
            "signature_scope": "JCS(pack_manifest_without_signature)",
        }

        # 7. Sign: JCS(unsigned_manifest) -> Ed25519
        canonical_unsigned = to_jcs_bytes(unsigned_manifest)
        signature_b64 = ks.sign_b64(canonical_unsigned, self.signer_id)

        # 8. Create signed manifest.
        # D12: pack_root_sha256 = attestation_sha256, making the attestation
        # the single immutable identifier for the evidence unit.
        pack_root_sha256 = attestation_sha256
        signed_manifest = {
            **unsigned_manifest,
            "signature": signature_b64,
            "pack_root_sha256": pack_root_sha256,
        }

        # 8b. Schema validation (build-time enforcement)
        from assay.manifest_schema import validate_manifest
        schema_errors = validate_manifest(signed_manifest)
        if schema_errors:
            raise ValueError(
                f"Built manifest fails schema validation: {schema_errors[0]}"
            )

        manifest_bytes = json.dumps(signed_manifest, indent=2).encode("utf-8")
        (output_dir / "pack_manifest.json").write_bytes(manifest_bytes)

        # 9. Write detached signature
        sig_raw = base64.b64decode(signature_b64)
        (output_dir / "pack_signature.sig").write_bytes(sig_raw)

        # 10. Write PACK_SUMMARY.md (presentation layer, not part of
        # the 5-file verification kernel). Safe to import here since
        # explain has no dependency on proof_pack.
        try:
            from assay.explain import explain_pack, render_md
            info = explain_pack(output_dir)
            summary = render_md(info)
            (output_dir / "PACK_SUMMARY.md").write_text(summary, encoding="utf-8")
        except Exception:
            pass  # Never fail a pack build over a summary file

        return output_dir


def build_proof_pack(
    trace_id: str,
    output_dir: Optional[Path] = None,
    *,
    keystore: Optional[AssayKeyStore] = None,
    mode: str = "shadow",
    claims: Optional[List[ClaimSpec]] = None,
) -> Path:
    """Convenience function: load trace from store and build a Proof Pack.

    Args:
        trace_id: The trace to package.
        output_dir: Where to write the 5 files (default: proof_pack_{trace_id}/).
        keystore: Optional key store (default: ~/.loom/assay/keys/).
        mode: shadow | enforced | breakglass.
        claims: Optional list of ClaimSpecs for semantic verification.

    Returns:
        Path to the output directory.
    """
    from assay.store import get_default_store

    store = get_default_store()
    entries = store.read_trace(trace_id)
    if not entries:
        raise ValueError(f"Trace not found: {trace_id}")

    if output_dir is None:
        output_dir = Path(f"proof_pack_{trace_id}")

    pack = ProofPack(run_id=trace_id, entries=entries, mode=mode, claims=claims)
    return pack.build(output_dir, keystore=keystore)


__all__ = [
    "ProofPack",
    "build_proof_pack",
]
