"""Example 18 — Human Feedback and Confidence Calibration.

Demonstrates how user feedback flows into the decision graph to
calibrate future agent confidence:

1. Agent runs and produces a decision
2. User provides satisfaction feedback ("were you satisfied?")
3. Fork/rewind events are tracked as implicit dissatisfaction
4. ConfidenceCalibrator reads feedback history to adjust confidence
5. Next agent run sees the feedback summary in its context

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 (
    ConfidenceCalibrator,
    DecisionTracer,
    FeedbackType,
    GraphAugmentedAgent,
    GraphConfig,
    create_graph_client,
    record_feedback,
)
from koan.llm.client import LLMClient
from koan.llm.config import ProviderSettings
from koan.memory.fork import ForkableThread
from koan.tools import ToolRegistry, tool

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

ACCOUNTS = {
    "acct-001": {"name": "Alice Corp", "balance": 15_000.00, "status": "active"},
    "acct-002": {"name": "Bob LLC", "balance": 3_200.50, "status": "active"},
}


@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} | Status: {acct['status']}"


# ── Seed some runs + feedback ────────────────────────────────────────

async def seed_data(graph_client):
    """Seed the graph with past runs that have feedback attached."""
    # Run 1: Successful, user satisfied
    async with DecisionTracer(
        graph_client, run_id="seed-fb-1", agent_id="support-agent", thread_id="t-fb-1",
    ) as tracer:
        d1 = await tracer.record_decision(
            description="Handled billing inquiry for Alice Corp",
            decision_type="classification",
            confidence=0.9,
            reasoning="Keyword match: billing",
        )
        await tracer.link_entity(decision_id=d1, entity_id="acct-001", entity_type="account")
        await tracer.record_outcome(status="success", result="Resolved billing question")

    await record_feedback(
        graph_client,
        run_id="seed-fb-1",
        feedback_type=FeedbackType.THUMBS_UP,
        score=1.0,
        comment="Very helpful, resolved quickly!",
    )

    # Run 2: User was not satisfied
    async with DecisionTracer(
        graph_client, run_id="seed-fb-2", agent_id="support-agent", thread_id="t-fb-2",
    ) as tracer:
        d2 = await tracer.record_decision(
            description="Handled refund request for Bob LLC",
            decision_type="classification",
            confidence=0.85,
            reasoning="Keyword match: refund",
        )
        await tracer.link_entity(decision_id=d2, entity_id="acct-002", entity_type="account")
        await tracer.record_outcome(status="success", result="Processed refund")

    await record_feedback(
        graph_client,
        run_id="seed-fb-2",
        feedback_type=FeedbackType.THUMBS_DOWN,
        score=-1.0,
        comment="Refund amount was wrong",
    )

    # Run 3: User forked the conversation (implicit dissatisfaction)
    async with DecisionTracer(
        graph_client, run_id="seed-fb-3", agent_id="support-agent", thread_id="t-fb-3",
    ) as tracer:
        d3 = await tracer.record_decision(
            description="Handled account status check",
            decision_type="classification",
            confidence=0.95,
            reasoning="Direct question about status",
        )
        await tracer.record_outcome(status="success", result="Returned account status")

    # Simulate fork event as implicit negative signal
    await record_feedback(
        graph_client,
        run_id="seed-fb-3",
        feedback_type=FeedbackType.FORK,
        score=-0.5,
        comment="",
    )

    print("[seed] Graph seeded with 3 runs + feedback (1 positive, 1 negative, 1 fork)")


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

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

    registry = ToolRegistry()
    registry.add(lookup_account)

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

        # Step 2: Check calibration
        print("\n" + "=" * 60)
        print("STEP 2: Confidence Calibration")
        print("=" * 60)

        calibrator = ConfidenceCalibrator(graph)
        cal = await calibrator.get_calibration(agent_id="support-agent")
        print(f"  Calibration factor: {cal.factor}")
        print(f"  Sample count: {cal.sample_count}")
        print(f"  Positive rate: {cal.positive_rate:.0%}")
        print(f"  Negative rate: {cal.negative_rate:.0%}")
        print(f"  Fork rate: {cal.fork_rate:.0%}")

        summary = await calibrator.get_feedback_summary(agent_id="support-agent")
        if summary:
            print(f"\n  Feedback summary:")
            print(f"    Total: {summary['total_feedback']}")
            print(f"    Satisfaction rate: {summary['satisfaction_rate']:.0%}")
            print(f"    Recent comments: {summary.get('recent_comments', [])}")

        # Step 3: Run graph-augmented agent — it now sees feedback in context
        print("\n" + "=" * 60)
        print("STEP 3: Graph-augmented agent with feedback context")
        print("=" * 60)

        agent = GraphAugmentedAgent(
            "support-agent",
            client=llm,
            config=AgentConfig(
                provider="openai",
                model="gpt-4o",
                system_prompt=(
                    "You are a support agent. Use available tools. "
                    "Pay attention to past feedback in your context — "
                    "if users have been dissatisfied, be extra careful."
                ),
            ),
            tools=registry,
            graph_client=graph,
            entity_ids=["acct-001"],
            decision_type="classification",
        )

        result = await agent.run("What is the balance on account acct-001?")
        print(f"\nAgent output: {result.output}")
        print(f"Graph context: {result.metadata.get('graph_context', {})}")

        # Step 4: Submit feedback on this new run
        print("\n" + "=" * 60)
        print("STEP 4: Submit user feedback")
        print("=" * 60)

        fid = await agent.submit_feedback(satisfied=True, comment="Accurate answer!")
        print(f"  Feedback recorded: {fid}")

        # Step 5: Demonstrate fork tracking
        print("\n" + "=" * 60)
        print("STEP 5: Fork tracking as implicit dissatisfaction")
        print("=" * 60)

        fork_events: list[tuple[str, str, str]] = []

        def on_fork(thread_id: str, event_type: str, checkpoint: str) -> None:
            fork_events.append((thread_id, event_type, checkpoint))
            print(f"  [event] {event_type} on thread {thread_id} at '{checkpoint}'")

        thread = ForkableThread("demo-thread", on_fork=on_fork)
        thread.append(role="user", content="What's my balance?")
        thread.checkpoint("after_question")
        thread.append(role="assistant", content="Your balance is $15,000.")
        thread.checkpoint("after_answer")

        # User is unhappy, forks to try again
        print("  User forks the conversation (implicit dissatisfaction)...")
        branch = thread.fork_at("after_question")
        branch.append(role="assistant", content="Let me look that up more carefully...")

        # User rewinds to try again
        print("  User rewinds (implicit dissatisfaction)...")
        thread.rewind_to("after_question")

        print(f"\n  Total fork/rewind events captured: {len(fork_events)}")
        print("  These events can be recorded as graph feedback using record_feedback()")

        # Record the fork events as graph feedback
        run_id = result.metadata.get("run_id", "")
        if run_id:
            for _, event_type, _ in fork_events:
                ftype = FeedbackType.FORK if event_type == "fork" else FeedbackType.REWIND
                await record_feedback(
                    graph,
                    run_id=run_id,
                    feedback_type=ftype,
                    score=-0.5,
                )
            print(f"  Recorded {len(fork_events)} fork/rewind signals to graph")

        # Step 6: Re-check calibration after new feedback
        print("\n" + "=" * 60)
        print("STEP 6: Updated calibration after feedback")
        print("=" * 60)

        cal2 = await calibrator.get_calibration(agent_id="support-agent")
        print(f"  Calibration factor: {cal2.factor} (was {cal.factor})")
        print(f"  Sample count: {cal2.sample_count} (was {cal.sample_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)
