"""Example 17 — Graph-Augmented Agent.

Demonstrates the closed read-write loop:
1. Agent queries the decision graph for past context before reasoning
2. Graph context (precedents, entity history, policies) is injected into the system prompt
3. Agent reasons with awareness of its own decision history
4. New decisions are traced back to the graph

Requires:
    - FalkorDB running: docker compose -f docker/docker-compose.graph.yml --profile falkordb up -d
    - OPENAI_API_KEY set in .env
"""

from __future__ import annotations

import asyncio
import sys

from dotenv import load_dotenv

load_dotenv()

from koan.agents import AgentConfig
from koan.graph import (
    DecisionTracer,
    GraphAugmentedAgent,
    GraphConfig,
    create_graph_client,
    get_run_trace,
)
from koan.llm.client import LLMClient
from koan.llm.config import ProviderSettings
from koan.tools import ToolRegistry, tool

# ── Tools ────────────────────────────────────────────────────────────

ACCOUNTS = {
    "acct-001": {"name": "Alice Corp", "balance": 15_000.00, "status": "active", "tier": "gold"},
    "acct-002": {"name": "Bob LLC", "balance": 3_200.50, "status": "active", "tier": "silver"},
    "acct-003": {"name": "Charlie Inc", "balance": -500.00, "status": "overdue", "tier": "bronze"},
}


@tool
def lookup_account(account_id: str) -> str:
    """Look up account details by account ID."""
    acct = ACCOUNTS.get(account_id)
    if not acct:
        return f"Account {account_id} not found."
    return (
        f"Account: {acct['name']} | Balance: ${acct['balance']:,.2f} | "
        f"Status: {acct['status']} | Tier: {acct['tier']}"
    )


@tool
def check_payment_history(account_id: str) -> str:
    """Check the payment history for an account."""
    histories = {
        "acct-001": "12 on-time payments in last 12 months. Perfect record.",
        "acct-002": "10 on-time, 2 late payments in last 12 months.",
        "acct-003": "5 on-time, 4 late, 3 missed payments. Account flagged.",
    }
    return histories.get(account_id, f"No payment history for {account_id}.")


@tool
def apply_credit(account_id: str, amount: float) -> str:
    """Apply a credit to an account balance."""
    acct = ACCOUNTS.get(account_id)
    if not acct:
        return f"Account {account_id} not found."
    return f"Applied ${amount:.2f} credit to {acct['name']}. New balance: ${acct['balance'] + amount:,.2f}"


# ── Seed the graph with prior decisions ──────────────────────────────

async def seed_graph(graph_client):
    """Seed the graph with some past decisions so the agent has history to read."""
    # Run 1: Past billing inquiry for acct-001 (successful)
    async with DecisionTracer(
        graph_client, run_id="seed-run-1", agent_id="billing-agent", thread_id="seed-t-1",
    ) as tracer:
        d1 = await tracer.record_decision(
            description="Classified request as billing inquiry for Alice Corp",
            decision_type="classification",
            confidence=0.95,
            reasoning="Keywords: invoice, billing. Account: acct-001",
        )
        await tracer.record_action(
            decision_id=d1, action_type="tool_call", tool_name="lookup_account",
        )
        await tracer.link_entity(
            decision_id=d1, entity_id="acct-001", entity_type="account", source_system="crm",
        )
        await tracer.record_outcome(
            status="success", result="Resolved billing inquiry. Applied $50 credit.", evaluation_score=0.9,
        )

    # Run 2: Past billing inquiry for acct-003 (escalated)
    async with DecisionTracer(
        graph_client, run_id="seed-run-2", agent_id="billing-agent", thread_id="seed-t-2",
    ) as tracer:
        d2 = await tracer.record_decision(
            description="Classified request as billing dispute for Charlie Inc",
            decision_type="classification",
            confidence=0.85,
            reasoning="Keywords: dispute, overdue. Account: acct-003",
        )
        await tracer.record_action(
            decision_id=d2, action_type="tool_call", tool_name="check_payment_history",
        )
        await tracer.link_entity(
            decision_id=d2, entity_id="acct-003", entity_type="account", source_system="crm",
        )
        await tracer.record_outcome(
            status="escalated", result="Account overdue with missed payments. Escalated to supervisor.",
        )

    # Add a policy
    await graph_client.execute_write(
        """
        MERGE (p:Policy {policy_id: $pid})
        SET p.name = $name, p.description = $desc, p.version = $ver
        """,
        {
            "pid": "pol-credit-limit",
            "name": "Credit Limit Policy",
            "desc": "Credits over $100 require supervisor approval. Overdue accounts cannot receive credits.",
            "ver": "1.0",
        },
    )

    print("[seed] Graph seeded with 2 prior runs + 1 policy")


