#!/usr/bin/env python3
"""
Customer Support Agent -- Prebuilt Evaluations

Demonstrates using Cascade SDK's built-in scorers on a real-time agent:

  init_tracing(evals=[...])  -- auto-evaluates every new trace

  That's it. Every trace produced by this agent will be automatically
  evaluated by the backend. Results appear in the dashboard.

Agent context:
  A customer support agent for an e-commerce platform ("ShopWave").
  The agent handles order lookup, return/refund processing, and
  inventory checks through a triage → resolution pipeline.

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

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, wrap_llm_client, tool, function,
)
from anthropic import Anthropic


# ---------------------------------------------------------------------------
# Helper functions (@function)
# ---------------------------------------------------------------------------

@function
def format_order_id(raw: str) -> str:
    """Normalize an order ID to standard format."""
    cleaned = raw.strip().upper().replace("-", "").replace(" ", "")
    if not cleaned.startswith("SW"):
        cleaned = "SW" + cleaned
    return cleaned


@function
def calculate_days_since(order_date: str) -> int:
    """Calculate days elapsed since the order date."""
    from datetime import datetime
    try:
        ordered = datetime.strptime(order_date, "%Y-%m-%d")
        return (datetime.now() - ordered).days
    except ValueError:
        return -1


@function
def apply_discount(amount: float, tier: str) -> float:
    """Apply loyalty-tier discount to a refund amount."""
    rates = {"gold": 1.0, "silver": 0.95, "bronze": 0.90, "standard": 0.85}
    return round(amount * rates.get(tier.lower(), 0.85), 2)


# ---------------------------------------------------------------------------
# Tools (@tool)
# ---------------------------------------------------------------------------

@tool
def search_orders(customer_email: str) -> str:
    """Look up recent orders for a customer by email."""
    orders_db = {
        "alice@example.com": [
            {"id": "SW-90210", "date": "2026-02-01", "items": "Wireless Headphones ($89.99), USB-C Hub ($34.99)", "status": "delivered", "total": 124.98},
            {"id": "SW-90185", "date": "2026-01-15", "items": "Running Shoes ($129.00)", "status": "delivered", "total": 129.00},
        ],
        "bob@example.com": [
            {"id": "SW-90322", "date": "2026-02-10", "items": "Mechanical Keyboard ($149.99), Desk Pad ($24.99)", "status": "shipped", "total": 174.98},
        ],
    }
    results = orders_db.get(customer_email.lower(), [])
    if not results:
        return f"No orders found for {customer_email}"
    lines = [f"Orders for {customer_email}:"]
    for o in results:
        oid = format_order_id(o["id"])
        age = calculate_days_since(o["date"])
        lines.append(f"  {oid} | {o['date']} ({age}d ago) | {o['status']} | ${o['total']:.2f}")
        lines.append(f"    Items: {o['items']}")
    return "\n".join(lines)


@tool
def check_inventory(product_name: str) -> str:
    """Check current stock level for a product."""
    stock = {
        "wireless headphones": {"in_stock": True, "qty": 142, "warehouse": "West"},
        "usb-c hub":           {"in_stock": True, "qty": 38,  "warehouse": "East"},
        "running shoes":       {"in_stock": False, "qty": 0,  "warehouse": "N/A", "restock_date": "2026-03-01"},
        "mechanical keyboard": {"in_stock": True, "qty": 67,  "warehouse": "West"},
        "desk pad":            {"in_stock": True, "qty": 210, "warehouse": "West"},
    }
    info = stock.get(product_name.lower())
    if not info:
        return f"Product '{product_name}' not found in catalog."
    if info["in_stock"]:
        return f"{product_name}: IN STOCK ({info['qty']} units, {info['warehouse']} warehouse)"
    return f"{product_name}: OUT OF STOCK (expected restock: {info.get('restock_date', 'TBD')})"


@tool
def get_return_policy(order_id: str) -> str:
    """Get applicable return policy for an order."""
    oid = format_order_id(order_id)
    policies = {
        "SW90210": {"window_days": 30, "condition": "unopened or defective", "method": "prepaid label"},
        "SW90185": {"window_days": 60, "condition": "unworn, tags attached", "method": "in-store or prepaid label"},
        "SW90322": {"window_days": 30, "condition": "unopened", "method": "prepaid label"},
    }
    key = oid.replace("-", "").replace("SW", "SW").replace("SW", "")
    key = "SW" + key.replace("SW", "")
    lookup = key.replace("-", "")
    p = policies.get(lookup)
    if not p:
        return f"No policy found for order {oid}. Standard 30-day return window applies."
    return (
        f"Return policy for {oid}:\n"
        f"  Window: {p['window_days']} days from delivery\n"
        f"  Condition: {p['condition']}\n"
        f"  Method: {p['method']}"
    )


@tool
def calculate_refund(order_id: str, item_name: str, reason: str) -> str:
    """Calculate refund amount for a returned item."""
    prices = {
        "wireless headphones": 89.99,
        "usb-c hub": 34.99,
        "running shoes": 129.00,
        "mechanical keyboard": 149.99,
        "desk pad": 24.99,
    }
    price = prices.get(item_name.lower(), 0)
    if price == 0:
        return f"Item '{item_name}' not found. Cannot calculate refund."

    oid = format_order_id(order_id)
    tier = "gold" if "90210" in oid else "silver"
    refund_base = price if reason.lower() in ("defective", "damaged", "wrong item") else price * 0.90
    final = apply_discount(refund_base, tier)
    return (
        f"Refund estimate for {item_name} (order {oid}):\n"
        f"  Original price: ${price:.2f}\n"
        f"  Reason: {reason}\n"
        f"  Customer tier: {tier}\n"
        f"  Refund amount: ${final:.2f}"
    )


# ---------------------------------------------------------------------------
# Tool schema + dispatch
# ---------------------------------------------------------------------------

TOOLS_SCHEMA = [
    {"name": "search_orders", "description": "Look up recent orders for a customer by email.",
     "input_schema": {"type": "object", "properties": {"customer_email": {"type": "string"}}, "required": ["customer_email"]}},
    {"name": "check_inventory", "description": "Check current stock level for a product.",
     "input_schema": {"type": "object", "properties": {"product_name": {"type": "string"}}, "required": ["product_name"]}},
    {"name": "get_return_policy", "description": "Get applicable return policy for an order.",
     "input_schema": {"type": "object", "properties": {"order_id": {"type": "string"}}, "required": ["order_id"]}},
    {"name": "calculate_refund", "description": "Calculate refund amount for a returned item.",
     "input_schema": {"type": "object", "properties": {"order_id": {"type": "string"}, "item_name": {"type": "string"}, "reason": {"type": "string"}}, "required": ["order_id", "item_name", "reason"]}},
]

TOOL_DISPATCH = {
    "search_orders": search_orders,
    "check_inventory": check_inventory,
    "get_return_policy": get_return_policy,
    "calculate_refund": calculate_refund,
}


# ---------------------------------------------------------------------------
# Agent runner
# ---------------------------------------------------------------------------

def run_support_agent(client, customer_request: str):
    """Run the customer support agent. Returns (trace_id, root_span)."""

    with trace_run("CustomerSupportSession", metadata={"channel": "chat", "department": "returns"}) as root:
        trace_id = format(root.get_span_context().trace_id, "032x")
        messages = [{"role": "user", "content": customer_request}]

        # Phase 1 -- Triage: identify the issue and look up context
        with trace_agent("TriageAgent"):
            for _ in range(5):
                resp = client.messages.create(
                    model="claude-3-5-haiku-20241022",
                    max_tokens=600,
                    system=(
                        "You are a customer support triage agent for ShopWave e-commerce. "
                        "Use the tools to look up the customer's orders and gather context. "
                        "Be thorough -- check orders, inventory, and return policies as needed."
                    ),
                    tools=TOOLS_SCHEMA,
                    messages=messages,
                )
                if resp.stop_reason == "end_turn":
                    break
                results = []
                for b in resp.content:
                    if b.type == "tool_use":
                        fn = TOOL_DISPATCH.get(b.name)
                        r = fn(**b.input) if fn else f"Unknown tool: {b.name}"
                        results.append({"type": "tool_result", "tool_use_id": b.id, "content": r})
                if not results:
                    break
                messages.append({"role": "assistant", "content": resp.content})
                messages.append({"role": "user", "content": results})

        # Phase 2 -- Resolution: compose a helpful response
        triage_text = "".join(b.text for b in resp.content if hasattr(b, "text"))
        with trace_agent("ResolutionAgent"):
            client.messages.create(
                model="claude-3-5-haiku-20241022",
                max_tokens=800,
                system=(
                    "You are a customer support resolution agent for ShopWave. "
                    "Based on the triage information, compose a clear, empathetic "
                    "response with specific next steps for the customer."
                ),
                messages=[{"role": "user", "content": f"Triage summary:\n{triage_text}\n\nCompose a final response for the customer."}],
            )

    return trace_id, root


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

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

    print("=" * 60)
    print("  Customer Support Agent -- Prebuilt Evaluations")
    print("=" * 60)

    # ------------------------------------------------------------------
    # 1. Init tracing with prebuilt evals (auto-evaluate every trace)
    # ------------------------------------------------------------------
    print("\n[1] Initializing tracing with prebuilt evals...")
    init_tracing(
        project="customer_support",
        evals=["helpfulness", "conciseness"],
    )
    client = wrap_llm_client(Anthropic(api_key=api_key))
    print("    project = customer_support")
    print("    evals   = ['helpfulness', 'conciseness']")
    print("    Every new trace will be auto-evaluated.\n")

    # ------------------------------------------------------------------
    # 2. Run agent -- return request
    # ------------------------------------------------------------------
    print("[2] Running support agent -- return request...")
    request_1 = (
        "Hi, I'm Alice (alice@example.com). I bought Wireless Headphones "
        "a couple weeks ago (order SW-90210) and they have a buzzing issue "
        "in the left ear. I'd like to return them for a refund. Can you "
        "help me with that?"
    )
    trace_id_1, _ = run_support_agent(client, request_1)
    print(f"    Trace: {trace_id_1}")

    # ------------------------------------------------------------------
    # 3. Run agent -- stock inquiry
    # ------------------------------------------------------------------
    print("\n[3] Running support agent -- stock inquiry...")
    request_2 = (
        "Hey, I'm Bob (bob@example.com). I ordered a Mechanical Keyboard "
        "last week and it still says shipped. Also, are Running Shoes back "
        "in stock yet? I want to order a pair."
    )
    trace_id_2, _ = run_support_agent(client, request_2)
    print(f"    Trace: {trace_id_2}")

    # ------------------------------------------------------------------
    # Done
    # ------------------------------------------------------------------
    print("\n" + "=" * 60)
    print("  Done!")
    print("=" * 60)
    print("  - Prebuilt scorers will auto-evaluate both traces via init_tracing(evals=...)")
    print("  - View results at: http://localhost:3000  (project: customer_support)")
    print()


if __name__ == "__main__":
    main()
