"""
StepRunner + ProofWriter — The seal-or-fail execution contract.

StepRunner wraps any business logic function with the guarantee:
  "A step produces a proof artifact or it didn't happen."

The runner:
  1. Builds a StepEnvelope (inputs fingerprinted)
  2. Checks the IdempotencyGate (short-circuit if seen before)
  3. Executes the step function
  4. Finalizes the envelope (outputs fingerprinted)
  5. Hands to ProofWriter for sealing
  6. Records the proof in the gate
  7. Returns the ProofArtifact

If execution raises: a FAILED proof is sealed.  Either way, a proof exists.
No ambiguity for downstream nodes.

ProofWriter abstracts the sealing backend:
  - SbnSealBackend: SnapChore seal + SmartBlock publish (production)
  - LocalSealBackend: SHA-256 local seal (dev/test, no SBN dependency)
"""

from __future__ import annotations

import hashlib
import json
import logging
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any, Awaitable, Callable, Dict, List, Optional, Protocol, runtime_checkable

from .artifact import ProofArtifact, ProofStatus, canonical_fingerprint
from .envelope import StepEnvelope, StepEnvelopeBuilder
from .gate import IdempotencyGate

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Step function protocol
# ---------------------------------------------------------------------------

@dataclass
class StepContext:
    """Context passed to every step function.

    Products extend this with product-specific fields by subclassing
    or by putting product context in the `config` dict.
    """
    inputs: Dict[str, Any]
    prior_proofs: Dict[str, ProofArtifact] = field(default_factory=dict)
    config: Dict[str, Any] = field(default_factory=dict)
    actor: str = ""
    run_id: str = ""
    slot_id: Optional[str] = None
    node_id: str = ""


# The canonical step function signature: takes context, returns outputs dict.
StepFn = Callable[[StepContext], Awaitable[Dict[str, Any]]]


# ---------------------------------------------------------------------------
# Seal backend protocol
# ---------------------------------------------------------------------------

@runtime_checkable
class SealBackend(Protocol):
    """Protocol for sealing step envelopes into proof artifacts.

    Two implementations provided:
      - SbnSealBackend (SnapChore seal + SmartBlock publish)
      - LocalSealBackend (SHA-256 local, no SBN dependency)

    Products can implement custom backends for specialized sealing.
    """
    def seal_envelope(
        self,
        envelope: StepEnvelope,
        *,
        status: str,
        domain: str,
        block_type: str,
        metrics: Optional[Dict[str, Any]],
    ) -> Dict[str, Any]:
        """Seal an envelope. Returns {"proof_hash": ..., "block_ref": ...}."""
        ...


# ---------------------------------------------------------------------------
# Local seal backend (no SBN dependency)
# ---------------------------------------------------------------------------