# ── Main ─────────────────────────────────────────────────────────────

async def main():
    graph_config = GraphConfig(
        provider="falkordb",
        uri="redis://localhost:6379",
        database="koan_augmented_demo",
    )

    registry = ToolRegistry()
    registry.add(lookup_account)
    registry.add(check_payment_history)
    registry.add(apply_credit)

    async with create_graph_client(graph_config) as graph, LLMClient(ProviderSettings()) as llm:
        # Step 1: Seed the graph with past decisions
        print("=" * 60)
        print("STEP 1: Seeding graph with past decisions")
        print("=" * 60)
        await seed_graph(graph)

        # Step 2: Run a graph-augmented agent for acct-001
        print("\n" + "=" * 60)
        print("STEP 2: Graph-augmented agent for acct-001 (has prior history)")
        print("=" * 60)

        agent = GraphAugmentedAgent(
            "billing-agent",
            client=llm,
            config=AgentConfig(
                provider="openai",
                model="gpt-4o",
                system_prompt=(
                    "You are a billing support agent. Look up accounts and help resolve billing issues. "
                    "Use the tools provided. Check payment history before applying credits. "
                    "Follow all active policies."
                ),
            ),
            tools=registry,
            graph_client=graph,
            entity_ids=["acct-001"],
            decision_type="classification",
        )

        result = await agent.run(
            "I'd like to request a $25 credit on account acct-001 for a billing error."
        )

        print(f"\nAgent output:\n{result.output}")
        print(f"\nMetadata: {result.metadata}")

        # Step 3: Run for acct-003 (overdue — agent should see past escalation)
        print("\n" + "=" * 60)
        print("STEP 3: Graph-augmented agent for acct-003 (overdue, past escalation)")
        print("=" * 60)

        agent2 = GraphAugmentedAgent(
            "billing-agent",
            client=llm,
            config=AgentConfig(
                provider="openai",
                model="gpt-4o",
                system_prompt=(
                    "You are a billing support agent. Look up accounts and help resolve billing issues. "
                    "Use the tools provided. Check payment history before applying credits. "
                    "Follow all active policies. If an account was previously escalated, note this."
                ),
            ),
            tools=registry,
            graph_client=graph,
            entity_ids=["acct-003"],
            decision_type="classification",
        )

        result2 = await agent2.run(
            "Can I get a $75 credit on account acct-003? I think there was a billing mistake."
        )

        print(f"\nAgent output:\n{result2.output}")
        print(f"\nMetadata: {result2.metadata}")

        # Step 4: Query the graph to see all decisions
        print("\n" + "=" * 60)
        print("STEP 4: Query all decision traces")
        print("=" * 60)

        for run_id_key in ["run_id"]:
            for r in [result, result2]:
                rid = r.metadata.get(run_id_key, "")
                if rid:
                    trace = await get_run_trace(graph, rid)
                    print(f"\nTrace for run {rid[:8]}...:")
                    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']}")

        # Show graph stats
        stats = await graph.execute(
            "MATCH (n) RETURN labels(n)[0] AS label, count(n) AS count"
        )
        print("\nGraph statistics:")
        for s in stats:
            print(f"  {s.data['label']}: {s.data['count']}")


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        pass
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        print("Make sure FalkorDB is running and OPENAI_API_KEY is set.", file=sys.stderr)
        sys.exit(1)
