#!/usr/bin/env python3
"""
Multi-Turn CLI Chatbot with Session Tracking

A fully interactive personal finance advisor that you can chat with in the
terminal. Each conversation is a separate session — every user message
creates its own trace, and all traces are grouped under one session_id.

This demonstrates:
  - init_tracing / trace_run / trace_agent / set_session_id
  - wrap_llm_client for automatic LLM tracing
  - @tool decorated functions for automatic tool tracing
  - @function decorated helpers as child spans
  - Multi-turn memory carried across turns
  - Session grouping visible on the dashboard Sessions page

Usage:
    export CASCADE_API_KEY="csk_live_..."
    export CASCADE_ENDPOINT="http://localhost:8000/v1/traces"
    export ANTHROPIC_API_KEY="sk-ant-..."
    python tests/test_session_chatbot.py
"""
import os
import sys
import uuid
import time
from datetime import datetime

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

try:
    from dotenv import load_dotenv
    load_dotenv(os.path.join(os.path.dirname(__file__), "..", ".env"))
except ImportError:
    pass

from cascade import (
    init_tracing, trace_run, trace_agent, set_session_id, end_session,
    wrap_llm_client, tool, function,
)
from anthropic import Anthropic


# ---------------------------------------------------------------------------
# Helper functions (@function) — show up as child spans under tools
# ---------------------------------------------------------------------------

@function
def calculate_compound_growth(principal: float, rate: float, years: int) -> dict:
    """Project compound growth over N years."""
    projections = []
    balance = principal
    for y in range(1, years + 1):
        balance *= (1 + rate)
        projections.append({"year": y, "balance": round(balance, 2)})
    return {"final_balance": round(balance, 2), "total_return": round(balance - principal, 2), "projections": projections[-3:]}


@function
def assess_risk_score(age: int, income: float, debt: float) -> dict:
    """Compute a simple risk tolerance score 1-10."""
    age_factor = max(0, (65 - age)) / 65
    income_factor = min(income / 200_000, 1.0)
    debt_ratio = min(debt / max(income, 1), 1.0)
    score = round((age_factor * 4 + income_factor * 4 + (1 - debt_ratio) * 2), 1)
    score = max(1, min(10, score))
    label = "conservative" if score <= 3 else "moderate" if score <= 6 else "aggressive"
    return {"score": score, "label": label}


@function
def format_currency(amount: float) -> str:
    """Format a number as USD."""
    return f"${amount:,.2f}"


# ---------------------------------------------------------------------------
# Tools (@tool) — each call becomes a traced span
# ---------------------------------------------------------------------------

@tool
def get_portfolio(user_id: str) -> str:
    """Look up a user's current investment portfolio."""
    portfolios = {
        "alice": {
            "stocks": [
                {"ticker": "AAPL", "shares": 50, "avg_cost": 142.00, "current_price": 189.50},
                {"ticker": "MSFT", "shares": 30, "avg_cost": 285.00, "current_price": 415.20},
                {"ticker": "GOOGL", "shares": 20, "avg_cost": 120.00, "current_price": 175.80},
            ],
            "bonds": {"total_value": 25_000, "yield_pct": 4.2},
            "cash": 12_500,
        },
        "default": {
            "stocks": [
                {"ticker": "VOO", "shares": 100, "avg_cost": 380.00, "current_price": 502.30},
                {"ticker": "QQQ", "shares": 40, "avg_cost": 350.00, "current_price": 480.10},
            ],
            "bonds": {"total_value": 15_000, "yield_pct": 3.8},
            "cash": 8_000,
        },
    }
    p = portfolios.get(user_id.lower(), portfolios["default"])
    lines = [f"Portfolio for {user_id}:"]
    total = 0.0
    for s in p["stocks"]:
        value = s["shares"] * s["current_price"]
        gain_pct = ((s["current_price"] - s["avg_cost"]) / s["avg_cost"]) * 100
        total += value
        lines.append(f"  {s['ticker']:6s}  {s['shares']:>4d} shares  @{format_currency(s['current_price'])}  ({gain_pct:+.1f}%)")
    total += p["bonds"]["total_value"] + p["cash"]
    lines.append(f"  Bonds: {format_currency(p['bonds']['total_value'])} (yield {p['bonds']['yield_pct']}%)")
    lines.append(f"  Cash:  {format_currency(p['cash'])}")
    lines.append(f"  Total: {format_currency(total)}")
    return "\n".join(lines)


