"""Demo 5 — Self-Organizing Graph: Full Policy Lifecycle.

A CLI demo that walks through the complete self-organizing policy lifecycle:

  1. Cold start     — no policies, all tools available
  2. Accumulation   — multiple runs build up decision traces in the graph
  3. Emergence      — after enough runs, patterns are discovered and promoted to policies
  4. Tool gating    — new runs see restricted tools based on active policies
  5. Feedback       — user feedback affects confidence calibration
  6. Retirement     — underperforming policies are automatically deleted

Requirements:
  - FalkorDB running:  docker compose -f docker/docker-compose.graph.yml --profile falkordb up -d
  - LLM API key in .env (OPENAI_API_KEY or ANTHROPIC_API_KEY)

Usage:
  uv run python demos/demo5_self_organizing.py
"""

from __future__ import annotations

import asyncio
import json
import sys
from pathlib import Path
from typing import TYPE_CHECKING, Any

sys.path.insert(0, str(Path(__file__).resolve().parent))

from dotenv import load_dotenv

from koan.agents.events import AgentComplete, AgentError, TextDelta, ToolCallEnd, ToolCallStart
from koan.agents.types import AgentConfig
from koan.graph import (
    FeedbackType,
    GraphAugmentedAgent,
    GraphConfig,
    PolicyGatedTools,
    PolicyManager,
    create_graph_client,
    record_feedback,
)
from koan.graph.context import GraphContextProvider
from koan.graph.feedback import ConfidenceCalibrator
from koan.llm.client import LLMClient
from koan.llm.config import ProviderSettings
from koan.tools import ToolRegistry, tool

if TYPE_CHECKING:
    from koan.graph.client import GraphClient

load_dotenv()


# ── Helpers ───────────────────────────────────────────────────────────────

BOLD = "\033[1m"
DIM = "\033[2m"
GREEN = "\033[32m"
YELLOW = "\033[33m"
CYAN = "\033[36m"
RED = "\033[31m"
MAGENTA = "\033[35m"
RESET = "\033[0m"


def banner(text: str) -> None:
    width = 70
    print(f"\n{BOLD}{CYAN}{'═' * width}{RESET}")
    print(f"{BOLD}{CYAN}  {text}{RESET}")
    print(f"{BOLD}{CYAN}{'═' * width}{RESET}\n")


def step(n: int, text: str) -> None:
    print(f"\n{BOLD}{YELLOW}┌─ Step {n}: {text}{RESET}")
    print(f"{YELLOW}│{RESET}")


def info(label: str, value: Any) -> None:
    print(f"{YELLOW}│{RESET}  {DIM}{label}:{RESET} {value}")


def success(text: str) -> None:
    print(f"{YELLOW}│{RESET}  {GREEN}✓ {text}{RESET}")


def warn(text: str) -> None:
    print(f"{YELLOW}│{RESET}  {RED}✗ {text}{RESET}")


def done() -> None:
    print(f"{YELLOW}└{'─' * 50}{RESET}")


# ── Stub tools ────────────────────────────────────────────────────────────

ACCOUNTS: dict[str, dict[str, Any]] = {
    "acct-1001": {"name": "Alice Johnson", "plan": "Pro", "status": "active", "balance": 142.50},
    "acct-1002": {"name": "Bob Smith", "plan": "Starter", "status": "active", "balance": 27.00},
    "acct-1003": {"name": "Carol Davis", "plan": "Enterprise", "status": "suspended", "balance": 980.00},
}


@tool
def lookup_account(account_id: str) -> str:
    """Look up account details by account ID.

    Args:
        account_id: The account identifier, e.g. 'acct-1001'.
    """
    acct = ACCOUNTS.get(account_id)
    if acct is None:
        return f"Account {account_id} not found."
    return json.dumps({"account_id": account_id, **acct})


@tool
def check_balance(account_id: str) -> str:
    """Check the current balance for an account.

    Args:
        account_id: The account identifier, e.g. 'acct-1001'.
    """
    acct = ACCOUNTS.get(account_id)
    if acct is None:
        return f"Account {account_id} not found."
    return f"Account {account_id} balance: ${acct['balance']:.2f} (plan: {acct['plan']})"


