"""Demo 6 — Legal Assistant with Self-Organizing Graph.

A graph-augmented legal assistant that helps with case lookup, document
search, motion drafting, and hearing scheduling.  Policies emerge from
repeated usage patterns — no hardcoded rules.

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/demo6_legal.py
  Open http://localhost:8006
"""

from __future__ import annotations

import json
import re
import sys
import uuid
from pathlib import Path
from typing import TYPE_CHECKING, Any

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

import uvicorn
from dotenv import load_dotenv
from fastapi import FastAPI
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from starlette.responses import StreamingResponse
from ui import chat_page, info_row

from koan.agents.events import AgentComplete, AgentError, TextDelta, ToolCallEnd, ToolCallStart
from koan.agents.types import AgentConfig
from koan.graph import (
    FeedbackType,
    GraphAugmentedAgent,
    GraphConfig,
    create_graph_client,
    record_feedback,
)
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()

# ---------------------------------------------------------------------------
# Stub data — cases, documents, hearings
# ---------------------------------------------------------------------------

CASES: dict[str, dict[str, Any]] = {
    "case-2024-001": {
        "title": "Smith v. Metro Insurance Co.",
        "type": "civil",
        "status": "active",
        "client": "John Smith",
        "opposing": "Metro Insurance Co.",
        "judge": "Hon. Patricia Clark",
        "filed": "2024-06-15",
        "summary": "Personal injury claim arising from auto accident. Defendant disputes liability.",
    },
    "case-2024-002": {
        "title": "In re: Estate of Margaret Williams",
        "type": "probate",
        "status": "active",
        "client": "Williams Family Trust",
        "opposing": "N/A",
        "judge": "Hon. Robert Chen",
        "filed": "2024-08-20",
        "summary": "Probate of estate valued at $2.3M. Contested will by two beneficiaries.",
    },
    "case-2024-003": {
        "title": "TechStart Inc. v. DataFlow Corp.",
        "type": "ip",
        "status": "discovery",
        "client": "TechStart Inc.",
        "opposing": "DataFlow Corp.",
        "judge": "Hon. Sarah Kim",
        "filed": "2024-03-10",
        "summary": "Patent infringement claim re: US Patent 11,234,567 (data pipeline architecture).",
    },
}

DOCUMENTS: dict[str, list[dict[str, str]]] = {
    "case-2024-001": [
        {"doc_id": "doc-101", "title": "Complaint", "type": "pleading", "date": "2024-06-15"},
        {"doc_id": "doc-102", "title": "Answer and Counterclaim", "type": "pleading", "date": "2024-07-20"},
        {"doc_id": "doc-103", "title": "Police Report #PR-44821", "type": "evidence", "date": "2024-06-01"},
        {"doc_id": "doc-104", "title": "Medical Records — Dr. Patel", "type": "evidence", "date": "2024-06-20"},
    ],
    "case-2024-002": [
        {"doc_id": "doc-201", "title": "Last Will and Testament", "type": "estate", "date": "2019-03-12"},
        {"doc_id": "doc-202", "title": "Petition for Probate", "type": "pleading", "date": "2024-08-20"},
        {"doc_id": "doc-203", "title": "Asset Inventory", "type": "financial", "date": "2024-09-01"},
    ],
    "case-2024-003": [
        {"doc_id": "doc-301", "title": "Patent US 11,234,567", "type": "ip", "date": "2022-01-15"},
        {"doc_id": "doc-302", "title": "Infringement Analysis Report", "type": "expert", "date": "2024-04-05"},
        {"doc_id": "doc-303", "title": "Claim Construction Brief", "type": "pleading", "date": "2024-05-10"},
    ],
}

HEARINGS: dict[str, list[dict[str, str]]] = {
    "case-2024-001": [
        {"date": "2025-04-10", "type": "Motion Hearing", "location": "Courtroom 4B"},
    ],
    "case-2024-003": [
        {"date": "2025-03-28", "type": "Markman Hearing", "location": "Courtroom 7A"},
        {"date": "2025-06-15", "type": "Trial", "location": "Courtroom 7A"},
    ],
}


