#!/usr/bin/env python3
"""
Span-Level Evaluations -- LLM Spans & Tool Spans

Demonstrates evaluate_spans() on both LLM and tool spans:

  1. Create a span-level scorer for LLM calls
  2. Create a span-level scorer for tool calls
  3. Run an agent to produce a trace
  4. evaluate_spans(span_type="llm")  -- evaluate all LLM spans
  5. evaluate_spans(span_type="tool") -- evaluate all tool spans

Agent context:
  A recipe assistant that looks up recipes, checks ingredient
  availability, calculates nutrition, and suggests substitutions.

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

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,
    CascadeEval,
)
from anthropic import Anthropic


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

@function
def parse_ingredients(raw: str) -> list:
    """Parse a comma-separated ingredient list."""
    return [i.strip() for i in raw.split(",") if i.strip()]


@function
def calculate_calories(protein_g: float, carbs_g: float, fat_g: float) -> int:
    """Calculate total calories from macronutrients."""
    return int(protein_g * 4 + carbs_g * 4 + fat_g * 9)


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

@tool
def search_recipes(query: str, dietary_preference: str = "none") -> str:
    """Search for recipes matching a query and dietary preference."""
    recipes_db = {
        "pasta": [
            {"name": "Classic Carbonara", "time": "25 min", "difficulty": "Medium",
             "ingredients": "spaghetti, eggs, pancetta, parmesan, black pepper"},
            {"name": "Pesto Penne", "time": "20 min", "difficulty": "Easy",
             "ingredients": "penne, basil, pine nuts, garlic, olive oil, parmesan"},
        ],
        "salad": [
            {"name": "Caesar Salad", "time": "15 min", "difficulty": "Easy",
             "ingredients": "romaine, croutons, parmesan, caesar dressing, lemon"},
            {"name": "Greek Salad", "time": "10 min", "difficulty": "Easy",
             "ingredients": "cucumber, tomato, red onion, olives, feta, olive oil"},
        ],
        "soup": [
            {"name": "Tomato Basil Soup", "time": "30 min", "difficulty": "Easy",
             "ingredients": "tomatoes, basil, onion, garlic, cream, vegetable broth"},
            {"name": "French Onion Soup", "time": "45 min", "difficulty": "Medium",
             "ingredients": "onions, beef broth, gruyere, baguette, butter, thyme"},
        ],
    }
    key = query.lower().strip()
    matches = recipes_db.get(key, [])
    if not matches:
        for k, v in recipes_db.items():
            if key in k:
                matches = v
                break
    if not matches:
        return f"No recipes found for '{query}'."
    lines = [f"Recipes for '{query}':"]
    for r in matches:
        lines.append(f"  - {r['name']} ({r['time']}, {r['difficulty']})")
        lines.append(f"    Ingredients: {r['ingredients']}")
    return "\n".join(lines)


@tool
def check_ingredients(ingredient_list: str) -> str:
    """Check availability and freshness of ingredients."""
    items = parse_ingredients(ingredient_list)
    stock = {
        "spaghetti": {"available": True, "fresh": True},
        "eggs": {"available": True, "fresh": True},
        "pancetta": {"available": True, "fresh": True},
        "parmesan": {"available": True, "fresh": True},
        "black pepper": {"available": True, "fresh": True},
        "penne": {"available": True, "fresh": True},
        "basil": {"available": True, "fresh": False},
        "pine nuts": {"available": False, "fresh": False},
        "garlic": {"available": True, "fresh": True},
        "olive oil": {"available": True, "fresh": True},
        "tomatoes": {"available": True, "fresh": True},
        "onion": {"available": True, "fresh": True},
        "cream": {"available": True, "fresh": True},
        "vegetable broth": {"available": True, "fresh": True},
    }
    lines = ["Ingredient check:"]
    for item in items:
        info = stock.get(item.lower(), {"available": False, "fresh": False})
        status = "IN STOCK" if info["available"] else "OUT OF STOCK"
        fresh = ", fresh" if info["fresh"] else ", not fresh" if info["available"] else ""
        lines.append(f"  {item}: {status}{fresh}")
    return "\n".join(lines)


@tool
def get_nutrition(recipe_name: str) -> str:
    """Get nutritional information for a recipe (per serving)."""
    nutrition_db = {
        "classic carbonara": {"protein": 22, "carbs": 65, "fat": 18, "fiber": 3, "servings": 4},
        "pesto penne": {"protein": 14, "carbs": 58, "fat": 24, "fiber": 4, "servings": 4},
        "caesar salad": {"protein": 8, "carbs": 12, "fat": 15, "fiber": 3, "servings": 2},
        "tomato basil soup": {"protein": 4, "carbs": 18, "fat": 8, "fiber": 4, "servings": 4},
    }
    info = nutrition_db.get(recipe_name.lower())
    if not info:
        return f"No nutrition data for '{recipe_name}'."
    cals = calculate_calories(info["protein"], info["carbs"], info["fat"])
    return (
        f"Nutrition for {recipe_name} (per serving, {info['servings']} servings):\n"
        f"  Calories: {cals} kcal\n"
        f"  Protein: {info['protein']}g | Carbs: {info['carbs']}g | Fat: {info['fat']}g\n"
        f"  Fiber: {info['fiber']}g"
    )


@tool
def suggest_substitutions(ingredient: str, reason: str) -> str:
    """Suggest ingredient substitutions."""
    subs = {
        "pine nuts": ["walnuts", "cashews", "sunflower seeds", "almonds"],
        "pancetta": ["bacon", "prosciutto", "smoked tofu (vegan)"],
        "cream": ["coconut cream", "cashew cream", "silken tofu blend"],
        "parmesan": ["pecorino romano", "nutritional yeast (vegan)", "aged asiago"],
        "basil": ["dried basil (use 1/3 amount)", "spinach + mint mix", "arugula"],
    }
    options = subs.get(ingredient.lower(), [])
    if not options:
        return f"No substitutions found for '{ingredient}'."
    lines = [f"Substitutions for {ingredient} ({reason}):"]
    for s in options:
        lines.append(f"  - {s}")
    return "\n".join(lines)


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

TOOLS_SCHEMA = [
    {"name": "search_recipes", "description": "Search for recipes matching a query and dietary preference.",
     "input_schema": {"type": "object", "properties": {"query": {"type": "string"}, "dietary_preference": {"type": "string"}}, "required": ["query"]}},
    {"name": "check_ingredients", "description": "Check availability and freshness of ingredients.",
     "input_schema": {"type": "object", "properties": {"ingredient_list": {"type": "string"}}, "required": ["ingredient_list"]}},
    {"name": "get_nutrition", "description": "Get nutritional information for a recipe (per serving).",
     "input_schema": {"type": "object", "properties": {"recipe_name": {"type": "string"}}, "required": ["recipe_name"]}},
    {"name": "suggest_substitutions", "description": "Suggest ingredient substitutions.",
     "input_schema": {"type": "object", "properties": {"ingredient": {"type": "string"}, "reason": {"type": "string"}}, "required": ["ingredient", "reason"]}},
]

TOOL_DISPATCH = {
    "search_recipes": search_recipes,
    "check_ingredients": check_ingredients,
    "get_nutrition": get_nutrition,
    "suggest_substitutions": suggest_substitutions,
}


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

def run_recipe_agent(client, user_request: str):
    """Run the recipe assistant. Returns (trace_id, root_span)."""

    with trace_run("RecipeAssistant", metadata={"request_type": "meal_planning"}) as root:
        trace_id = format(root.get_span_context().trace_id, "032x")
        messages = [{"role": "user", "content": user_request}]

        # Phase 1 -- Research: search recipes and check ingredients
        with trace_agent("ResearchAgent"):
            for _ in range(6):
                resp = client.messages.create(
                    model="claude-3-5-haiku-20241022",
                    max_tokens=600,
                    system=(
                        "You are a recipe research assistant. Use the tools to find recipes, "
                        "check ingredient availability, get nutrition info, and suggest "
                        "substitutions for missing ingredients. Be thorough."
                    ),
                    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 -- Summary: compile a meal plan
        research_text = "".join(b.text for b in resp.content if hasattr(b, "text"))
        with trace_agent("PlannerAgent"):
            client.messages.create(
                model="claude-3-5-haiku-20241022",
                max_tokens=800,
                system=(
                    "You are a meal planner. Based on the research, compile a clear "
                    "meal recommendation with recipe, nutrition summary, ingredient "
                    "notes, and any substitutions needed."
                ),
                messages=[{"role": "user", "content": f"Research results:\n{research_text}\n\nCompile a meal recommendation."}],
            )

    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("  Span-Level Evaluations -- LLM & Tool Spans")
    print("=" * 60)

    # ------------------------------------------------------------------
    # 1. Init tracing
    # ------------------------------------------------------------------
    print("\n[1] Initializing tracing...")
    init_tracing(project="recipe_assistant")
    client = wrap_llm_client(Anthropic(api_key=api_key))
    print("    project = recipe_assistant\n")

    # ------------------------------------------------------------------
    # 2. Create span-level scorers
    # ------------------------------------------------------------------
    print("[2] Creating span-level scorers...")
    evals = CascadeEval()
    created_ids = []

    # Scorer A: LLM span scorer
    llm_scorer = evals.create_scorer(
        name="LLM Response Relevance",
        scorer_type="llm_judge",
        scope="span",
        threshold=0.5,
        description="Evaluates whether each LLM response is relevant and well-structured",
        llm_provider="anthropic",
        llm_model="claude-3-5-haiku-20241022",
        output_type="categorical",
        choices=[
            {"label": "Y", "score": 1.0, "description": "Response is relevant and useful"},
            {"label": "N", "score": 0.0, "description": "Response is off-topic or poorly structured"},
        ],
        llm_template=(
            "Evaluate the relevance of this LLM response.\n\n"
            "Prompt:\n{{prompt}}\n\n"
            "Response:\n{{completion}}\n\n"
            "Is the response directly relevant to the prompt, well-structured, "
            "and does it make appropriate use of available tools?\n"
            "Answer Y if relevant and useful, N if off-topic or poorly structured.\n\n"
            "Answer (Y/N):"
        ),
        variable_mappings={"prompt": "llm.prompt", "completion": "llm.completion"},
        tags=["relevance", "span-level", "llm"],
    )
    created_ids.append(llm_scorer["scorer_id"])
    print(f"    Created: {llm_scorer['name']} (span-level, for LLM spans)")

    # Scorer B: Tool span scorer
    tool_scorer = evals.create_scorer(
        name="Tool Call Appropriateness",
        scorer_type="llm_judge",
        scope="span",
        threshold=0.5,
        description="Evaluates whether each tool call used correct parameters and returned useful results",
        llm_provider="anthropic",
        llm_model="claude-3-5-haiku-20241022",
        output_type="categorical",
        choices=[
            {"label": "Y", "score": 1.0, "description": "Tool call was appropriate"},
            {"label": "N", "score": 0.0, "description": "Tool call was unnecessary or had wrong params"},
        ],
        llm_template=(
            "Evaluate this tool call.\n\n"
            "Tool: {{tool_name}}\n"
            "Input: {{tool_input}}\n"
            "Output: {{tool_output}}\n\n"
            "Was the tool called with appropriate parameters for the task at hand? "
            "Did it return useful, non-empty results?\n"
            "Answer Y if the call was appropriate, N if it was unnecessary or had wrong parameters.\n\n"
            "Answer (Y/N):"
        ),
        variable_mappings={
            "tool_name": "tool.name",
            "tool_input": "tool.input",
            "tool_output": "tool.output",
        },
        tags=["appropriateness", "span-level", "tool"],
    )
    created_ids.append(tool_scorer["scorer_id"])
    print(f"    Created: {tool_scorer['name']} (span-level, for tool spans)")
    print()

    # ------------------------------------------------------------------
    # 3. Run agent
    # ------------------------------------------------------------------
    print("[3] Running recipe assistant agent...")
    request = (
        "I want to make a nice pasta dinner tonight. Can you find me a good "
        "recipe, check if I have the ingredients, give me the nutrition info, "
        "and suggest substitutions for anything that's missing or not fresh?"
    )
    trace_id, _ = run_recipe_agent(client, request)
    print(f"    Trace: {trace_id}")

    # ------------------------------------------------------------------
    # 4. Wait for ingestion
    # ------------------------------------------------------------------
    print("\n[4] Waiting 5s for trace ingestion...")
    time.sleep(5)

    # ------------------------------------------------------------------
    # 5. Evaluate LLM spans
    # ------------------------------------------------------------------
    print("\n[5] Evaluating LLM spans...")
    llm_resp = evals.evaluate_spans(
        trace_id=trace_id,
        scorer_ids=[llm_scorer["scorer_id"]],
        span_type="llm",
    )
    print(f"    Success: {llm_resp.get('success')}")
    for r in llm_resp.get("results", []):
        tag = "PASS" if r.get("passed") else "FAIL"
        span_label = r.get("span_name", "?")[:35]
        print(f"    [{tag}] {span_label}: score={r.get('score')}")
    if llm_resp.get("errors"):
        for e in llm_resp["errors"]:
            print(f"    ERROR: {e}")

    # ------------------------------------------------------------------
    # 6. Evaluate tool spans
    # ------------------------------------------------------------------
    print("\n[6] Evaluating tool spans...")
    tool_resp = evals.evaluate_spans(
        trace_id=trace_id,
        scorer_ids=[tool_scorer["scorer_id"]],
        span_type="tool",
    )
    print(f"    Success: {tool_resp.get('success')}")
    for r in tool_resp.get("results", []):
        tag = "PASS" if r.get("passed") else "FAIL"
        span_label = r.get("span_name", "?")[:35]
        print(f"    [{tag}] {span_label}: score={r.get('score')}")
    if tool_resp.get("errors"):
        for e in tool_resp["errors"]:
            print(f"    ERROR: {e}")

    # ------------------------------------------------------------------
    # Done
    # ------------------------------------------------------------------
    print("\n" + "=" * 60)
    print("  Done!")
    print("=" * 60)
    print("  - Created 2 span-level scorers (1 for LLM, 1 for tool)")
    print("  - Ran recipe assistant agent")
    print("  - evaluate_spans(span_type='llm')  -- evaluated all LLM call spans")
    print("  - evaluate_spans(span_type='tool') -- evaluated all tool call spans")
    print("  - View results at: http://localhost:3000  (project: recipe_assistant)")
    print()


if __name__ == "__main__":
    main()
