#!/usr/bin/env python3
"""16 — Context graphs: trace agent decisions into a graph database.

Demonstrates the full context graph lifecycle:
1. Run a KOAN agent with tools (stock prices, HR lookups)
2. Capture every decision, tool call, and outcome into FalkorDB
3. Query the graph for decision traces, precedent search, and gap detection
4. Show how a second run can cite a prior decision as precedent

Architecture:
    Agent (ReactAgent) → MCP Server (tools)
         ↓  decisions, tool calls, outcomes
    DecisionTracer → FalkorDB (context graph)
         ↓  Cypher queries
    Precedent Search / Trace / Gap Detection

Requirements:
    pip install koan[llm,mcp,graph-falkordb]
    export OPENAI_API_KEY=sk-...
    docker compose -f docker/docker-compose.graph.yml --profile falkordb up -d
"""

from __future__ import annotations

import asyncio
import sys
import time
import uuid
from contextlib import AsyncExitStack

from mcp.client.session import ClientSession
from mcp.client.stdio import StdioServerParameters, stdio_client

from koan.agents import AgentConfig, ReactAgent
from koan.graph import (
    DecisionTracer,
    GraphConfig,
    PrecedentSearch,
    create_graph_client,
    detect_policy_gaps,
    get_run_trace,
)
from koan.llm.client import LLMClient
from koan.mcp import MCPToolSource

# --- Config ---

GRAPH_CONFIG = GraphConfig(
    provider="falkordb",
    uri="redis://localhost:6379",
    database="koan_context_graph_demo",
)


async def build_mcp_agent(
    stack: AsyncExitStack,
    client: LLMClient,
    server_script: str,
    agent_name: str,
    system_prompt: str,
) -> tuple[ReactAgent, MCPToolSource]:
    """Launch an MCP server and create a connected agent."""
    params = StdioServerParameters(command=sys.executable, args=[server_script])
    transport = await stack.enter_async_context(stdio_client(params))
    session = await stack.enter_async_context(ClientSession(*transport))
    await session.initialize()

    source = MCPToolSource(session)
    await source.discover()

    agent = ReactAgent(
        agent_name,
        client=client,
        config=AgentConfig(
            provider="openai",
            model="gpt-4o-mini",
            system_prompt=system_prompt,
        ),
        tools=source,
    )
    return agent, source


async def traced_agent_run(
    agent: ReactAgent,
    query: str,
    *,
    graph_client: object,
    thread_id: str,
    entity_refs: list[tuple[str, str, str]] | None = None,
    policy_id: str | None = None,
    precedent_decision_id: str | None = None,
) -> tuple[str, str]:
    """Run an agent with full decision tracing.

    Returns (run_id, first_decision_id) for downstream use.
    """
    run_id = str(uuid.uuid4())

    async with DecisionTracer(
        graph_client,  # type: ignore[arg-type]
        run_id=run_id,
        agent_id=agent.name,
        thread_id=thread_id,
    ) as tracer:
        # 1. Record routing decision
        routing_id = await tracer.record_decision(
            description=f"Processing query: {query[:80]}",
            decision_type="routing",
            confidence=1.0,
            reasoning="Direct agent invocation",
        )

        # 2. Run the agent
        start = time.monotonic()
        result = await agent.run(query)
        latency = int((time.monotonic() - start) * 1000)

        # 3. Record tool selection decision (if tools were called)
        tool_decision_id = routing_id
        if result.tool_calls_made > 0:
            tool_decision_id = await tracer.record_decision(
                description=f"Selected {result.tool_calls_made} tool(s) to answer",
                decision_type="tool_selection",
                confidence=0.9,
                reasoning=f"Agent used {result.tool_calls_made} tool call(s) across {result.iterations} iterations",
                policy_id=policy_id,
                precedent_decision_id=precedent_decision_id,
            )

            # Record the tool call action
            await tracer.record_action(
                decision_id=tool_decision_id,
                action_type="tool_call",
                tool_name=f"{result.tool_calls_made} tool(s)",
                latency_ms=latency,
            )

        # 4. Record output decision
        output_decision_id = await tracer.record_decision(
            description=f"Generated response: {result.output[:100]}",
            decision_type="output",
            confidence=0.85,
            reasoning="LLM produced final answer from tool results",
        )

        # 5. Record the LLM call action
        await tracer.record_action(
            decision_id=output_decision_id,
            action_type="llm_call",
            tool_name=None,
            latency_ms=latency,
        )

        # 6. Link entities if provided
        if entity_refs:
            for entity_id, entity_type, source_system in entity_refs:
                await tracer.link_entity(
                    decision_id=tool_decision_id,
                    entity_id=entity_id,
                    entity_type=entity_type,
                    source_system=source_system,
                )

        # 7. Record outcome
        status = "failure" if result.is_error else "success"
        await tracer.record_outcome(status=status, result=result.output[:200])

    return run_id, tool_decision_id