# ---------------------------------------------------------------------------
# Tools
# ---------------------------------------------------------------------------


@tool
def lookup_case(case_id: str) -> str:
    """Look up a legal case by case ID.

    Args:
        case_id: The case identifier, e.g. 'case-2024-001'.
    """
    case = CASES.get(case_id)
    if case is None:
        return f"Case {case_id} not found."
    return json.dumps({"case_id": case_id, **case})


@tool
def search_documents(case_id: str) -> str:
    """Search for documents filed in a case.

    Args:
        case_id: The case identifier.
    """
    docs = DOCUMENTS.get(case_id)
    if docs is None:
        return f"No documents found for {case_id}."
    return json.dumps({"case_id": case_id, "documents": docs, "total": len(docs)})


@tool
def draft_motion(case_id: str, motion_type: str, grounds: str) -> str:
    """Draft a motion for a case.

    Args:
        case_id: The case identifier.
        motion_type: Type of motion (e.g. 'summary_judgment', 'dismiss', 'compel').
        grounds: Legal grounds for the motion.
    """
    case = CASES.get(case_id)
    if case is None:
        return f"Case {case_id} not found."
    return (
        f"DRAFT MOTION — {motion_type.upper()}\n"
        f"Case: {case['title']}\n"
        f"Before: {case['judge']}\n"
        f"Grounds: {grounds}\n"
        f"Status: Draft created. Review and file when ready."
    )


@tool
def check_hearing_schedule(case_id: str) -> str:
    """Check upcoming hearings for a case.

    Args:
        case_id: The case identifier.
    """
    hearings = HEARINGS.get(case_id)
    if not hearings:
        return f"No upcoming hearings scheduled for {case_id}."
    return json.dumps({"case_id": case_id, "hearings": hearings})


# ---------------------------------------------------------------------------
# App setup
# ---------------------------------------------------------------------------

app = FastAPI(title="Demo 6: Legal Assistant")
app.mount("/static", StaticFiles(directory="demos/static"), name="static")

settings = ProviderSettings()
llm_client = LLMClient(settings)

registry = ToolRegistry()
registry.register(lookup_case)
registry.register(search_documents)
registry.register(draft_motion)
registry.register(check_hearing_schedule)

graph_config = GraphConfig(
    provider="falkordb",
    uri="redis://localhost:6379",
    database="koan_demo6",
)

graph_client: GraphClient | None = None

agent_config = AgentConfig(
    provider="openai",
    model="gpt-4o-mini",
    temperature=0.3,
    system_prompt=(
        "You are a legal assistant at a law firm. Help attorneys with case lookups, "
        "document searches, motion drafting, and hearing schedules. Be precise and "
        "cite case details when relevant. Use the provided case ID directly."
    ),
)

threads: dict[str, list[dict[str, str]]] = {}


@app.on_event("startup")
async def startup() -> None:
    global graph_client  # noqa: PLW0603
    graph_client = create_graph_client(graph_config)
    await graph_client.connect()


@app.on_event("shutdown")
async def shutdown() -> None:
    if graph_client is not None:
        await graph_client.close()


# ---------------------------------------------------------------------------
# Request / response models
# ---------------------------------------------------------------------------


class ChatRequest(BaseModel):
    thread_id: str
    message: str


class FeedbackRequest(BaseModel):
    run_id: str
    satisfied: bool


class CreateThreadResponse(BaseModel):
    thread_id: str


# ---------------------------------------------------------------------------
# Routes
# ---------------------------------------------------------------------------


@app.get("/", response_class=HTMLResponse)
async def index() -> HTMLResponse:
    return HTMLResponse(content=chat_page(
        title="Demo 6: Legal Assistant",
        badge_text="Law",
        badge_color="purple",
        port=8006,
        info_html=(
            info_row("Model", "gpt-4o-mini")
            + info_row("Tools", "lookup_case, search_documents, draft_motion, check_hearing_schedule")
            + info_row("Domain", "Legal / Law")
            + info_row("Graph", "FalkorDB (koan_demo6)")
        ),
    ))


