"""DecisionTracer — capture decision traces from agent runs into the graph."""

from __future__ import annotations

import uuid
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from koan.graph.client import GraphClient


class DecisionTracer:
    """Context manager that captures decision traces and persists them to a graph.

    Usage::

        async with DecisionTracer(graph_client, run_id="r-1", agent_id="a-1", thread_id="t-1") as tracer:
            decision_id = await tracer.record_decision(
                description="Classified as billing",
                decision_type="classification",
                confidence=0.92,
                reasoning="Keyword match",
            )
            action_id = await tracer.record_action(
                decision_id=decision_id,
                action_type="tool_call",
                tool_name="lookup_account",
            )
            await tracer.record_outcome(status="success", result="Resolved")
    """

    def __init__(
        self,
        client: GraphClient,
        *,
        run_id: str,
        agent_id: str,
        thread_id: str,
    ) -> None:
        self.client = client
        self.run_id = run_id
        self.agent_id = agent_id
        self.thread_id = thread_id
        self._started_at: datetime | None = None

    async def __aenter__(self) -> DecisionTracer:
        self._started_at = datetime.now(UTC)
        await self.client.execute_write(
            """
            MERGE (a:Agent {agent_id: $agent_id})
            CREATE (r:Run {
                run_id: $run_id,
                thread_id: $thread_id,
                started_at: $started_at,
                status: 'running'
            })
            CREATE (a)-[:EXECUTED]->(r)
            """,
            {
                "agent_id": self.agent_id,
                "run_id": self.run_id,
                "thread_id": self.thread_id,
                "started_at": self._started_at.isoformat(),
            },
        )
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: object,
    ) -> None:
        status = "failed" if exc_type else "completed"
        await self.client.execute_write(
            """
            MATCH (r:Run {run_id: $run_id})
            SET r.status = $status, r.ended_at = $ended_at
            """,
            {
                "run_id": self.run_id,
                "status": status,
                "ended_at": datetime.now(UTC).isoformat(),
            },
        )

    async def record_decision(
        self,
        *,
        description: str,
        decision_type: str,
        confidence: float = 1.0,
        reasoning: str = "",
        inputs: dict[str, Any] | None = None,
        policy_id: str | None = None,
        precedent_decision_id: str | None = None,
    ) -> str:
        """Record a decision node and link it to the current run.

        Returns the generated decision_id.
        """
        decision_id = str(uuid.uuid4())

        await self.client.execute_write(
            """
            MATCH (r:Run {run_id: $run_id})
            CREATE (d:Decision {
                decision_id: $decision_id,
                run_id: $run_id,
                timestamp: $timestamp,
                description: $description,
                decision_type: $decision_type,
                confidence: $confidence,
                reasoning: $reasoning
            })
            CREATE (r)-[:CONTAINS]->(d)
            """,
            {
                "run_id": self.run_id,
                "decision_id": decision_id,
                "timestamp": datetime.now(UTC).isoformat(),
                "description": description,
                "decision_type": decision_type,
                "confidence": confidence,
                "reasoning": reasoning,
            },
        )

        if policy_id:
            await self.client.execute_write(
                """
                MATCH (d:Decision {decision_id: $decision_id})
                MATCH (p:Policy {policy_id: $policy_id})
                CREATE (d)-[:GOVERNED_BY]->(p)
                """,
                {"decision_id": decision_id, "policy_id": policy_id},
            )

        if precedent_decision_id:
            await self.client.execute_write(
                """
                MATCH (d:Decision {decision_id: $decision_id})
                MATCH (prev:Decision {decision_id: $prev_id})
                CREATE (d)-[:CITED_PRECEDENT]->(prev)
                """,
                {"decision_id": decision_id, "prev_id": precedent_decision_id},
            )

        return decision_id

    async def record_action(
        self,
        *,
        decision_id: str,
        action_type: str,
        tool_name: str | None = None,
        input_data: dict[str, Any] | None = None,
        output_data: Any = None,
        latency_ms: int | None = None,
        error: str | None = None,
    ) -> str:
        """Record an action resulting from a decision.

        Returns the generated action_id.
        """
        action_id = str(uuid.uuid4())
        await self.client.execute_write(
            """
            MATCH (d:Decision {decision_id: $decision_id})
            CREATE (a:Action {
                action_id: $action_id,
                type: $type,
                tool_name: $tool_name,
                latency_ms: $latency_ms,
                error: $error
            })
            CREATE (d)-[:LED_TO]->(a)
            """,
            {
                "decision_id": decision_id,
                "action_id": action_id,
                "type": action_type,
                "tool_name": tool_name,
                "latency_ms": latency_ms,
                "error": error,
            },
        )
        return action_id

    async def record_outcome(
        self,
        *,
        status: str,
        result: Any = None,
        evaluation_score: float | None = None,
    ) -> str:
        """Record the outcome of the current run.

        Returns the generated outcome_id.
        """
        outcome_id = str(uuid.uuid4())
        await self.client.execute_write(
            """
            MATCH (r:Run {run_id: $run_id})
            CREATE (o:Outcome {
                outcome_id: $outcome_id,
                run_id: $run_id,
                status: $status,
                evaluation_score: $eval_score
            })
            MERGE (r)-[:PRODUCED]->(o)
            """,
            {
                "run_id": self.run_id,
                "outcome_id": outcome_id,
                "status": status,
                "eval_score": evaluation_score,
            },
        )
        return outcome_id

    async def link_entity(
        self,
        *,
        decision_id: str,
        entity_id: str,
        entity_type: str,
        source_system: str = "",
    ) -> None:
        """Link a business entity to a decision."""
        await self.client.execute_write(
            """
            MERGE (e:Entity {entity_id: $entity_id})
            ON CREATE SET e.entity_type = $entity_type, e.source_system = $source
            WITH e
            MATCH (d:Decision {decision_id: $decision_id})
            CREATE (d)-[:REFERENCED]->(e)
            """,
            {
                "decision_id": decision_id,
                "entity_id": entity_id,
                "entity_type": entity_type,
                "source": source_system,
            },
        )