async def demo_traced_agent_runs() -> dict[str, str]:
    """Run agents with decision tracing and return run metadata."""
    print("=" * 60)
    print("PART 1: Traced Agent Runs")
    print("=" * 60)

    run_ids: dict[str, str] = {}

    async with AsyncExitStack() as stack:
        llm = await stack.enter_async_context(LLMClient())
        graph = await stack.enter_async_context(create_graph_client(GRAPH_CONFIG))

        # Create finance agent connected to MCP server
        finance_agent, finance_src = await build_mcp_agent(
            stack,
            llm,
            "examples/15_mcp_finance_server.py",
            "finance-agent",
            "You are a finance specialist. Use tools to answer. Be concise.",
        )
        print(f"\nFinance tools: {[t.name for t in finance_src.definitions()]}")

        # Create HR agent connected to MCP server
        hr_agent, hr_src = await build_mcp_agent(
            stack,
            llm,
            "examples/15_mcp_hr_server.py",
            "hr-agent",
            "You are an HR specialist. Use tools to answer. Be concise.",
        )
        print(f"HR tools: {[t.name for t in hr_src.definitions()]}")

        # --- Run 1: Finance query ---
        print("\n--- Run 1: Stock price lookup ---")
        run_id_1, decision_id_1 = await traced_agent_run(
            finance_agent,
            "What is the stock price of AAPL?",
            graph_client=graph,
            thread_id="session-001",
            entity_refs=[("AAPL", "stock", "market_data")],
        )
        run_ids["finance_1"] = run_id_1
        print(f"  Run ID: {run_id_1[:8]}...")

        # --- Run 2: HR query ---
        print("\n--- Run 2: Employee lookup ---")
        run_id_2, decision_id_2 = await traced_agent_run(
            hr_agent,
            "Look up employee Alice and check her PTO balance.",
            graph_client=graph,
            thread_id="session-001",
            entity_refs=[
                ("alice", "employee", "hr_system"),
                ("team-alpha", "team", "hr_system"),
            ],
        )
        run_ids["hr_1"] = run_id_2
        print(f"  Run ID: {run_id_2[:8]}...")

        # --- Run 3: Second finance query — cites Run 1 as precedent ---
        print("\n--- Run 3: Portfolio lookup (cites precedent from Run 1) ---")
        run_id_3, _ = await traced_agent_run(
            finance_agent,
            "Show me Alice's portfolio.",
            graph_client=graph,
            thread_id="session-001",
            entity_refs=[("alice", "employee", "hr_system")],
            precedent_decision_id=decision_id_1,
        )
        run_ids["finance_2"] = run_id_3
        print(f"  Run ID: {run_id_3[:8]}...")

        # --- Run 4: Policy-governed decision ---
        # First create a policy node
        await graph.execute_write(
            """
            CREATE (p:Policy {
                policy_id: 'pol-data-access',
                name: 'Data Access Policy',
                version: '1.0',
                description: 'All portfolio lookups must be logged'
            })
            """,
        )
        print("\n--- Run 4: Currency conversion (with policy governance) ---")
        run_id_4, _ = await traced_agent_run(
            finance_agent,
            "Convert 500 USD to GBP.",
            graph_client=graph,
            thread_id="session-002",
            policy_id="pol-data-access",
        )
        run_ids["finance_3"] = run_id_4
        print(f"  Run ID: {run_id_4[:8]}...")

    return run_ids


