"""DeepSigma LangChain Tools — expose governance operations as LangChain-compatible tools.

Four tools:
    IRISQueryTool       — query the IRIS provenance engine (WHY / WHAT_DRIFTED / STATUS / ...)
    CoherenceScorerTool — score a set of episodes for coherence (0–100, A–F)
    AuditTool           — cross-artifact consistency audit with findings
    ReconcilerTool      — detect inconsistencies and propose repairs

Usage::

    from langchain_deepsigma import deepsigma_tools, DeepSigmaCallbackHandler

    handler = DeepSigmaCallbackHandler()
    # ... run chain with handler ...

    tools = deepsigma_tools(
        mg=handler.mg,
        dlr=handler.dlr,
        rs=handler.rs,
        ds=handler.ds,
    )
    # Bind tools to an agent
    agent = create_tool_calling_agent(llm, tools, prompt)
"""
from __future__ import annotations

import json
import uuid
from typing import Any, Dict, List, Optional, Type

from langchain_core.tools import BaseTool
from pydantic import BaseModel, Field

from coherence_ops import (
    CoherenceAuditor,
    CoherenceScorer,
    DLRBuilder,
    DriftSignalCollector,
    IRISEngine,
    IRISQuery,
    MemoryGraph,
    QueryType,
    Reconciler,
    ReflectionSession,
)
from coherence_ops.manifest import ArtifactDeclaration, ArtifactKind, ComplianceLevel, CoherenceManifest


# ---------------------------------------------------------------------------
# Input schemas (Pydantic v2 compatible via langchain_core shim)
# ---------------------------------------------------------------------------

class IRISQueryInput(BaseModel):
    query_type: str = Field(
        default="STATUS",
        description="One of: WHY, WHAT_CHANGED, WHAT_DRIFTED, RECALL, STATUS",
    )
    target: str = Field(
        default="",
        description="Episode ID or Claim ID to query about (required for WHY / WHAT_CHANGED / WHAT_DRIFTED)",
    )
    text: str = Field(
        default="",
        description="Free-text search term (used with RECALL)",
    )


class EpisodesInput(BaseModel):
    episodes: str = Field(
        description="JSON array of sealed decision episodes",
    )
    drift_events: str = Field(
        default="[]",
        description="JSON array of drift event objects (optional)",
    )


# ---------------------------------------------------------------------------
# IRISQueryTool
# ---------------------------------------------------------------------------

class IRISQueryTool(BaseTool):
    """Query the DeepSigma IRIS provenance engine.

    Answers questions like:
    - WHY was episode ep-001 decided the way it was?
    - WHAT_DRIFTED in the last set of decisions?
    - STATUS — what is the current system health?
    - RECALL — what do we know about past decisions?
    - WHAT_CHANGED — what patches or supersessions exist?
    """

    name: str = "iris_query"
    description: str = (
        "Query the DeepSigma IRIS engine to retrieve governance provenance. "
        "Use query_type=WHY with a target episode ID to understand why a decision was made. "
        "Use WHAT_DRIFTED to find stale assumptions. "
        "Use STATUS for system health. "
        "Use RECALL for general memory lookup. "
        "Use WHAT_CHANGED to find patches."
    )
    args_schema: Type[BaseModel] = IRISQueryInput

    _mg: Optional[MemoryGraph] = None
    _dlr: Optional[DLRBuilder] = None

    def __init__(self, mg: Optional[MemoryGraph] = None, dlr: Optional[DLRBuilder] = None, **kwargs: Any) -> None:
        super().__init__(**kwargs)
        self._mg = mg or MemoryGraph()
        self._dlr = dlr or DLRBuilder()

    def _run(
        self,
        query_type: str = "STATUS",
        target: str = "",
        text: str = "",
    ) -> str:
        engine = IRISEngine(memory_graph=self._mg, dlr_entries=self._dlr.to_dict_list())
        query = IRISQuery(
            query_type=query_type.upper(),
            episode_id=target if not target.startswith("CLAIM") else "",
            claim_id=target if target.startswith("CLAIM") else "",
            text=text,
        )
        response = engine.resolve(query)
        lines = [
            f"[IRIS {response.query_type}] Status: {response.status}",
            f"Summary: {response.summary}",
        ]
        if response.warnings:
            lines.append(f"Warnings: {'; '.join(response.warnings)}")
        if response.provenance:
            lines.append(f"Provenance chain: {len(response.provenance)} node(s)")
        lines.append(f"({response.elapsed_ms:.0f}ms)")
        return "\n".join(lines)

    async def _arun(self, **kwargs: Any) -> str:
        return self._run(**kwargs)


