from __future__ import annotations

from dataclasses import dataclass, asdict
from typing import Any, Dict, List
import hashlib
import json

from .dsl import PlanDelta


@dataclass
class CSFDFrame:
    timestamp_offset: int
    invariant_id_list: List[str]
    minimal_state_snapshot_hash: str
    expected_resource_bounds: Dict[str, Any]
    proof_stub_hash: str


@dataclass
class CSFDDigest:
    frames: List[CSFDFrame]
    signer: str = ""
    digest_hex: str = ""


def _hash_obj(obj: Any) -> str:
    s = json.dumps(obj, sort_keys=True, separators=(",", ":"))
    return hashlib.sha256(s.encode()).hexdigest()


def generate_csfd(plan_delta: PlanDelta, signer: str = "", n_frames: int = 3) -> CSFDDigest:
    """Generate a compact Critical-Safety Frame Digest (CSFD) for a PlanDelta.

    The digest contains up to `n_frames` ordered frames. Each frame is a
    compact summary referencing invariant ids and a minimal snapshot hash.

    This implementation is intentionally small and deterministic for the
    MVP; a production implementation should include proper signing (ed25519)
    and configurable snapshot selectors.
    """
    # Source invariants from PlanDelta.metadata (if present)
    invariants = plan_delta.metadata.get("invariants", {}) if plan_delta and plan_delta.metadata else {}

    # Flatten invariant ids deterministically
    invariant_ids = sorted(list(map(str, invariants.keys())))

    frames: List[CSFDFrame] = []
    # Create N frames. For small plan deltas we may repeat or slice invariant ids.
    for i in range(n_frames):
        # take a slice of invariant ids for this frame
        slice_start = (i * 1) % max(1, len(invariant_ids)) if invariant_ids else 0
        # choose up to 3 invariants per frame
        ids = invariant_ids[slice_start : slice_start + 3] if invariant_ids else []

        # minimal state snapshot: hash of a few plan delta fields
        snapshot = {
            "version": plan_delta.version,
            "delta_keys": sorted(list(map(str, plan_delta.delta.keys()))) if plan_delta.delta else [],
            "meta_keys": sorted(list(map(str, plan_delta.metadata.keys()))) if plan_delta.metadata else [],
            "frame_index": i,
        }

        minimal_hash = _hash_obj(snapshot)

        # proof_stub_hash: if metadata provides a proof stub map use it, else hash the snapshot
        proof_stub_hash = (
            plan_delta.metadata.get("proof_stub_hash", "")
            or _hash_obj({"snapshot": snapshot, "invariants": ids})
        )

        # expected_resource_bounds: best-effort from metadata
        expected_bounds = plan_delta.metadata.get("expected_resource_bounds", {}) or {}

        frame = CSFDFrame(
            timestamp_offset=i,  # simplistic offset in frames
            invariant_id_list=ids,
            minimal_state_snapshot_hash=minimal_hash,
            expected_resource_bounds=expected_bounds,
            proof_stub_hash=proof_stub_hash,
        )
        frames.append(frame)

    # digest: compact hex over serialized frames
    digest_src = [asdict(f) for f in frames]
    digest_hex = hashlib.sha256(json.dumps(digest_src, sort_keys=True, separators=(",", ":")).encode()).hexdigest()

    csfd = CSFDDigest(frames=frames, signer=signer, digest_hex=digest_hex)

    # Safety: ensure serialized size is reasonably small (<= 512 bytes for MVP)
    serialized = json.dumps(asdict(csfd), sort_keys=True, separators=(",", ":"))
    if len(serialized.encode()) > 512:
        # If too large, truncate frames to last frame only and recompute digest
        frames = frames[:1]
        digest_src = [asdict(f) for f in frames]
        digest_hex = hashlib.sha256(json.dumps(digest_src, sort_keys=True, separators=(",", ":")).encode()).hexdigest()
        csfd = CSFDDigest(frames=frames, signer=signer, digest_hex=digest_hex)

    return csfd