"""Demo 8 — Investment Advisor with Self-Organizing Graph.

A graph-augmented investment advisor that helps with portfolio lookup,
market data, trade execution, and risk assessment.  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/demo8_finance.py
  Open http://localhost:8008
"""

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 — portfolios, market data, positions
# ---------------------------------------------------------------------------

PORTFOLIOS: dict[str, dict[str, Any]] = {
    "port-7001": {
        "owner": "Alice Chen",
        "type": "growth",
        "risk_tolerance": "moderate",
        "total_value": 245_800.00,
        "cash_available": 18_500.00,
        "ytd_return": 12.3,
        "positions": [
            {"symbol": "AAPL", "shares": 150, "avg_cost": 145.20, "current_price": 198.50},
            {"symbol": "MSFT", "shares": 80, "avg_cost": 310.00, "current_price": 425.80},
            {"symbol": "AMZN", "shares": 40, "avg_cost": 128.50, "current_price": 185.20},
            {"symbol": "VTI", "shares": 200, "avg_cost": 205.00, "current_price": 248.60},
        ],
    },
    "port-7002": {
        "owner": "Robert Park",
        "type": "income",
        "risk_tolerance": "conservative",
        "total_value": 520_400.00,
        "cash_available": 42_000.00,
        "ytd_return": 6.8,
        "positions": [
            {"symbol": "JNJ", "shares": 300, "avg_cost": 155.40, "current_price": 162.30},
            {"symbol": "PG", "shares": 250, "avg_cost": 140.80, "current_price": 168.50},
            {"symbol": "BND", "shares": 500, "avg_cost": 72.50, "current_price": 73.20},
            {"symbol": "SCHD", "shares": 400, "avg_cost": 74.00, "current_price": 82.40},
        ],
    },
    "port-7003": {
        "owner": "Jessica Torres",
        "type": "aggressive",
        "risk_tolerance": "high",
        "total_value": 89_200.00,
        "cash_available": 5_300.00,
        "ytd_return": 28.5,
        "positions": [
            {"symbol": "NVDA", "shares": 50, "avg_cost": 480.00, "current_price": 890.50},
            {"symbol": "TSLA", "shares": 60, "avg_cost": 195.00, "current_price": 245.80},
            {"symbol": "COIN", "shares": 100, "avg_cost": 85.00, "current_price": 210.40},
        ],
    },
}

MARKET_DATA: dict[str, dict[str, Any]] = {
    "AAPL": {
        "price": 198.50, "change": 2.30, "change_pct": 1.17,
        "volume": "52.3M", "pe_ratio": 31.2, "dividend_yield": 0.53,
    },
    "MSFT": {
        "price": 425.80, "change": -1.40, "change_pct": -0.33,
        "volume": "22.1M", "pe_ratio": 36.8, "dividend_yield": 0.72,
    },
    "AMZN": {
        "price": 185.20, "change": 3.80, "change_pct": 2.10,
        "volume": "48.7M", "pe_ratio": 62.5, "dividend_yield": 0.0,
    },
    "NVDA": {
        "price": 890.50, "change": 15.20, "change_pct": 1.74,
        "volume": "35.6M", "pe_ratio": 72.1, "dividend_yield": 0.03,
    },
    "TSLA": {
        "price": 245.80, "change": -8.50, "change_pct": -3.34,
        "volume": "98.2M", "pe_ratio": 68.3, "dividend_yield": 0.0,
    },
    "VTI": {
        "price": 248.60, "change": 1.10, "change_pct": 0.44,
        "volume": "3.8M", "pe_ratio": 22.4, "dividend_yield": 1.35,
    },
    "JNJ": {
        "price": 162.30, "change": 0.80, "change_pct": 0.50,
        "volume": "6.2M", "pe_ratio": 15.8, "dividend_yield": 2.95,
    },
    "BND": {
        "price": 73.20, "change": 0.10, "change_pct": 0.14,
        "volume": "7.5M", "pe_ratio": 0, "dividend_yield": 4.52,
    },
}


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


@tool
def lookup_portfolio(portfolio_id: str) -> str:
    """Look up a client's investment portfolio including all positions.

    Args:
        portfolio_id: The portfolio identifier, e.g. 'port-7001'.
    """
    port = PORTFOLIOS.get(portfolio_id)
    if port is None:
        return f"Portfolio {portfolio_id} not found."
    return json.dumps({"portfolio_id": portfolio_id, **port})


@tool
def get_market_data(symbol: str) -> str:
    """Get current market data for a stock or ETF symbol.

    Args:
        symbol: The ticker symbol, e.g. 'AAPL', 'MSFT'.
    """
    data = MARKET_DATA.get(symbol.upper())
    if data is None:
        return f"No market data available for {symbol}."
    return json.dumps({"symbol": symbol.upper(), **data})