# ---------------------------------------------------------------------------
# CoherenceScorerTool
# ---------------------------------------------------------------------------

class CoherenceScorerTool(BaseTool):
    """Score a set of decision episodes for governance coherence.

    Returns an overall score (0–100) with letter grade (A–F) and a
    per-dimension breakdown across:
    - Policy adherence
    - Outcome health
    - Drift control
    - Memory completeness
    """

    name: str = "coherence_score"
    description: str = (
        "Score a JSON array of sealed decision episodes for coherence (0–100, A–F). "
        "Returns policy adherence, outcome health, drift control, and memory completeness scores. "
        "Input: episodes (JSON array), drift_events (optional JSON array)."
    )
    args_schema: Type[BaseModel] = EpisodesInput

    def _run(self, episodes: str, drift_events: str = "[]") -> str:
        try:
            ep_list: List[Dict[str, Any]] = json.loads(episodes)
            drift_list: List[Dict[str, Any]] = json.loads(drift_events)
        except json.JSONDecodeError as exc:
            return f"Error: invalid JSON input — {exc}"

        dlr = DLRBuilder()
        dlr.from_episodes(ep_list)

        rs = ReflectionSession(f"rs-score-{uuid.uuid4().hex[:8]}")
        rs.ingest(ep_list)

        ds = DriftSignalCollector()
        if drift_list:
            ds.ingest(drift_list)

        mg = MemoryGraph()
        for ep in ep_list:
            mg.add_episode(ep)
        for d in drift_list:
            mg.add_drift(d)

        scorer = CoherenceScorer(dlr_builder=dlr, rs=rs, ds=ds, mg=mg)
        report = scorer.score()

        lines = [
            f"Coherence Score: {report.overall_score:.1f}/100  Grade: {report.grade}",
            f"Computed: {report.computed_at}",
            "",
            "Dimensions:",
        ]
        for dim in report.dimensions:
            lines.append(f"  {dim.name:25s} {dim.score:.1f}  (weight {dim.weight:.0%})")
        return "\n".join(lines)

    async def _arun(self, **kwargs: Any) -> str:
        return self._run(**kwargs)


# ---------------------------------------------------------------------------
# AuditTool
# ---------------------------------------------------------------------------

class AuditTool(BaseTool):
    """Run a cross-artifact consistency audit on decision episodes.

    Checks DLR completeness, drift resolution, memory graph orphans,
    verification consistency, and manifest coverage.
    """

    name: str = "coherence_audit"
    description: str = (
        "Run a cross-artifact consistency audit on a JSON array of decision episodes. "
        "Returns pass/fail with a list of findings (info / warning / critical). "
        "Input: episodes (JSON array), drift_events (optional JSON array)."
    )
    args_schema: Type[BaseModel] = EpisodesInput

    def _run(self, episodes: str, drift_events: str = "[]") -> str:
        try:
            ep_list: List[Dict[str, Any]] = json.loads(episodes)
            drift_list: List[Dict[str, Any]] = json.loads(drift_events)
        except json.JSONDecodeError as exc:
            return f"Error: invalid JSON input — {exc}"

        dlr = DLRBuilder()
        dlr.from_episodes(ep_list)

        rs = ReflectionSession(f"rs-audit-{uuid.uuid4().hex[:8]}")
        rs.ingest(ep_list)

        ds = DriftSignalCollector()
        if drift_list:
            ds.ingest(drift_list)

        mg = MemoryGraph()
        for ep in ep_list:
            mg.add_episode(ep)
        for d in drift_list:
            mg.add_drift(d)

        manifest = _make_manifest()
        auditor = CoherenceAuditor(manifest=manifest, dlr_builder=dlr, rs=rs, ds=ds, mg=mg)
        report = auditor.run(f"audit-lc-{uuid.uuid4().hex[:8]}")

        result_label = "PASSED" if report.passed else "FAILED"
        lines = [
            f"Audit {report.audit_id}: {result_label}",
            f"Run at: {report.run_at}",
            f"Findings: {len(report.findings)} ({report.summary})",
            "",
        ]
        for f in report.findings:
            lines.append(f"  [{f.severity.upper()}] {f.check_name}: {f.message}")
        if not report.findings:
            lines.append("  No issues found.")
        return "\n".join(lines)

    async def _arun(self, **kwargs: Any) -> str:
        return self._run(**kwargs)