@tool
def apply_credit(account_id: str, amount: float) -> str:
    """Apply a credit (discount) to an account balance.

    Args:
        account_id: The account identifier, e.g. 'acct-1001'.
        amount: The dollar amount to credit.
    """
    acct = ACCOUNTS.get(account_id)
    if acct is None:
        return f"Account {account_id} not found."
    old = acct["balance"]
    acct["balance"] = max(0.0, old - amount)
    return f"Applied ${amount:.2f} credit. ${old:.2f} → ${acct['balance']:.2f}"


@tool
def escalate_ticket(account_id: str, reason: str) -> str:
    """Escalate a support ticket to a human supervisor.

    Args:
        account_id: The account identifier.
        reason: Why the ticket should be escalated.
    """
    return f"Ticket escalated for {account_id}: {reason}"


# ── Graph inspection helpers ──────────────────────────────────────────────

async def show_graph_state(gc: GraphClient, label: str) -> None:
    """Print a summary of what's in the graph."""
    print(f"\n{MAGENTA}  ┌─ Graph State: {label}{RESET}")

    # Count nodes by label
    for node_label in ("Agent", "Run", "Decision", "Action", "Entity", "Policy", "Outcome", "Feedback"):
        records = await gc.execute(
            f"MATCH (n:{node_label}) RETURN count(n) AS c", {},
        )
        count = records[0].data.get("c", 0) if records else 0
        if count > 0:
            print(f"{MAGENTA}  │{RESET}  {node_label}: {count}")

    # Show policies in detail
    policies = await gc.execute(
        "MATCH (p:Policy) RETURN p.name AS name, p.source AS source, p.rules AS rules", {},
    )
    if policies:
        print(f"{MAGENTA}  │{RESET}")
        print(f"{MAGENTA}  │  {BOLD}Active Policies:{RESET}")
        for p in policies:
            name = p.data.get("name", "?")
            source = p.data.get("source", "?")
            rules = p.data.get("rules", [])
            print(f"{MAGENTA}  │{RESET}    • {name} ({source})")
            if rules:
                for r in rules:
                    print(f"{MAGENTA}  │{RESET}      → {r}")

    # Show GOVERNED_BY edges
    gov = await gc.execute(
        """
        MATCH (d:Decision)-[:GOVERNED_BY]->(p:Policy)
        RETURN p.name AS policy, count(d) AS decisions
        """,
        {},
    )
    if gov:
        print(f"{MAGENTA}  │{RESET}")
        print(f"{MAGENTA}  │  {BOLD}GOVERNED_BY Links:{RESET}")
        for g in gov:
            print(f"{MAGENTA}  │{RESET}    • {g.data.get('policy', '?')}: {g.data.get('decisions', 0)} decision(s)")

    print(f"{MAGENTA}  └{'─' * 45}{RESET}")


async def show_augmented_prompt(gc: GraphClient, agent_name: str) -> None:
    """Show what the agent would see on next run (context preview)."""
    calibrator = ConfidenceCalibrator(gc)
    provider = GraphContextProvider(gc, calibrator=calibrator)
    ctx = await provider.gather(
        agent_id=agent_name,
        decision_type="billing",
        limit=5,
    )

    prompt = ctx.to_prompt()
    if prompt:
        print(f"\n{DIM}  ┌─ Context that next run will see:{RESET}")
        for line in prompt.strip().split("\n"):
            print(f"{DIM}  │ {line}{RESET}")
        print(f"{DIM}  └{'─' * 45}{RESET}")
    else:
        print(f"\n{DIM}  (No graph context yet — cold start){RESET}")


# ── Run an agent interaction ──────────────────────────────────────────────