@tool
def execute_trade(portfolio_id: str, action: str, symbol: str, shares: int) -> str:
    """Execute a buy or sell trade for a portfolio.

    Args:
        portfolio_id: The portfolio identifier.
        action: 'buy' or 'sell'.
        symbol: The ticker symbol.
        shares: Number of shares to trade.
    """
    port = PORTFOLIOS.get(portfolio_id)
    if port is None:
        return f"Portfolio {portfolio_id} not found."

    market = MARKET_DATA.get(symbol.upper())
    if market is None:
        return f"Cannot trade {symbol} — no market data."

    price = market["price"]
    total = price * shares

    if action.lower() == "buy":
        if total > port["cash_available"]:
            return (
                f"TRADE REJECTED — Insufficient cash. "
                f"Need ${total:,.2f} but only ${port['cash_available']:,.2f} available."
            )
        return (
            f"TRADE EXECUTED — BUY\n"
            f"Portfolio: {port['owner']} ({portfolio_id})\n"
            f"{shares} shares of {symbol.upper()} @ ${price:,.2f} = ${total:,.2f}\n"
            f"Remaining cash: ${port['cash_available'] - total:,.2f}"
        )

    # Sell
    position = next((p for p in port["positions"] if p["symbol"] == symbol.upper()), None)
    if position is None:
        return f"TRADE REJECTED — {portfolio_id} has no position in {symbol.upper()}."
    if shares > position["shares"]:
        return f"TRADE REJECTED — Only {position['shares']} shares available, cannot sell {shares}."

    gain = (price - position["avg_cost"]) * shares
    return (
        f"TRADE EXECUTED — SELL\n"
        f"Portfolio: {port['owner']} ({portfolio_id})\n"
        f"{shares} shares of {symbol.upper()} @ ${price:,.2f} = ${total:,.2f}\n"
        f"Realized gain/loss: ${gain:,.2f}"
    )


@tool
def assess_risk(portfolio_id: str) -> str:
    """Assess the risk profile and concentration of a portfolio.

    Args:
        portfolio_id: The portfolio identifier.
    """
    port = PORTFOLIOS.get(portfolio_id)
    if port is None:
        return f"Portfolio {portfolio_id} not found."

    positions = port["positions"]
    total_invested = sum(p["current_price"] * p["shares"] for p in positions)

    concentrations = []
    for p in positions:
        value = p["current_price"] * p["shares"]
        pct = (value / total_invested) * 100 if total_invested > 0 else 0
        concentrations.append({
            "symbol": p["symbol"],
            "value": round(value, 2),
            "pct_of_portfolio": round(pct, 1),
            "gain_loss_pct": round(((p["current_price"] - p["avg_cost"]) / p["avg_cost"]) * 100, 1),
        })

    # Flag risks
    flags = []
    for c in concentrations:
        if c["pct_of_portfolio"] > 30:
            flags.append(f"{c['symbol']} is {c['pct_of_portfolio']}% of portfolio — over-concentrated")
        if c["gain_loss_pct"] > 50:
            flags.append(f"{c['symbol']} up {c['gain_loss_pct']}% — consider taking profits")
        if c["gain_loss_pct"] < -20:
            flags.append(f"{c['symbol']} down {c['gain_loss_pct']}% — review thesis")

    return json.dumps({
        "portfolio_id": portfolio_id,
        "owner": port["owner"],
        "type": port["type"],
        "risk_tolerance": port["risk_tolerance"],
        "total_invested": round(total_invested, 2),
        "cash_pct": round((port["cash_available"] / port["total_value"]) * 100, 1),
        "concentrations": concentrations,
        "risk_flags": flags,
    })


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

app = FastAPI(title="Demo 8: Investment Advisor")
app.mount("/static", StaticFiles(directory="demos/static"), name="static")

settings = ProviderSettings()
llm_client = LLMClient(settings)

registry = ToolRegistry()
registry.register(lookup_portfolio)
registry.register(get_market_data)
registry.register(execute_trade)
registry.register(assess_risk)

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

graph_client: GraphClient | None = None

agent_config = AgentConfig(
    provider="openai",
    model="gpt-4o-mini",
    temperature=0.3,
    system_prompt=(
        "You are an investment advisor assistant. Help clients with portfolio reviews, "
        "market data lookups, trade execution, and risk assessment. Always consider the "
        "client's risk tolerance before recommending trades. Flag any concentration risks. "
        "Use the provided portfolio 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 8: Investment Advisor",
        badge_text="Finance",
        badge_color="orange",
        port=8008,
        info_html=(
            info_row("Model", "gpt-4o-mini")
            + info_row("Tools", "lookup_portfolio, get_market_data, execute_trade, assess_risk")
            + info_row("Domain", "Finance / Investments")
            + info_row("Graph", "FalkorDB (koan_demo8)")
        ),
    ))


@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"port-\d{4}", user_message)

    assert graph_client is not None  # noqa: S101

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

    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=8008)