# ---------------------------------------------------------------------------
# ReconcilerTool
# ---------------------------------------------------------------------------

class ReconcilerTool(BaseTool):
    """Detect cross-artifact inconsistencies and propose repairs.

    Finds missing Memory Graph nodes, orphan drift events, missing policy
    stamps, and unresolved red drift signals, then proposes repair actions.
    """

    name: str = "coherence_reconcile"
    description: str = (
        "Detect inconsistencies across DeepSigma governance artifacts and propose repairs. "
        "Returns a list of repair proposals (auto-fixable or manual). "
        "Input: episodes (JSON array), drift_events (optional JSON array)."
    )
    args_schema: Type[BaseModel] = EpisodesInput

    def _run(self, episodes: str, drift_events: str = "[]") -> str:
        try:
            ep_list: List[Dict[str, Any]] = json.loads(episodes)
            drift_list: List[Dict[str, Any]] = json.loads(drift_events)
        except json.JSONDecodeError as exc:
            return f"Error: invalid JSON input — {exc}"

        dlr = DLRBuilder()
        dlr.from_episodes(ep_list)

        ds = DriftSignalCollector()
        if drift_list:
            ds.ingest(drift_list)

        mg = MemoryGraph()
        for ep in ep_list:
            mg.add_episode(ep)
        for d in drift_list:
            mg.add_drift(d)

        recon = Reconciler(dlr_builder=dlr, ds=ds, mg=mg)
        result = recon.reconcile()

        lines = [
            f"Reconciliation complete: {len(result.proposals)} proposal(s)",
            f"  Auto-fixable: {result.auto_fixable_count}",
            f"  Manual:       {result.manual_count}",
            f"  Run at:       {result.run_at}",
            "",
        ]
        for p in result.proposals:
            tag = "[AUTO]" if p.auto_fixable else "[MANUAL]"
            lines.append(f"  {tag} {p.kind.value}: {p.description}")
        if not result.proposals:
            lines.append("  No inconsistencies found.")
        return "\n".join(lines)

    async def _arun(self, **kwargs: Any) -> str:
        return self._run(**kwargs)


# ---------------------------------------------------------------------------
# Convenience factory
# ---------------------------------------------------------------------------

def deepsigma_tools(
    mg: Optional[MemoryGraph] = None,
    dlr: Optional[DLRBuilder] = None,
    rs: Optional[ReflectionSession] = None,
    ds: Optional[DriftSignalCollector] = None,
) -> List[BaseTool]:
    """Return all four DeepSigma tools pre-wired to the provided artifacts.

    Typically called after a chain run to give an agent visibility into
    the governance state::

        handler = DeepSigmaCallbackHandler()
        chain.invoke(..., config={"callbacks": [handler]})
        tools = deepsigma_tools(mg=handler.mg, dlr=handler.dlr)

    Args:
        mg: MemoryGraph instance (used by IRISQueryTool).
        dlr: DLRBuilder instance (used by IRISQueryTool).
        rs: ReflectionSession (currently unused — included for future extension).
        ds: DriftSignalCollector (currently unused — included for future extension).

    Returns:
        List of [IRISQueryTool, CoherenceScorerTool, AuditTool, ReconcilerTool].
    """
    return [
        IRISQueryTool(mg=mg, dlr=dlr),
        CoherenceScorerTool(),
        AuditTool(),
        ReconcilerTool(),
    ]


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

def _make_manifest() -> CoherenceManifest:
    manifest = CoherenceManifest(
        system_id="langchain-deepsigma",
        version="0.1.0",
    )
    for kind in (ArtifactKind.DLR, ArtifactKind.RS, ArtifactKind.DS, ArtifactKind.MG):
        manifest.declare(ArtifactDeclaration(
            kind=kind,
            schema_version="1.0",
            compliance=ComplianceLevel.FULL,
            source="langchain_deepsigma.callbacks",
        ))
    return manifest