class LocalSealBackend:
    """Local-only sealing via SHA-256.  No network calls, no SBN.

    Suitable for development, testing, and degraded-mode operation.
    Proofs are valid but not SnapChore-attested.
    """

    def __init__(self, slot_manager: Any = None) -> None:
        self._slot_manager = slot_manager

    def seal_envelope(
        self,
        envelope: StepEnvelope,
        *,
        status: str = "completed",
        domain: str = "general",
        block_type: str = "event",
        metrics: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        payload = envelope.to_seal_payload()
        payload["status"] = status
        payload["domain"] = domain
        if metrics:
            payload["metrics"] = metrics
        raw = json.dumps(payload, sort_keys=True, default=str, separators=(",", ":"))
        proof_hash = hashlib.sha256(raw.encode()).hexdigest()
        block_ref = f"local:{uuid.uuid4().hex[:12]}"
        return {"proof_hash": proof_hash, "block_ref": block_ref}


# ---------------------------------------------------------------------------
# SBN seal backend (production)
# ---------------------------------------------------------------------------

class SbnSealBackend:
    """Production sealing via SnapChore seal + SmartBlock publish.

    Requires an SBN client with `snapchore` and `blocks` sub-clients.
    """

    def __init__(
        self,
        sbn_client: Any,
        slot_id: Optional[str] = None,
    ) -> None:
        self._sbn = sbn_client
        self._slot_id = slot_id

    def seal_envelope(
        self,
        envelope: StepEnvelope,
        *,
        status: str = "completed",
        domain: str = "general",
        block_type: str = "event",
        metrics: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        payload = envelope.to_seal_payload()
        payload["status"] = status
        if metrics:
            payload["metrics"] = metrics

        # SnapChore seal
        proof_hash = ""
        try:
            seal_result = self._sbn.snapchore.seal(
                payload,
                domain=domain,
                block_type=block_type,
                metadata={"step_id": envelope.step_id},
                metrics=metrics,
            )
            proof_hash = seal_result.get("hash", seal_result.get("snapchore_hash", ""))
        except Exception as exc:
            logger.warning("SnapChore seal failed for %s: %s", envelope.step_id, exc)
            raw = json.dumps(payload, sort_keys=True, default=str, separators=(",", ":"))
            proof_hash = hashlib.sha256(raw.encode()).hexdigest()

        # SmartBlock publish
        block_ref = ""
        if hasattr(self._sbn, "blocks"):
            try:
                block_payload = {
                    "step_id": envelope.step_id,
                    "proof_hash": proof_hash,
                    "status": status,
                    "domain": domain,
                }
                block_result = self._sbn.blocks.submit(
                    slot_id=self._slot_id or "",
                    payload=block_payload,
                    domain=domain,
                    metadata={"step_id": envelope.step_id},
                )
                block_ref = block_result.get("block_id", block_result.get("hash", ""))
            except Exception as exc:
                logger.warning("SmartBlock publish failed for %s: %s", envelope.step_id, exc)
                block_ref = f"unpublished:{uuid.uuid4().hex[:12]}"
        else:
            block_ref = f"no-blocks:{uuid.uuid4().hex[:12]}"

        return {"proof_hash": proof_hash, "block_ref": block_ref}


# ---------------------------------------------------------------------------
# ProofWriter
# ---------------------------------------------------------------------------

class ProofWriter:
    """Writes finalized envelopes into ProofArtifacts via a SealBackend.

    The writer doesn't know how sealing works — that's the backend's job.
    It just takes a finalized envelope and produces a ProofArtifact.
    """

    def __init__(self, backend: Optional[SealBackend] = None) -> None:
        self._backend = backend or LocalSealBackend()

    def write(
        self,
        envelope: StepEnvelope,
        *,
        status: ProofStatus = ProofStatus.COMPLETED,
        domain: str = "general",
        block_type: str = "event",
        metrics: Optional[Dict[str, Any]] = None,
    ) -> ProofArtifact:
        """Seal an envelope and produce a ProofArtifact."""
        result = self._backend.seal_envelope(
            envelope,
            status=status.value,
            domain=domain,
            block_type=block_type,
            metrics=metrics,
        )
        return ProofArtifact(
            step_id=envelope.step_id,
            spec_version=envelope.spec_version,
            inputs_fingerprint=envelope.inputs_fingerprint,
            outputs_fingerprint=envelope.outputs_fingerprint,
            proof_hash=result.get("proof_hash", ""),
            block_ref=result.get("block_ref", ""),
            status=status,
            metrics=metrics or {},
        )


# ---------------------------------------------------------------------------
# StepRunner
# ---------------------------------------------------------------------------

class StepRunner:
    """Wraps business logic with the seal-or-fail contract.

    Usage:
        runner = StepRunner(
            envelope_builder=StepEnvelopeBuilder(snapchore),
            proof_writer=ProofWriter(SbnSealBackend(sbn_client)),
            domain="finance.payroll",
        )

        proof = await runner.run(
            step_id="reserve_float",
            spec_version="1.0",
            step_fn=my_reserve_float_fn,
            context=StepContext(inputs={...}),
        )
    """

    def __init__(
        self,
        envelope_builder: StepEnvelopeBuilder,
        proof_writer: ProofWriter,
        gate: Optional[IdempotencyGate] = None,
        *,
        domain: str = "general",
    ) -> None:
        self._builder = envelope_builder
        self._writer = proof_writer
        self._gate = gate or IdempotencyGate()
        self._domain = domain

    async def run(
        self,
        step_id: str,
        spec_version: str,
        step_fn: StepFn,
        context: StepContext,
        *,
        metrics: Optional[Dict[str, Any]] = None,
        domain: Optional[str] = None,
    ) -> ProofArtifact:
        """Execute a step function and produce a ProofArtifact.

        1. Build envelope (inputs fingerprinted)
        2. Check idempotency gate (short-circuit if seen)
        3. Execute step function
        4. Finalize envelope (outputs fingerprinted)
        5. Seal via ProofWriter
        6. Record in gate
        """
        step_domain = domain or self._domain

        # Build envelope
        envelope = self._builder.build(
            step_id,
            spec_version,
            context.inputs,
            config=context.config,
            actor=context.actor,
        )

        # Idempotency check
        prior = self._gate.check(
            envelope.inputs_fingerprint,
            step_id,
            spec_version,
        )
        if prior is not None:
            logger.debug("Idempotency hit for %s — returning cached proof", step_id)
            return prior

        # Execute — re-raise WaitForEvent so the DAG runner can handle it
        from .dag import WaitForEvent
        try:
            outputs = await step_fn(context)
            envelope = self._builder.finalize(envelope, outputs)
            proof = self._writer.write(
                envelope,
                status=ProofStatus.COMPLETED,
                domain=step_domain,
                metrics=metrics,
            )
        except WaitForEvent:
            raise  # Let the DAG runner handle pause/resume
        except Exception as exc:
            # Seal a FAILED proof — the step happened, it just failed
            logger.warning("Step %s failed: %s", step_id, exc)
            envelope = self._builder.finalize(envelope, {"error": str(exc)})
            proof = self._writer.write(
                envelope,
                status=ProofStatus.FAILED,
                domain=step_domain,
                metrics={"error": str(exc), **(metrics or {})},
            )

        # Record
        self._gate.store(proof)
        return proof