@tool
def get_stock_quote(ticker: str) -> str:
    """Get the latest quote for a stock ticker."""
    quotes = {
        "AAPL": {"price": 189.50, "change": +1.23, "pct": +0.65, "volume": "52.3M", "pe": 29.4, "high_52w": 199.62, "low_52w": 143.90},
        "MSFT": {"price": 415.20, "change": -2.10, "pct": -0.50, "volume": "21.1M", "pe": 35.1, "high_52w": 430.82, "low_52w": 309.45},
        "GOOGL": {"price": 175.80, "change": +3.45, "pct": +2.00, "volume": "28.7M", "pe": 24.8, "high_52w": 180.40, "low_52w": 120.21},
        "VOO": {"price": 502.30, "change": +1.05, "pct": +0.21, "volume": "4.8M", "pe": 22.1, "high_52w": 510.00, "low_52w": 390.00},
        "QQQ": {"price": 480.10, "change": +2.70, "pct": +0.57, "volume": "38.2M", "pe": 30.5, "high_52w": 495.00, "low_52w": 345.00},
        "TSLA": {"price": 245.60, "change": -5.40, "pct": -2.15, "volume": "95.6M", "pe": 62.3, "high_52w": 299.29, "low_52w": 138.80},
        "AMZN": {"price": 185.30, "change": +0.80, "pct": +0.43, "volume": "44.1M", "pe": 58.7, "high_52w": 191.70, "low_52w": 118.35},
    }
    q = quotes.get(ticker.upper())
    if not q:
        return f"No quote found for {ticker.upper()}. Try AAPL, MSFT, GOOGL, VOO, QQQ, TSLA, AMZN."
    sign = "+" if q["change"] >= 0 else ""
    return (
        f"{ticker.upper()}  ${q['price']:.2f}  {sign}{q['change']:.2f} ({sign}{q['pct']:.2f}%)\n"
        f"  Volume: {q['volume']}  |  P/E: {q['pe']}  |  52w: ${q['low_52w']:.2f} - ${q['high_52w']:.2f}"
    )


@tool
def get_market_news(topic: str) -> str:
    """Get latest market news for a topic (e.g., 'tech', 'bonds', 'inflation')."""
    news = {
        "tech": [
            "AI spending drives cloud revenue higher for major tech firms",
            "Semiconductor stocks rally on strong demand forecasts",
            "Apple announces record services revenue in Q4 earnings",
        ],
        "bonds": [
            "10-year Treasury yield steady at 4.25% amid Fed pause",
            "Investment-grade corporate bonds see inflows for 8th straight week",
            "Municipal bond supply tightens, pushing prices higher",
        ],
        "inflation": [
            "CPI comes in at 2.9% YoY, slightly below expectations",
            "Core PCE remains sticky at 2.8%, Fed maintains cautious stance",
            "Wage growth moderates to 3.5% annualized in latest jobs report",
        ],
        "markets": [
            "S&P 500 closes at fresh all-time high above 5,800",
            "Small-cap stocks outperform as rate cut expectations build",
            "Volatility index (VIX) drops to lowest level since January",
        ],
    }
    articles = news.get(topic.lower(), news["markets"])
    lines = [f"Latest news on '{topic}':"]
    for i, headline in enumerate(articles, 1):
        lines.append(f"  {i}. {headline}")
    return "\n".join(lines)