async def demo_graph_queries(run_ids: dict[str, str]) -> None:
    """Query the context graph for traces, precedents, and gaps."""
    print("\n" + "=" * 60)
    print("PART 2: Querying the Context Graph")
    print("=" * 60)

    async with create_graph_client(GRAPH_CONFIG) as graph:
        # --- 1. Full decision trace for Run 1 ---
        print("\n--- Decision trace for Run 1 (finance) ---")
        trace = await get_run_trace(graph, run_ids["finance_1"])
        for record in trace:
            d = record.data
            print(f"  [{d.get('decision_type', '?')}] {d.get('description', '')}")
            if d.get("tool_name"):
                print(f"    -> Tool: {d['tool_name']}")
            if d.get("entity_id"):
                print(f"    -> Entity: {d['entity_id']} ({d.get('entity_type', '')})")

        # --- 2. Full trace for Run 2 (HR) ---
        print("\n--- Decision trace for Run 2 (HR) ---")
        trace = await get_run_trace(graph, run_ids["hr_1"])
        for record in trace:
            d = record.data
            print(f"  [{d.get('decision_type', '?')}] {d.get('description', '')}")
            if d.get("entity_id"):
                print(f"    -> Entity: {d['entity_id']} ({d.get('entity_type', '')})")

        # --- 3. Precedent search ---
        print("\n--- Precedent search: past decisions for entity 'alice' ---")
        search = PrecedentSearch(graph)
        precedents = await search.find(
            entity_id="alice",
            decision_type="tool_selection",
            limit=5,
        )
        for p in precedents:
            print(f"  [{p.decision_id[:8]}...] {p.description}")
            print(f"    Confidence: {p.confidence}, Outcome: {p.outcome_status}")

        # --- 4. Cross-entity query ---
        print("\n--- All agents that ran in this session ---")
        agents = await graph.execute(
            """
            MATCH (a:Agent)-[:EXECUTED]->(r:Run)
            RETURN a.agent_id AS agent, count(r) AS runs
            """,
        )
        for r in agents:
            print(f"  {r.data['agent']}: {r.data['runs']} run(s)")

        # --- 5. Runs with outcomes ---
        print("\n--- Run outcomes ---")
        outcomes = await graph.execute(
            """
            MATCH (a:Agent)-[:EXECUTED]->(r:Run)-[:PRODUCED]->(o:Outcome)
            RETURN a.agent_id AS agent, r.run_id AS run_id, o.status AS status,
                   r.thread_id AS thread
            ORDER BY r.started_at
            """,
        )
        for r in outcomes:
            d = r.data
            print(f"  [{d['agent']}] run {d['run_id'][:8]}... → {d['status']} (thread: {d['thread']})")

        # --- 6. Policy gap detection ---
        print("\n--- Policy gap detection (decisions without governance) ---")
        gaps = await detect_policy_gaps(graph)
        if gaps:
            for r in gaps:
                d = r.data
                print(f"  [GAP] {d.get('description', '')}")
                print(f"    Type: {d.get('decision_type', '')}")
        else:
            print("  No ungoverned policy decisions found.")

        # --- 7. Graph statistics ---
        print("\n--- Graph statistics ---")
        for label in ["Agent", "Run", "Decision", "Action", "Outcome", "Entity", "Policy"]:
            count = await graph.execute(
                f"MATCH (n:{label}) RETURN count(n) AS count"
            )
            n = count[0].data["count"] if count else 0
            print(f"  {label}: {n} nodes")

        edges = await graph.execute(
            "MATCH ()-[r]->() RETURN type(r) AS type, count(r) AS count"
        )
        print("\n  Relationships:")
        for r in edges:
            print(f"    {r.data['type']}: {r.data['count']}")


async def main() -> None:
    # Clean slate
    async with create_graph_client(GRAPH_CONFIG) as graph:
        healthy = await graph.health_check()
        if not healthy:
            print("ERROR: FalkorDB not reachable at redis://localhost:6379")
            print("Start with: docker compose -f docker/docker-compose.graph.yml --profile falkordb up -d")
            return
        # Clear previous demo data
        await graph.execute_write("MATCH (n) DETACH DELETE n")
        print("Connected to FalkorDB. Graph cleared.\n")

    run_ids = await demo_traced_agent_runs()
    await demo_graph_queries(run_ids)


if __name__ == "__main__":
    asyncio.run(main())