@app.get("/api/threads")
async def list_threads() -> list[dict[str, str]]:
    return [{"thread_id": tid} for tid in threads]


@app.post("/api/threads")
async def create_thread() -> CreateThreadResponse:
    thread_id = str(uuid.uuid4())
    threads[thread_id] = []
    return CreateThreadResponse(thread_id=thread_id)


@app.get("/api/threads/{thread_id}/messages")
async def get_messages(thread_id: str) -> list[dict]:
    return [
        {"role": m["role"], "content": m["content"], "metadata": m.get("metadata", {})}
        for m in threads.get(thread_id, [])
        if m.get("role") in ("user", "assistant")
    ]


@app.post("/api/chat")
async def chat(req: ChatRequest) -> StreamingResponse:
    thread_id = req.thread_id
    user_message = req.message

    if thread_id not in threads:
        threads[thread_id] = []

    threads[thread_id].append({"role": "user", "content": user_message})

    conversation_parts: list[str] = []
    for m in threads[thread_id]:
        if m["role"] == "user":
            conversation_parts.append(f"User: {m['content']}")
        elif m["role"] == "assistant":
            conversation_parts.append(f"Assistant: {m['content']}")

    input_text = "\n".join(conversation_parts) if len(conversation_parts) > 1 else user_message

    entity_ids = re.findall(r"case-\d{4}-\d{3}", user_message)

    assert graph_client is not None  # noqa: S101

    agent = GraphAugmentedAgent(
        "legal-agent",
        client=llm_client,
        config=agent_config,
        tools=registry,
        graph_client=graph_client,
        entity_ids=entity_ids,
        decision_type="legal",
    )

    async def event_stream():  # noqa: ANN202
        full_response = ""
        run_id = ""
        graph_context_info: dict[str, Any] = {}
        tool_calls: list[dict[str, str]] = []
        try:
            async for event in agent.run_stream(input_text):
                if isinstance(event, TextDelta):
                    full_response += event.text
                    yield f"data: {json.dumps({'type': 'token', 'token': event.text})}\n\n"
                elif isinstance(event, ToolCallStart):
                    data = json.dumps({"type": "tool_start", "tool": event.tool_name, "arguments": event.arguments})
                    yield f"data: {data}\n\n"
                elif isinstance(event, ToolCallEnd):
                    tool_calls.append({"tool": event.tool_name, "result": event.result})
                    data = json.dumps({"type": "tool_end", "tool": event.tool_name, "result": event.result})
                    yield f"data: {data}\n\n"
                elif isinstance(event, AgentComplete):
                    run_id = event.metadata.get("run_id", "")
                    graph_context_info = event.metadata.get("graph_context", {})
                    meta: dict[str, Any] = {
                        "iterations": event.iterations,
                        "tool_calls_made": event.tool_calls_made,
                        "graph_context": graph_context_info,
                        "tool_calls": tool_calls,
                        "run_id": run_id,
                    }
                    threads[thread_id].append({"role": "assistant", "content": event.output, "metadata": meta})
                    yield f"data: {json.dumps({'type': 'complete', 'run_id': run_id, 'metadata': meta})}\n\n"
                elif isinstance(event, AgentError):
                    threads[thread_id].append({"role": "assistant", "content": event.output})
                    yield f"data: {json.dumps({'type': 'error', 'error': event.error})}\n\n"
        except Exception as exc:
            yield f"data: {json.dumps({'type': 'error', 'error': str(exc)})}\n\n"

    return StreamingResponse(event_stream(), media_type="text/event-stream")


@app.post("/api/feedback")
async def submit_feedback(req: FeedbackRequest) -> dict[str, str]:
    assert graph_client is not None  # noqa: S101
    score = 1.0 if req.satisfied else -1.0
    ftype = FeedbackType.THUMBS_UP if req.satisfied else FeedbackType.THUMBS_DOWN
    feedback_id = await record_feedback(graph_client, run_id=req.run_id, feedback_type=ftype, score=score)
    return {"feedback_id": feedback_id, "status": "recorded"}


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8006)
