from __future__ import annotations
from dataclasses import dataclass
from typing import Any, Callable, Dict, Optional

from .context import SentinelContext
from ..control.guardrails import Guardrails
from ..control.decision import DecisionPolicy, ConfidenceInputs, DecisionOutcome
from ..audit.event import AuditEvent
from ..audit.store import JsonlAuditStore, AuditStore

@dataclass
class SentinelResult:
    outcome: DecisionOutcome
    confidence: float
    output: Any
    correlation_id: str
    run_id: str
    ctx_snapshot: Dict[str, Any]

class SentinelRunner:
    def __init__(self, *, name: str = "sentinel_run",
                 guardrails: Optional[Guardrails] = None,
                 decision_policy: Optional[DecisionPolicy] = None,
                 audit_store: Optional[AuditStore] = None,
                 audit_hash_chain: bool = True,
                 metadata: Optional[Dict[str, Any]] = None):
        self.name = name
        self.guardrails = guardrails or Guardrails()
        self.decision_policy = decision_policy or DecisionPolicy()
        self.audit_store = audit_store or JsonlAuditStore(hash_chain=audit_hash_chain)
        self.audit_hash_chain = audit_hash_chain
        self.metadata = metadata or {}

    def _log(self, ctx: SentinelContext, event_type: str, payload: Dict[str, Any]) -> None:
        prev = self.audit_store.last_hash(ctx.correlation_id)
        event = AuditEvent.create(
            correlation_id=ctx.correlation_id,
            run_id=ctx.run_id,
            event_type=event_type,
            payload=payload,
            prev_hash=prev,
            hash_chain=self.audit_hash_chain,
        )
        self.audit_store.append(event)

    def run(self, agent_fn: Callable[[SentinelContext, Any], Any], input_data: Any,
            *, confidence_inputs: Optional[ConfidenceInputs] = None) -> SentinelResult:
        ctx = SentinelContext(metadata={"name": self.name, **(self.metadata or {})})
        self._log(ctx, "run_start", {"input_type": str(type(input_data)), "ctx": ctx.snapshot()})

        try:
            ctx.bump_step()
            self.guardrails.validate_limits(ctx.state.steps, ctx.state.tool_calls, ctx.state.cost_usd)

            output = agent_fn(ctx, input_data)
            self._log(ctx, "agent_output", {"output_type": str(type(output))})

            ci = confidence_inputs or ConfidenceInputs()
            conf = self.decision_policy.score(ci)
            outcome = self.decision_policy.decide(ci)

            self._log(ctx, "decision", {"confidence": conf, "outcome": outcome.value, "inputs": ci.__dict__})
            self._log(ctx, "run_end", {"status": "ok", "ctx": ctx.snapshot()})

            return SentinelResult(
                outcome=outcome,
                confidence=conf,
                output=output,
                correlation_id=ctx.correlation_id,
                run_id=ctx.run_id,
                ctx_snapshot=ctx.snapshot(),
            )
        except Exception as e:
            self._log(ctx, "run_error", {"error_type": type(e).__name__, "message": str(e), "ctx": ctx.snapshot()})
            self._log(ctx, "run_end", {"status": "error", "ctx": ctx.snapshot()})
            raise