async def run_interaction(
    llm: LLMClient,
    gc: GraphClient,
    registry: ToolRegistry,
    query: str,
    entity_ids: list[str],
    *,
    run_number: int,
) -> str:
    """Run one agent interaction and return the run_id."""
    info("Query", f'"{query}"')
    info("Entities", entity_ids)

    agent = GraphAugmentedAgent(
        "billing-agent",
        client=llm,
        config=AgentConfig(
            provider="openai",
            model="gpt-4o-mini",
            temperature=0.2,
            system_prompt=(
                "You are a billing support agent. Help customers with account lookups, "
                "balance inquiries, and applying credits. Be concise and professional. "
                "Use the provided account ID directly — do not ask the user to confirm it."
            ),
        ),
        tools=registry,
        graph_client=gc,
        entity_ids=entity_ids,
        decision_type="billing",
    )

    run_id = ""
    full_response = ""
    tool_calls: list[str] = []

    async for event in agent.run_stream(query):
        if isinstance(event, TextDelta):
            full_response += event.text
        elif isinstance(event, ToolCallStart):
            tool_calls.append(event.tool_name)
        elif isinstance(event, ToolCallEnd):
            pass
        elif isinstance(event, AgentComplete):
            run_id = event.metadata.get("run_id", "")
            gc_meta = event.metadata.get("graph_context", {})
            info("Tools used", tool_calls or "(none)")
            info("Context seen", f"{gc_meta.get('precedents', 0)} precedents, "
                 f"{gc_meta.get('policies', 0)} policies, "
                 f"{gc_meta.get('entities', 0)} entities")
            if gc_meta.get("new_policies"):
                success(f"New policies promoted: {gc_meta['new_policies']}")
            if gc_meta.get("linked_policies"):
                success(f"Decision linked to {len(gc_meta['linked_policies'])} policy(ies)")
            if gc_meta.get("retired_policies"):
                warn(f"Policies retired: {gc_meta['retired_policies']}")
        elif isinstance(event, AgentError):
            warn(f"Error: {event.error}")
            full_response = event.output or ""

    # Truncate for display
    preview = full_response[:200].replace("\n", " ")
    info("Response", f"{preview}{'...' if len(full_response) > 200 else ''}")

    return run_id


# ── Main demo flow ────────────────────────────────────────────────────────