@tool
def project_savings(monthly_contribution: float, annual_return: float, years: int) -> str:
    """Project how savings grow with monthly contributions over time."""
    result = calculate_compound_growth(0, annual_return / 12, years * 12)
    total_contributed = monthly_contribution * years * 12
    monthly_rate = annual_return / 12
    future_value = 0.0
    for _ in range(years * 12):
        future_value = (future_value + monthly_contribution) * (1 + monthly_rate)
    return (
        f"Savings projection over {years} years:\n"
        f"  Monthly contribution: {format_currency(monthly_contribution)}\n"
        f"  Annual return: {annual_return:.1%}\n"
        f"  Total contributed: {format_currency(total_contributed)}\n"
        f"  Projected value: {format_currency(future_value)}\n"
        f"  Growth: {format_currency(future_value - total_contributed)}"
    )


@tool
def assess_risk(age: int, annual_income: float, total_debt: float) -> str:
    """Assess an investor's risk tolerance based on age, income, and debt."""
    result = assess_risk_score(age, annual_income, total_debt)
    return (
        f"Risk Assessment:\n"
        f"  Score: {result['score']}/10 ({result['label']})\n"
        f"  Based on: age={age}, income={format_currency(annual_income)}, debt={format_currency(total_debt)}\n"
        f"  Recommendation: {'Growth-oriented ETFs and individual stocks' if result['label'] == 'aggressive' else 'Balanced mix of stocks and bonds' if result['label'] == 'moderate' else 'Focus on bonds, dividend stocks, and cash reserves'}"
    )


# ---------------------------------------------------------------------------
# Tool schema + dispatch for Anthropic
# ---------------------------------------------------------------------------

TOOLS_SCHEMA = [
    {
        "name": "get_portfolio",
        "description": "Look up a user's current investment portfolio. Use 'default' if no user specified.",
        "input_schema": {"type": "object", "properties": {"user_id": {"type": "string"}}, "required": ["user_id"]},
    },
    {
        "name": "get_stock_quote",
        "description": "Get the latest quote for a stock ticker (AAPL, MSFT, GOOGL, VOO, QQQ, TSLA, AMZN).",
        "input_schema": {"type": "object", "properties": {"ticker": {"type": "string"}}, "required": ["ticker"]},
    },
    {
        "name": "get_market_news",
        "description": "Get latest market news for a topic: tech, bonds, inflation, or markets.",
        "input_schema": {"type": "object", "properties": {"topic": {"type": "string"}}, "required": ["topic"]},
    },
    {
        "name": "project_savings",
        "description": "Project how savings grow with monthly contributions over time.",
        "input_schema": {
            "type": "object",
            "properties": {
                "monthly_contribution": {"type": "number"},
                "annual_return": {"type": "number", "description": "As a decimal, e.g. 0.07 for 7%"},
                "years": {"type": "integer"},
            },
            "required": ["monthly_contribution", "annual_return", "years"],
        },
    },
    {
        "name": "assess_risk",
        "description": "Assess investor risk tolerance based on age, annual income, and total debt.",
        "input_schema": {
            "type": "object",
            "properties": {
                "age": {"type": "integer"},
                "annual_income": {"type": "number"},
                "total_debt": {"type": "number"},
            },
            "required": ["age", "annual_income", "total_debt"],
        },
    },
]

TOOL_DISPATCH = {
    "get_portfolio": get_portfolio,
    "get_stock_quote": get_stock_quote,
    "get_market_news": get_market_news,
    "project_savings": project_savings,
    "assess_risk": assess_risk,
}


# ---------------------------------------------------------------------------
# Chat agent — one turn at a time, each turn is its own trace
# ---------------------------------------------------------------------------

SYSTEM_PROMPT = (
    "You are a personal finance advisor. You help users understand their "
    "portfolios, check stock quotes, read market news, project savings growth, "
    "and assess risk tolerance. Be concise and actionable. Use the tools when "
    "you need data — don't make up numbers. If the user hasn't identified "
    "themselves, use 'default' as the user_id for portfolio lookups."
)


