#!/usr/bin/env python3
"""
Real-Time Custom Evaluations

Demonstrates creating custom scorers and running them in real-time
on every incoming trace via init_tracing(evals=[...]):

  1. CascadeEval.create_scorer() -- create two custom scorers
  2. init_tracing(evals=[scorer_id_1, scorer_id_2]) -- pass their IDs
  3. Run agents -- every trace is auto-evaluated with the custom scorers

The key insight: init_tracing(evals=...) accepts both built-in names
("helpfulness") AND custom scorer UUIDs. So you can create your own
rubrics and have them run automatically on every new trace.

Agent context:
  A tech support agent for a SaaS platform ("CloudBase") that
  diagnoses issues, checks system status, and guides users through
  troubleshooting steps.

Usage:
    export CASCADE_API_KEY="csk_live_..."
    export CASCADE_ENDPOINT="http://localhost:8000/v1/traces"
    export ANTHROPIC_API_KEY="sk-ant-..."
    python tests/test_realtime_custom_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,
    CascadeEval,
)
from anthropic import Anthropic


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

@function
def parse_error_code(raw: str) -> dict:
    """Parse an error code into component parts."""
    parts = raw.strip().upper().split("-")
    return {
        "raw": raw,
        "service": parts[0] if len(parts) > 0 else "UNKNOWN",
        "severity": parts[1] if len(parts) > 1 else "UNKNOWN",
        "id": parts[2] if len(parts) > 2 else "0000",
    }


@function
def format_uptime(seconds: int) -> str:
    """Format seconds into human-readable uptime."""
    days = seconds // 86400
    hours = (seconds % 86400) // 3600
    mins = (seconds % 3600) // 60
    return f"{days}d {hours}h {mins}m"


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

@tool
def check_service_status(service_name: str) -> str:
    """Check the current status of a CloudBase service."""
    statuses = {
        "api": {"status": "operational", "uptime": 2592000, "latency_ms": 45, "error_rate": 0.001},
        "database": {"status": "degraded", "uptime": 1296000, "latency_ms": 320, "error_rate": 0.05,
                     "incident": "Elevated latency due to replica lag"},
        "auth": {"status": "operational", "uptime": 5184000, "latency_ms": 12, "error_rate": 0.0},
        "storage": {"status": "operational", "uptime": 3888000, "latency_ms": 85, "error_rate": 0.002},
        "cdn": {"status": "maintenance", "uptime": 0, "latency_ms": 0, "error_rate": 0,
                "incident": "Scheduled maintenance window until 06:00 UTC"},
    }
    info = statuses.get(service_name.lower())
    if not info:
        return f"Service '{service_name}' not found. Available: api, database, auth, storage, cdn"
    uptime_str = format_uptime(info["uptime"])
    lines = [f"Service: {service_name.upper()}"]
    lines.append(f"  Status: {info['status']}")
    lines.append(f"  Uptime: {uptime_str}")
    lines.append(f"  Latency: {info['latency_ms']}ms")
    lines.append(f"  Error rate: {info['error_rate']:.3%}")
    if info.get("incident"):
        lines.append(f"  Active incident: {info['incident']}")
    return "\n".join(lines)


@tool
def lookup_error(error_code: str) -> str:
    """Look up an error code in the knowledge base."""
    parsed = parse_error_code(error_code)
    errors_db = {
        "DB-WARN-4201": {
            "title": "Connection Pool Exhausted",
            "cause": "All database connections are in use",
            "fix": "1. Check for unclosed connections in your code\n"
                   "2. Increase pool size in config (max_connections)\n"
                   "3. Add connection timeouts (idle_timeout: 30s)",
        },
        "DB-ERR-5001": {
            "title": "Replica Sync Failure",
            "cause": "Read replica fell behind primary by >10s",
            "fix": "1. This usually resolves within 5-10 minutes\n"
                   "2. If persistent, check replica disk I/O\n"
                   "3. Consider routing reads to primary temporarily",
        },
        "AUTH-ERR-4010": {
            "title": "Token Validation Failed",
            "cause": "JWT signature mismatch or token expired",
            "fix": "1. Check if token was issued by correct auth service\n"
                   "2. Verify clock sync between services (<5s drift)\n"
                   "3. Refresh token and retry",
        },
        "API-WARN-4290": {
            "title": "Rate Limit Exceeded",
            "cause": "Client exceeded 1000 req/min threshold",
            "fix": "1. Implement exponential backoff\n"
                   "2. Cache responses where possible\n"
                   "3. Request rate limit increase via support portal",
        },
    }
    entry = errors_db.get(error_code.upper())
    if not entry:
        return f"Error code '{error_code}' not found in knowledge base. Parsed: {parsed}"
    return (
        f"Error: {error_code}\n"
        f"  Title: {entry['title']}\n"
        f"  Cause: {entry['cause']}\n"
        f"  Resolution:\n  {entry['fix']}"
    )


@tool
def run_diagnostic(service_name: str, check_type: str) -> str:
    """Run a diagnostic check on a service. check_type: connectivity, performance, logs."""
    diagnostics = {
        ("database", "connectivity"): "Connectivity OK: 3/3 replicas responding, primary reachable (2ms)",
        ("database", "performance"): "Performance DEGRADED: primary write latency 45ms (normal), replica read latency 320ms (HIGH, threshold: 100ms)",
        ("database", "logs"): "Recent logs:\n  [WARN] Replica lag 8.2s on replica-2\n  [WARN] Slow query: SELECT * FROM events WHERE... (2.1s)\n  [INFO] Autovacuum completed on table 'traces'",
        ("api", "connectivity"): "Connectivity OK: all endpoints responding, health check passing",
        ("api", "performance"): "Performance OK: p50=12ms, p95=45ms, p99=120ms",
        ("api", "logs"): "Recent logs:\n  [INFO] 12,450 requests in last hour\n  [WARN] 3 requests exceeded 5s timeout\n  [INFO] Cache hit rate: 89%",
        ("auth", "connectivity"): "Connectivity OK: auth service responding, JWKS endpoint cached",
        ("auth", "performance"): "Performance OK: token validation p95=8ms",
    }
    result = diagnostics.get((service_name.lower(), check_type.lower()))
    if not result:
        return f"No diagnostic available for {service_name}/{check_type}"
    return f"Diagnostic [{service_name}/{check_type}]:\n  {result}"


@tool
def suggest_resolution(issue_summary: str) -> str:
    """Generate resolution steps based on the issue summary."""
    return (
        f"Recommended resolution for: {issue_summary}\n"
        "  Step 1: Verify the issue is ongoing (check service status)\n"
        "  Step 2: Review recent changes or deployments\n"
        "  Step 3: Apply the specific fix from the error lookup\n"
        "  Step 4: Monitor for 15 minutes after applying fix\n"
        "  Step 5: If issue persists, escalate to on-call engineer"
    )


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

TOOLS_SCHEMA = [
    {"name": "check_service_status", "description": "Check the current status of a CloudBase service.",
     "input_schema": {"type": "object", "properties": {"service_name": {"type": "string"}}, "required": ["service_name"]}},
    {"name": "lookup_error", "description": "Look up an error code in the knowledge base.",
     "input_schema": {"type": "object", "properties": {"error_code": {"type": "string"}}, "required": ["error_code"]}},
    {"name": "run_diagnostic", "description": "Run a diagnostic check on a service. check_type: connectivity, performance, logs.",
     "input_schema": {"type": "object", "properties": {"service_name": {"type": "string"}, "check_type": {"type": "string"}}, "required": ["service_name", "check_type"]}},
    {"name": "suggest_resolution", "description": "Generate resolution steps based on the issue summary.",
     "input_schema": {"type": "object", "properties": {"issue_summary": {"type": "string"}}, "required": ["issue_summary"]}},
]

TOOL_DISPATCH = {
    "check_service_status": check_service_status,
    "lookup_error": lookup_error,
    "run_diagnostic": run_diagnostic,
    "suggest_resolution": suggest_resolution,
}


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

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

    with trace_run("TechSupportSession", metadata={"platform": "CloudBase"}) as root:
        trace_id = format(root.get_span_context().trace_id, "032x")
        messages = [{"role": "user", "content": user_issue}]

        # Phase 1 -- Diagnosis
        with trace_agent("DiagnosisAgent"):
            for _ in range(6):
                resp = client.messages.create(
                    model="claude-3-5-haiku-20241022",
                    max_tokens=600,
                    system=(
                        "You are a tech support diagnosis agent for CloudBase. "
                        "Use the tools to check service status, look up error codes, "
                        "and run diagnostics. Be thorough -- gather all relevant information."
                    ),
                    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
        diagnosis_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 tech support resolution agent. Based on the diagnosis, "
                    "provide a clear, step-by-step resolution guide for the user. "
                    "Include what caused the issue and how to prevent it in the future."
                ),
                messages=[{"role": "user", "content": f"Diagnosis:\n{diagnosis_text}\n\nProvide a resolution guide."}],
            )

    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("  Real-Time Custom Evaluations")
    print("=" * 60)

    # ------------------------------------------------------------------
    # 1. Create custom scorers FIRST (before init_tracing)
    # ------------------------------------------------------------------
    print("\n[1] Creating custom scorers...")
    evals = CascadeEval()

    # Scorer A: Trace-level -- how thorough was the diagnosis? (numeric 0-1)
    diagnosis_quality = evals.create_scorer(
        name="Diagnosis Thoroughness",
        scorer_type="llm_judge",
        scope="trace",
        threshold=0.6,
        description="Rates how thoroughly the agent diagnosed the issue on a 0-1 scale",
        llm_provider="anthropic",
        llm_model="claude-3-5-haiku-20241022",
        output_type="numeric",
        min_score=0,
        max_score=1,
        llm_template=(
            "Rate the thoroughness of this tech support agent's diagnosis on a "
            "scale of 0.0 to 1.0.\n\n"
            "Full execution trajectory:\n{{trajectory}}\n\n"
            "Scoring criteria:\n"
            "  0.0-0.3: Jumped to conclusions with little or no investigation\n"
            "  0.3-0.5: Checked one thing but missed other relevant diagnostics\n"
            "  0.5-0.7: Checked service status and looked up errors, but skipped diagnostics\n"
            "  0.7-0.9: Thorough -- checked status, looked up errors, and ran diagnostics\n"
            "  0.9-1.0: Exceptional -- exhaustive investigation across all relevant services\n\n"
            "Score (0.0-1.0):"
        ),
        variable_mappings={"trajectory": "trajectory"},
        tags=["tech-support", "diagnosis", "trace-level"],
    )
    print(f"    Created: {diagnosis_quality['name']} (trace-level, numeric 0-1)")
    print(f"    ID: {diagnosis_quality['scorer_id']}")

    # Scorer B: Span-level -- are tool calls appropriate?
    tool_quality = evals.create_scorer(
        name="Tool Usage Quality",
        scorer_type="llm_judge",
        scope="span",
        threshold=0.5,
        description="Evaluates if each tool call was appropriate and used correct parameters",
        llm_provider="anthropic",
        llm_model="claude-3-5-haiku-20241022",
        output_type="categorical",
        choices=[
            {"label": "Y", "score": 1.0, "description": "Appropriate tool usage"},
            {"label": "N", "score": 0.0, "description": "Unnecessary or incorrect usage"},
        ],
        llm_template=(
            "Evaluate this tool call in a tech support context.\n\n"
            "Tool: {{tool_name}}\n"
            "Input: {{tool_input}}\n"
            "Output: {{tool_output}}\n\n"
            "Was this tool call appropriate for the diagnosis? "
            "Were the parameters correct and specific enough?\n"
            "Answer Y if appropriate, N if unnecessary or poorly parameterized.\n\n"
            "Answer (Y/N):"
        ),
        variable_mappings={
            "tool_name": "tool.name",
            "tool_input": "tool.input",
            "tool_output": "tool.output",
        },
        tags=["tech-support", "tool-quality", "span-level"],
    )
    print(f"    Created: {tool_quality['name']} (span-level)")
    print(f"    ID: {tool_quality['scorer_id']}")

    # ------------------------------------------------------------------
    # 2. Init tracing with CUSTOM scorer IDs for real-time evaluation
    # ------------------------------------------------------------------
    print("\n[2] Initializing tracing with custom scorers for real-time eval...")
    init_tracing(
        project="tech_support_test_3",
        evals=["Redundant Repetition of User Prompt Without Added Progress", "Tool Usage Quality"],
    )
    client = wrap_llm_client(Anthropic(api_key=api_key))
    print("    project = tech_support")
    print(f"    evals = [{diagnosis_quality['name']}, {tool_quality['name']}]")
    print("    Every new trace will be auto-evaluated with BOTH custom scorers.\n")

    # ------------------------------------------------------------------
    # 3. Run agents -- these traces will be auto-evaluated in real-time
    # ------------------------------------------------------------------
    print("[3] Running tech support agent -- database issue...")
    issue_1 = (
        "Hi, I'm getting error DB-WARN-4201 when my app tries to connect "
        "to the database. Queries are timing out and some of my API calls "
        "are failing with 500 errors. Can you help?"
    )
    trace_id_1, _ = run_support_agent(client, issue_1)
    print(f"    Trace: {trace_id_1}")
    print("    (auto-evaluation triggered in background)")

    print("\n[4] Running tech support agent -- auth issue...")
    issue_2 = (
        "Our users are reporting AUTH-ERR-4010 when trying to log in. "
        "It started about an hour ago. Is there a known issue?"
    )
    trace_id_2, _ = run_support_agent(client, issue_2)
    print(f"    Trace: {trace_id_2}")
    print("    (auto-evaluation triggered in background)")

    # ------------------------------------------------------------------
    # Done
    # ------------------------------------------------------------------
    print("\n" + "=" * 60)
    print("  Done!")
    print("=" * 60)
    print("  - Created 2 custom scorers (1 trace-level, 1 span-level)")
    print("  - Passed their IDs to init_tracing(evals=[...]) for real-time eval")
    print("  - Ran 2 tech support traces -- both auto-evaluated by the backend")
    print("  - The trace-level scorer runs once per trace")
    print("  - The span-level scorer runs on every tool call span in the trace")
    print("  - No evaluate() calls needed -- it's all automatic!")
    print("  - View results at: http://localhost:3000  (project: tech_support)")
    print()


if __name__ == "__main__":
    main()