async def main() -> None:
    banner("KOAN — Self-Organizing Policy Lifecycle Demo")

    print(f"{DIM}This demo walks through the full lifecycle:{RESET}")
    print(f"{DIM}  1. Cold start (no policies)         → all tools available{RESET}")
    print(f"{DIM}  2. Accumulate runs (3+ similar)      → patterns discovered{RESET}")
    print(f"{DIM}  3. Policy emergence                  → auto-promoted to active{RESET}")
    print(f"{DIM}  4. Tool gating                       → restricted tools on next run{RESET}")
    print(f"{DIM}  5. Feedback + retirement              → self-correction loop{RESET}")
    print()

    # ── Connect ───────────────────────────────────────────────────────
    graph_config = GraphConfig(
        provider="falkordb",
        uri="redis://localhost:6379",
        database="koan_demo5",
    )

    gc = create_graph_client(graph_config)
    try:
        await gc.connect()
        healthy = await gc.health_check()
        if not healthy:
            print(f"{RED}ERROR: FalkorDB not reachable. Start it with:{RESET}")
            print("  docker compose -f docker/docker-compose.graph.yml --profile falkordb up -d")
            return
    except Exception as e:
        print(f"{RED}ERROR: Cannot connect to FalkorDB: {e}{RESET}")
        print("  docker compose -f docker/docker-compose.graph.yml --profile falkordb up -d")
        return

    success("Connected to FalkorDB")

    # Clean slate — drop previous demo data
    try:
        await gc.execute_write("MATCH (n) DETACH DELETE n", {})
        info("Database", "koan_demo5 (cleaned)")
    except Exception:
        info("Database", "koan_demo5 (fresh)")

    settings = ProviderSettings()
    llm = LLMClient(settings)

    registry = ToolRegistry()
    registry.register(lookup_account)
    registry.register(check_balance)
    registry.register(apply_credit)
    registry.register(escalate_ticket)

    info("Tools registered", [t.name for t in registry.definitions()])

    # ─────────────────────────────────────────────────────────────────
    # Step 1: Cold start — no policies, all tools available
    # ─────────────────────────────────────────────────────────────────
    step(1, "Cold Start — No Policies")
    info("Policies in graph", 0)
    info("Tools visible to agent", "ALL (no gating)")

    await show_augmented_prompt(gc, "billing-agent")

    run1 = await run_interaction(
        llm, gc, registry,
        query="Look up account acct-1001 and tell me their plan and balance.",
        entity_ids=["acct-1001"],
        run_number=1,
    )
    await record_feedback(gc, run_id=run1, feedback_type=FeedbackType.THUMBS_UP, score=1.0)
    success("Feedback: thumbs up")
    done()

    await show_graph_state(gc, "After Run 1")

    # ─────────────────────────────────────────────────────────────────
    # Step 2: Accumulation — repeat similar queries to build patterns
    # ─────────────────────────────────────────────────────────────────
    step(2, "Accumulation — Building Patterns (Runs 2-4)")

    queries = [
        ("What is the balance on acct-1002?", ["acct-1002"]),
        ("Look up acct-1003 and check their status.", ["acct-1003"]),
        ("Check the balance for acct-1001.", ["acct-1001"]),
    ]

    run_ids: list[str] = [run1]
    for i, (q, eids) in enumerate(queries, start=2):
        print(f"\n{YELLOW}│  --- Run {i} ---{RESET}")
        rid = await run_interaction(llm, gc, registry, q, eids, run_number=i)
        run_ids.append(rid)
        # Positive feedback on all
        await record_feedback(gc, run_id=rid, feedback_type=FeedbackType.THUMBS_UP, score=1.0)
        success("Feedback: thumbs up")

    done()
    await show_graph_state(gc, "After Run 4 (patterns should emerge)")

    # ─────────────────────────────────────────────────────────────────
    # Step 3: Verify policy emergence
    # ─────────────────────────────────────────────────────────────────
    step(3, "Policy Emergence — Check What Evolved")

    pm = PolicyManager(gc)
    policies = await pm.list_policies()

    if policies:
        success(f"{len(policies)} policy(ies) emerged automatically!")
        for p in policies:
            info("Policy", f"{p.get('name')} — {p.get('description', '')[:80]}")
    else:
        info("Policies", "None yet (need more runs with the same pattern)")
        # Force a discovery check to show what's happening
        patterns = await pm.discover_patterns(min_frequency=2, min_success_rate=0.5)
        if patterns:
            info("Discovered patterns", len(patterns))
            for pat in patterns:
                info("  Pattern", f"{pat.decision_type} → {pat.tool_sequence} (freq={pat.frequency})")
        else:
            info("Discovered patterns", "0 — agent may have used different tool sequences")

    done()

    # ─────────────────────────────────────────────────────────────────
    # Step 4: Tool gating — does the agent see restricted tools?
    # ─────────────────────────────────────────────────────────────────
    step(4, "Tool Gating — Policy-Restricted Tools")

    # Show what tools would be gated
    from koan.graph.context import GraphContextProvider as CtxProvider  # noqa: PLC0415
    calibrator = ConfidenceCalibrator(gc)
    provider = CtxProvider(gc, calibrator=calibrator)
    ctx = await provider.gather(agent_id="billing-agent", decision_type="billing", limit=5)

    if ctx.active_policies:
        allowed = PolicyGatedTools.extract_tools_from_policies(ctx.active_policies)
        if allowed:
            gated = PolicyGatedTools(registry, allowed_tools=allowed)
            visible = [d.name for d in gated.definitions()]
            all_tools = [d.name for d in registry.definitions()]
            blocked = set(all_tools) - set(visible)
            info("All tools", all_tools)
            info("Visible (gated)", visible)
            if blocked:
                warn(f"Blocked by policy: {sorted(blocked)}")
            success("Tool gating is ACTIVE")
        else:
            info("Policies exist but don't constrain specific tools", "all pass through")
    else:
        info("No policies yet", "all tools pass through (permissive cold start)")

    await show_augmented_prompt(gc, "billing-agent")

    # Run once more with gating active
    print(f"\n{YELLOW}│  --- Run 5 (with policy context) ---{RESET}")
    run5 = await run_interaction(
        llm, gc, registry,
        query="Apply a $10 credit to acct-1001.",
        entity_ids=["acct-1001"],
        run_number=5,
    )
    run_ids.append(run5)
    await record_feedback(gc, run_id=run5, feedback_type=FeedbackType.THUMBS_UP, score=1.0)
    success("Feedback: thumbs up")
    done()

    await show_graph_state(gc, "After Run 5")

    # ─────────────────────────────────────────────────────────────────
    # Step 5: Negative feedback + retirement
    # ─────────────────────────────────────────────────────────────────
    step(5, "Negative Feedback & Policy Retirement")

    info("Scenario", "Simulating bad outcomes to test retirement")

    # Run queries that will generate negative feedback
    bad_queries = [
        ("Escalate acct-1003 because they are unhappy.", ["acct-1003"]),
        ("Escalate acct-1002 — billing dispute.", ["acct-1002"]),
        ("Escalate acct-1001 — cannot resolve.", ["acct-1001"]),
    ]

    bad_run_ids: list[str] = []
    for i, (q, eids) in enumerate(bad_queries, start=6):
        print(f"\n{YELLOW}│  --- Run {i} (bad outcome) ---{RESET}")
        rid = await run_interaction(llm, gc, registry, q, eids, run_number=i)
        run_ids.append(rid)
        bad_run_ids.append(rid)
        # Negative feedback
        await record_feedback(gc, run_id=rid, feedback_type=FeedbackType.THUMBS_DOWN, score=-1.0)
        warn("Feedback: thumbs DOWN")

    # Mark those runs' outcomes as "failure" — simulates an evaluation step
    # (e.g. a human reviewer or automated evaluator marking the outcome as bad)
    info("Evaluation", "Marking bad runs as failures (simulating evaluator)")
    for rid in bad_run_ids:
        await gc.execute_write(
            """
            MATCH (r:Run {run_id: $run_id})-[:PRODUCED]->(o:Outcome)
            SET o.status = 'failure'
            """,
            {"run_id": rid},
        )
    warn(f"Flipped {len(bad_run_ids)} outcome(s) to 'failure'")

    # Now trigger retirement check explicitly
    retired = await pm.retire_underperforming(max_failure_rate=0.5, min_governed_decisions=2)
    if retired:
        warn(f"Retired {len(retired)} underperforming policy(ies): {retired}")
    else:
        info("Retirement", "No policies met retirement criteria (need ≥2 governed decisions with >50% failure)")

    done()
    await show_graph_state(gc, "After Negative Feedback Runs")

    # ─────────────────────────────────────────────────────────────────
    # Step 6: Recovery — tools ungated after policy retirement
    # ─────────────────────────────────────────────────────────────────
    step(6, "Recovery — Post-Retirement Run")

    # Check if tools are ungated now
    ctx2 = await provider.gather(agent_id="billing-agent", decision_type="billing", limit=5)
    if ctx2.active_policies:
        allowed2 = PolicyGatedTools.extract_tools_from_policies(ctx2.active_policies)
        if allowed2:
            info("Still gated to", allowed2)
        else:
            success("No tool constraints — all tools available again!")
    else:
        success("No policies remain — all tools available again!")

    print(f"\n{YELLOW}│  --- Run 9 (post-retirement) ---{RESET}")
    run9 = await run_interaction(
        llm, gc, registry,
        query="Apply a $10 credit to acct-1001.",
        entity_ids=["acct-1001"],
        run_number=9,
    )
    run_ids.append(run9)
    await record_feedback(gc, run_id=run9, feedback_type=FeedbackType.THUMBS_UP, score=1.0)
    success("Feedback: thumbs up")

    done()
    await show_graph_state(gc, "After Recovery Run")

    # ─────────────────────────────────────────────────────────────────
    # Step 7: Final summary
    # ─────────────────────────────────────────────────────────────────
    step(7, "Final Summary")

    total_runs = len(run_ids)
    final_policies = await pm.list_policies()
    patterns = await pm.discover_patterns(min_frequency=2, min_success_rate=0.5)

    info("Total runs", total_runs)
    info("Active policies", len(final_policies))
    info("Discovered patterns", len(patterns))

    for p in final_policies:
        info("  Policy", f"{p.get('name')} (source={p.get('source')})")

    for pat in patterns:
        info("  Pattern", f"{pat.decision_type} → {pat.tool_sequence} "
             f"(freq={pat.frequency}, success={pat.success_rate:.0%})")

    # Show the full graph context one last time
    await show_augmented_prompt(gc, "billing-agent")

    print()
    success("Demo complete!")
    print()
    print(f"{DIM}  Next steps:{RESET}")
    print(f"{DIM}  • Explore the graph visually: http://localhost:3000{RESET}")
    print(f"{DIM}    (FalkorDB Browser — connect to koan_demo5){RESET}")
    print(f"{DIM}  • Try the interactive web demo: uv run python demos/demo2_graph_chat.py{RESET}")
    print(f"{DIM}  • View the full graph:  MATCH (n) RETURN n{RESET}")
    print()

    done()
    await gc.close()


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