def run_turn(client, user_message: str, conversation_memory: list, turn_number: int) -> str:
    """
    Handle a single conversational turn.

    Each turn is wrapped in its own trace_run so it appears as a separate
    trace on the dashboard — but they're all grouped by the session_id that
    was set before this function is called.
    """
    with trace_run("FinanceAdvisorTurn", metadata={"turn": turn_number}):

        # Build system prompt with memory for continuity
        memory_block = ""
        if conversation_memory:
            memory_block = (
                "\n\nConversation so far (use this for continuity):\n"
                + "\n".join(conversation_memory[-12:])
            )

        messages = [{"role": "user", "content": user_message}]

        with trace_agent("AdvisorAgent"):
            # Agentic loop: LLM may call tools, then respond
            for _ in range(8):
                resp = client.messages.create(
                    model="claude-sonnet-4-20250514",
                    max_tokens=800,
                    system=SYSTEM_PROMPT + memory_block,
                    tools=TOOLS_SCHEMA,
                    messages=messages,
                )

                if resp.stop_reason == "end_turn":
                    break

                tool_results = []
                for block in resp.content:
                    if block.type == "tool_use":
                        fn = TOOL_DISPATCH.get(block.name)
                        result = fn(**block.input) if fn else f"Unknown tool: {block.name}"
                        tool_results.append({
                            "type": "tool_result",
                            "tool_use_id": block.id,
                            "content": result,
                        })

                if not tool_results:
                    break

                messages.append({"role": "assistant", "content": resp.content})
                messages.append({"role": "user", "content": tool_results})

        # Extract final text
        reply = "".join(
            b.text for b in resp.content if getattr(b, "type", None) == "text"
        ).strip()

        return reply


# ---------------------------------------------------------------------------
# Main — interactive CLI loop
# ---------------------------------------------------------------------------

def main():
    api_key = os.getenv("ANTHROPIC_API_KEY")
    if not api_key:
        print("ERROR: Set ANTHROPIC_API_KEY"); return
    if not os.getenv("CASCADE_API_KEY"):
        print("ERROR: Set CASCADE_API_KEY"); return

    # ── Initialize tracing ───────────────────────────────────────────────
    # init_tracing(
    #     project="finance_advisor_testing_1",
    #     session_evals=["Proceeding Despite Explicit User Refusal"],
    # )
    init_tracing(project="kevin_son_beast")
    client = wrap_llm_client(Anthropic(api_key=api_key))

    # ── Generate a unique session ID for this conversation ───────────────
    session_id = f"chat-{uuid.uuid4().hex[:8]}-{int(time.time())}"

    print()
    print("=" * 60)
    print("  Personal Finance Advisor  (multi-turn session demo)")
    print("=" * 60)
    print(f"  Session ID : {session_id}")
    print(f"  Project    : finance_advisor")
    print(f"  Dashboard  : http://localhost:3000")
    print()
    print("  Type your questions below. Type 'quit' or 'exit' to end.")
    print("  Each message becomes a separate trace, all grouped under")
    print("  the same session so you can see them together on the")
    print("  Sessions page.")
    print("=" * 60)
    print()

    # ── Set session ID in context ────────────────────────────────────────
    # Every trace_run() from here on inherits this session_id.
    set_session_id(session_id)

    conversation_memory = []
    turn = 0

    try:
        while True:
            try:
                user_input = input("You: ").strip()
            except EOFError:
                break

            if not user_input:
                continue
            if user_input.lower() in ("quit", "exit", "q"):
                break

            turn += 1
            conversation_memory.append(f"User: {user_input}")

            reply = run_turn(client, user_input, conversation_memory, turn)

            conversation_memory.append(f"Advisor: {reply}")
            print(f"\nAdvisor: {reply}\n")

    except KeyboardInterrupt:
        pass
    finally:
        # End session: clears context + triggers session-level evals in background
        end_session(session_id)

    print()
    print("=" * 60)
    print(f"  Session ended after {turn} turn(s).")
    print(f"  Session ID: {session_id}")
    print(f"  Session evals triggered in background.")
    print(f"  View on dashboard: http://localhost:3000  (Sessions page)")
    print("=" * 60)

    # Give the background session eval thread time to fire before exit
    time.sleep(8)
    print("  Done.")
    print()


if __name__ == "__main__":
    main()
