from __future__ import annotations

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

try:
    from nacl.signing import SigningKey, VerifyKey
    from nacl.exceptions import BadSignatureError
    _HAS_PYNACL = True
except Exception:
    # Keep module import-safe for environments without PyNaCl installed.
    SigningKey = None
    VerifyKey = None
    BadSignatureError = Exception
    _HAS_PYNACL = False

try:
    import cbor2
    _HAS_CBOR = True
except Exception:
    cbor2 = None
    _HAS_CBOR = False

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


def sign_csfd(csfd: CSFDDigest, signing_key_bytes: bytes) -> str:
    """Sign a CSFD digest with an ed25519 private key (SigningKey bytes).

    Returns the signature as a hex string. Requires PyNaCl.
    """
    if not _HAS_PYNACL:
        raise RuntimeError("PyNaCl is required for signing CSFDs")

    if not signing_key_bytes:
        raise ValueError("signing_key_bytes must be provided")

    sk = SigningKey(signing_key_bytes)
    # Sign the compact digest_hex to produce a small signature
    sig = sk.sign(bytes.fromhex(csfd.digest_hex)).signature
    return sig.hex()


def verify_csfd_signature(csfd: CSFDDigest, signature_hex: str, verify_key_bytes: bytes) -> bool:
    """Verify a previously-created signature on the CSFD digest.

    Returns True if valid, False otherwise. Requires PyNaCl.
    """
    if not _HAS_PYNACL:
        raise RuntimeError("PyNaCl is required for signature verification")

    vk = VerifyKey(verify_key_bytes)
    try:
        vk.verify(bytes.fromhex(csfd.digest_hex), bytes.fromhex(signature_hex))
        return True
    except BadSignatureError:
        return False


def csfd_to_cbor(csfd: CSFDDigest) -> bytes:
    """Serialize a CSFD to CBOR bytes. Requires cbor2 if you want CBOR output.

    Falls back to JSON bytes if cbor2 is not available.
    """
    obj = asdict(csfd)
    if _HAS_CBOR:
        return cbor2.dumps(obj)
    # fallback: compact JSON bytes
    return json.dumps(obj, sort_keys=True, separators=(",", ":")).encode()


def csfd_from_cbor(b: bytes) -> CSFDDigest:
    """Deserialize CBOR/JSON bytes into a CSFDDigest.

    Accepts CBOR if cbor2 is available; otherwise expects UTF-8 JSON bytes.
    """
    if _HAS_CBOR:
        obj = cbor2.loads(b)
    else:
        obj = json.loads(b.decode())

    frames = [CSFDFrame(**f) for f in obj.get("frames", [])]
    return CSFDDigest(frames=frames, signer=obj.get("signer", ""), digest_hex=obj.get("digest_hex", ""))