"""DeepSigmaCallbackHandler — LangChain callback integration for Σ OVERWATCH.

Maps LangChain chain/LLM/tool lifecycle events to DeepSigma Decision Episodes,
then runs the full coherence_ops pipeline (DLR → RS → DS → MG) on completion.

Usage::

    from langchain_deepsigma import DeepSigmaCallbackHandler
    from langchain_core.callbacks import CallbackManager

    handler = DeepSigmaCallbackHandler(decision_type="CustomerSupport")
    chain.invoke({"input": "..."}, config={"callbacks": [handler]})

    # After invocation, governance artifacts are available:
    handler.dlr        # DLRBuilder with sealed records
    handler.rs         # ReflectionSession summary
    handler.ds         # DriftSignalCollector summary
    handler.mg         # MemoryGraph with provenance nodes
"""
from __future__ import annotations

import uuid
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional, Union

from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.outputs import LLMResult

from coherence_ops import (
    DLRBuilder,
    DriftSignalCollector,
    MemoryGraph,
    ReflectionSession,
)


def _now() -> str:
    return datetime.now(timezone.utc).isoformat()


class DeepSigmaCallbackHandler(BaseCallbackHandler):
    """LangChain callback handler that wraps chain execution in a DeepSigma Decision Episode.

    Every chain run becomes a sealed episode recorded through the full
    DLR / RS / DS / MG governance pipeline.

    Args:
        decision_type: Label for the decision type (maps to DTE decision type).
            Defaults to the chain class name when available, else "LangChainExecution".
        policy_pack: Optional policy pack dict to stamp on the episode.
        emit_drift: If True, emit a drift signal when the chain errors. Default True.
        verbose: Log extra info. Default False.
    """

    raise_error: bool = False  # LangChain base setting — do not raise inside callbacks

    def __init__(
        self,
        decision_type: str = "LangChainExecution",
        policy_pack: Optional[Dict[str, Any]] = None,
        emit_drift: bool = True,
        verbose: bool = False,
    ) -> None:
        super().__init__()
        self.decision_type = decision_type
        self.policy_pack = policy_pack or {}
        self.emit_drift = emit_drift
        self.verbose = verbose

        # Governance artifact builders — populated after chain end
        self.dlr = DLRBuilder()
        self.rs = ReflectionSession(f"rs-{uuid.uuid4().hex[:8]}")
        self.ds = DriftSignalCollector()
        self.mg = MemoryGraph()

        # Internal episode state accumulated across events
        self._episode: Dict[str, Any] = {}
        self._drift_events: List[Dict[str, Any]] = []
        self._claim_index: Dict[str, int] = {}  # run_id -> claim list index
        self._action_index: Dict[str, int] = {}  # run_id -> action list index

    # ------------------------------------------------------------------
    # Chain events
    # ------------------------------------------------------------------

    def on_chain_start(
        self,
        serialized: Dict[str, Any],
        inputs: Dict[str, Any],
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        chain_name = (
            serialized.get("name")
            or serialized.get("id", [""])[-1]
            or self.decision_type
        )
        self._episode = {
            "episodeId": f"ep-lc-{run_id.hex[:12]}",
            "decisionType": chain_name,
            "startedAt": _now(),
            "actor": {"type": "langchain_chain", "id": chain_name, "version": "unknown"},
            "context": {
                "inputs": str(inputs),
                "evidenceRefs": [],
            },
            "truth": {"claims": []},
            "actions": [],
            "degrade": {"step": "ideal", "rationale": "LangChain execution"},
            "policy": self.policy_pack if self.policy_pack else {},
            "verification": {"result": "na", "verifierId": "langchain"},
            "outcome": {"code": "pending", "detail": ""},
            "seal": {},
        }
        if self.verbose:
            print(f"[DeepSigma] chain_start: {chain_name} (ep {self._episode['episodeId']})")

    def on_chain_end(
        self,
        outputs: Dict[str, Any],
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        tags: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> None:
        self._episode["outcome"] = {"code": "success", "detail": str(outputs)}
        self._episode["endedAt"] = _now()
        self._seal_episode()
        self._run_pipeline()

    def on_chain_error(
        self,
        error: Union[Exception, KeyboardInterrupt],
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        tags: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> None:
        self._episode["outcome"] = {"code": "fail", "detail": str(error)}
        self._episode["endedAt"] = _now()
        if self.emit_drift:
            self._drift_events.append(self._make_drift_event(
                episode_id=self._episode.get("episodeId", "unknown"),
                drift_type="outcome",
                severity="red",
                expected="success",
                actual=str(error),
                fingerprint_key=f"outcome/chain-error/{type(error).__name__}",
            ))
        self._seal_episode()
        self._run_pipeline()

    # ------------------------------------------------------------------
    # LLM events
    # ------------------------------------------------------------------

    def on_llm_start(
        self,
        serialized: Dict[str, Any],
        prompts: List[str],
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        claim = {
            "claimId": f"CLAIM-LC-{run_id.hex[:8].upper()}",
            "statement": prompts[0][:500] if prompts else "",
            "confidence": 0.5,
            "statusLight": "GREEN",
            "capturedAt": _now(),
            "ttlMs": 60000,
            "role": "llm_input",
        }
        claims = self._episode.setdefault("truth", {}).setdefault("claims", [])
        self._claim_index[str(run_id)] = len(claims)
        claims.append(claim)

    def on_llm_end(
        self,
        response: LLMResult,
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        tags: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> None:
        idx = self._claim_index.get(str(run_id))
        if idx is None:
            return
        claims = self._episode.get("truth", {}).get("claims", [])
        if idx >= len(claims):
            return

        output_text = ""
        if response.generations:
            gen = response.generations[0]
            if gen and hasattr(gen[0], "text"):
                output_text = gen[0].text[:500]

        # Estimate confidence from token usage if available
        confidence = 0.85
        if response.llm_output and "token_usage" in response.llm_output:
            usage = response.llm_output["token_usage"]
            if usage.get("total_tokens", 0) > 0:
                confidence = min(0.99, 0.70 + (usage.get("completion_tokens", 0) / 1000))

        claims[idx]["llm_output"] = output_text
        claims[idx]["confidence"] = confidence
        claims[idx]["statusLight"] = "GREEN"
        claims[idx]["completedAt"] = _now()

    def on_llm_error(
        self,
        error: Union[Exception, KeyboardInterrupt],
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        tags: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> None:
        idx = self._claim_index.get(str(run_id))
        if idx is not None:
            claims = self._episode.get("truth", {}).get("claims", [])
            if idx < len(claims):
                claims[idx]["confidence"] = 0.0
                claims[idx]["statusLight"] = "RED"
                claims[idx]["error"] = str(error)

        if self.emit_drift:
            self._drift_events.append(self._make_drift_event(
                episode_id=self._episode.get("episodeId", "unknown"),
                drift_type="verify",
                severity="red",
                expected="llm_response",
                actual=str(error),
                fingerprint_key=f"verify/llm-error/{type(error).__name__}",
            ))

    # ------------------------------------------------------------------
    # Tool events
    # ------------------------------------------------------------------

    def on_tool_start(
        self,
        serialized: Dict[str, Any],
        input_str: str,
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        tags: Optional[List[str]] = None,
        metadata: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> None:
        tool_name = serialized.get("name", "unknown_tool")
        action = {
            "type": tool_name,
            "input": input_str[:300],
            "startedAt": _now(),
            "blastRadiusTier": "low",
            "idempotencyKey": f"tool-{run_id.hex[:12]}",
            "authorization": {"mode": "auto"},
        }
        actions = self._episode.setdefault("actions", [])
        self._action_index[str(run_id)] = len(actions)
        actions.append(action)

    def on_tool_end(
        self,
        output: str,
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        tags: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> None:
        idx = self._action_index.get(str(run_id))
        if idx is None:
            return
        actions = self._episode.get("actions", [])
        if idx < len(actions):
            actions[idx]["output"] = str(output)[:300]
            actions[idx]["completedAt"] = _now()
            actions[idx]["result"] = "success"

    def on_tool_error(
        self,
        error: Union[Exception, KeyboardInterrupt],
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        tags: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> None:
        idx = self._action_index.get(str(run_id))
        if idx is not None:
            actions = self._episode.get("actions", [])
            if idx < len(actions):
                actions[idx]["result"] = "fail"
                actions[idx]["error"] = str(error)

        if self.emit_drift:
            self._drift_events.append(self._make_drift_event(
                episode_id=self._episode.get("episodeId", "unknown"),
                drift_type="outcome",
                severity="yellow",
                expected="tool_success",
                actual=str(error),
                fingerprint_key="outcome/tool-error",
            ))

    # ------------------------------------------------------------------
    # Agent events
    # ------------------------------------------------------------------

    def on_agent_action(
        self,
        action: Any,
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        tags: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> None:
        agent_action = {
            "type": f"agent_action:{getattr(action, 'tool', 'unknown')}",
            "input": str(getattr(action, "tool_input", ""))[:300],
            "startedAt": _now(),
            "blastRadiusTier": "medium",
            "idempotencyKey": f"agent-{run_id.hex[:12]}",
            "authorization": {"mode": "auto"},
        }
        self._episode.setdefault("actions", []).append(agent_action)

    def on_agent_finish(
        self,
        finish: Any,
        *,
        run_id: uuid.UUID,
        parent_run_id: Optional[uuid.UUID] = None,
        tags: Optional[List[str]] = None,
        **kwargs: Any,
    ) -> None:
        self._episode["agentReturn"] = str(getattr(finish, "return_values", ""))[:300]

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _seal_episode(self) -> None:
        import hashlib
        import json

        episode_json = json.dumps(self._episode, sort_keys=True, default=str)
        seal_hash = "sha256:" + hashlib.sha256(episode_json.encode()).hexdigest()
        self._episode["seal"] = {
            "sealHash": seal_hash,
            "sealedAt": _now(),
            "version": 1,
        }
        if self.verbose:
            print(f"[DeepSigma] sealed episode {self._episode['episodeId']} ({self._episode['outcome']['code']})")

    def _run_pipeline(self) -> None:
        episode = self._episode

        self.dlr.from_episodes([episode])

        self.rs.ingest([episode])

        if self._drift_events:
            self.ds.ingest(self._drift_events)

        self.mg.add_episode(episode)
        for drift in self._drift_events:
            self.mg.add_drift(drift)

        if self.verbose:
            summary = self.rs.summarise()
            print(f"[DeepSigma] pipeline complete — RS episodes: {summary.episode_count}")

    def _make_drift_event(
        self,
        episode_id: str,
        drift_type: str,
        severity: str,
        expected: str,
        actual: str,
        fingerprint_key: str,
    ) -> Dict[str, Any]:
        return {
            "driftId": f"drift-lc-{uuid.uuid4().hex[:12]}",
            "episodeId": episode_id,
            "driftType": drift_type,
            "severity": severity,
            "detectedAt": _now(),
            "expectedBehaviour": expected,
            "actualBehaviour": actual,
            "fingerprint": {"key": fingerprint_key, "version": 1},
            "recommendedPatchType": "outcome_review",
            "metadata": {},
        }

    @property
    def episode(self) -> Dict[str, Any]:
        """The sealed episode dict. Available after chain completion."""
        return self._episode

    @property
    def drift_events(self) -> List[Dict[str, Any]]:
        """Drift events emitted during this run."""
        return list(self._drift_events)
