"""Comprehensive test agent for Staso SDK + UI integration.

Exercises EVERY UI feature:
  - Multiple environments (production, staging, default)
  - Workspace scoping (derived from API key)
  - Multiple sessions with grouped traces
  - All span kinds (agent, tool, chain, retriever, llm)
  - Nested span trees (deep parent-child hierarchies)
  - Rich input/output for preview columns
  - Multiple models for cost estimation
  - Error and timeout scenarios
  - Tags on spans
  - High token counts for cost breakdown
  - Varied durations for waterfall visualization

Supports both Anthropic and OpenAI providers via PROVIDER env var.

Usage:
    # 1. Start the Staso backend
    # 2. Install dependencies:
    #      cd .. && pip install -e ".[anthropic,openai]"
    #      pip install -r requirements.txt
    # 3. Set environment variables:
    #      export PROVIDER="anthropic"              (or "openai")
    #      export ANTHROPIC_API_KEY="sk-ant-..."    (if provider=anthropic)
    #      export OPENAI_API_KEY="sk-..."           (if provider=openai)
    #      export STASO_API_KEY="ak_..."
    #      export STASO_BASE_URL="http://localhost:6999"
    # 4. Run:
    #      python agent.py
"""
from __future__ import annotations

import json
import os
import random
import sys
import time
from dataclasses import dataclass, field

from pydantic_settings import BaseSettings, SettingsConfigDict

import staso as st


# ---------------------------------------------------------------------------
# Configuration (loaded from .env file or environment variables)
# ---------------------------------------------------------------------------

class TestAgentSettings(BaseSettings):
    model_config = SettingsConfigDict(
        env_file=".env",
        env_file_encoding="utf-8",
        extra="ignore",
    )

    staso_api_key: str
    staso_base_url: str = "https://api.staso.ai"
    staso_environment: str = "dev"
    staso_workspace_slug: str = "staging"
    provider: str = "anthropic"  # "anthropic" or "openai"
    anthropic_api_key: str = ""
    openai_api_key: str = ""


_settings = TestAgentSettings()

PROVIDER = _settings.provider.lower()
ENVIRONMENT = _settings.staso_environment
WORKSPACE_SLUG = _settings.staso_workspace_slug
BASE_URL = _settings.staso_base_url
API_KEY = _settings.staso_api_key

if PROVIDER == "anthropic" and not _settings.anthropic_api_key:
    sys.exit("ANTHROPIC_API_KEY is required when PROVIDER=anthropic")
if PROVIDER == "openai" and not _settings.openai_api_key:
    sys.exit("OPENAI_API_KEY is required when PROVIDER=openai")

# Set provider-specific env vars
if _settings.anthropic_api_key:
    os.environ["ANTHROPIC_API_KEY"] = _settings.anthropic_api_key
if _settings.openai_api_key:
    os.environ["OPENAI_API_KEY"] = _settings.openai_api_key

# ---------------------------------------------------------------------------
# 1. Initialise Staso + patch the active provider
# ---------------------------------------------------------------------------

st.init(
    api_key=API_KEY,
    agent_id="test-support-agent",
    base_url=BASE_URL,
    environment=ENVIRONMENT,
    workspace_slug=WORKSPACE_SLUG,
    debug=True,
)


# ---------------------------------------------------------------------------
# Provider abstraction — handles tool format, LLM calls, response parsing
# ---------------------------------------------------------------------------

@dataclass
class ToolCall:
    """Provider-agnostic tool call."""
    id: str
    name: str
    arguments: dict


@dataclass
class LLMResponse:
    """Provider-agnostic LLM response."""
    text: str | None = None
    tool_calls: list[ToolCall] = field(default_factory=list)
    raw: object = None  # original response for message history


def _default_model() -> str:
    if PROVIDER == "openai":
        return "gpt-4o-mini"
    return "claude-haiku-4-5"


if PROVIDER == "openai":
    from openai import OpenAI
    st.integrations.patch_openai()
    _llm_client = OpenAI()
else:
    from anthropic import Anthropic
    st.integrations.patch_anthropic()
    _llm_client = Anthropic()


# ---------------------------------------------------------------------------
# 2. Tools — covering many different categories for filter diversity
# ---------------------------------------------------------------------------


@st.tool(name="search_knowledge_base", tags=["retrieval", "knowledge"])
def search_knowledge_base(query: str) -> str:
    """Search the knowledge base for articles matching the query."""
    time.sleep(random.uniform(0.05, 0.2))  # Simulate I/O latency
    kb = {
        "billing": "Billing FAQ: You can update your payment method in Settings > Billing. We accept Visa, Mastercard, and bank transfers. Changes take effect on your next billing cycle.",
        "refund": "Refund policy: Refunds within 30 days are processed automatically to the original payment method. Enterprise refunds require manager approval and take 5-7 business days.",
        "password": "Password reset: Click 'Forgot password' on the login page. A reset link valid for 1 hour will be sent to your registered email. Ensure your email is up to date in Settings.",
        "pricing": "Pricing tiers: Free ($0, 1K requests/mo), Pro ($49, 100K requests/mo), Enterprise (custom). All plans include basic analytics. Pro adds priority support.",
        "api": "API documentation: Base URL is https://api.example.com/v2. Authentication via Bearer token. Rate limit: 1000 req/min (Pro), 100 req/min (Free). See /docs for OpenAPI spec.",
        "sso": "SSO setup: Navigate to Settings > Security > SSO. We support SAML 2.0 and OIDC. Upload your IdP metadata XML or enter the discovery URL. Test before enforcing.",
        "data": "Data export: Go to Settings > Data > Export. Choose format (CSV, JSON, Parquet). Exports over 1GB are queued and delivered via email link within 24 hours.",
    }
    for key, answer in kb.items():
        if key in query.lower():
            return answer
    return f"No knowledge base articles found matching: '{query}'. Try broader search terms or contact support for personalized help."


@st.tool(name="create_ticket", tags=["ticketing", "workflow"])
def create_ticket(title: str, description: str, priority: str = "medium") -> dict:
    """Create a support ticket for issues that need follow-up."""
    time.sleep(random.uniform(0.1, 0.3))  # Simulate DB write
    ticket_id = f"TKT-{random.randint(1000, 9999)}"
    return {
        "ticket_id": ticket_id,
        "title": title,
        "description": description,
        "priority": priority,
        "status": "open",
        "assigned_to": "support-team",
        "sla_hours": {"low": 72, "medium": 24, "high": 4}.get(priority, 24),
        "created_at": "2026-03-20T12:00:00Z",
    }


@st.tool(name="get_account_info", tags=["account", "lookup"])
def get_account_info(account_id: str) -> dict:
    """Look up detailed account information."""
    time.sleep(random.uniform(0.02, 0.1))  # Simulate DB read
    accounts = {
        "ACC-789": {
            "account_id": "ACC-789",
            "name": "Jane Doe",
            "email": "jane.doe@example.com",
            "plan": "Pro",
            "status": "active",
            "balance": 42.50,
            "monthly_usage": {"api_calls": 45_230, "storage_gb": 12.4},
            "created_at": "2024-06-15",
        },
        "ACC-456": {
            "account_id": "ACC-456",
            "name": "Acme Corp",
            "email": "admin@acme.com",
            "plan": "Enterprise",
            "status": "active",
            "balance": 0.00,
            "monthly_usage": {"api_calls": 1_234_567, "storage_gb": 450.2},
            "created_at": "2023-01-10",
        },
        "ACC-SUSPENDED": {
            "account_id": "ACC-SUSPENDED",
            "name": "Suspended User",
            "email": "suspended@example.com",
            "plan": "Free",
            "status": "suspended",
            "balance": -15.00,
            "suspension_reason": "Payment overdue for 60+ days",
        },
    }
    if account_id not in accounts:
        raise ValueError(f"Account not found: {account_id}")
    return accounts[account_id]


@st.tool(name="send_email", tags=["communication", "notification"])
def send_email(to: str, subject: str, body: str) -> dict:
    """Send an email notification to a customer."""
    time.sleep(random.uniform(0.1, 0.4))  # Simulate email API
    return {
        "message_id": f"msg-{random.randint(10000, 99999)}",
        "to": to,
        "subject": subject,
        "status": "sent",
        "delivered_at": "2026-03-20T12:00:05Z",
    }


@st.tool(name="check_service_status", tags=["monitoring", "health"])
def check_service_status(service_name: str) -> dict:
    """Check the operational status of a platform service."""
    time.sleep(random.uniform(0.01, 0.05))
    services = {
        "api": {"status": "operational", "latency_p99_ms": 45, "uptime_30d": "99.98%"},
        "dashboard": {"status": "operational", "latency_p99_ms": 120, "uptime_30d": "99.95%"},
        "billing": {"status": "degraded", "latency_p99_ms": 890, "uptime_30d": "99.80%", "incident": "INC-2026-0312: Slow payment processing"},
        "auth": {"status": "operational", "latency_p99_ms": 22, "uptime_30d": "99.99%"},
    }
    return services.get(service_name, {"status": "unknown", "error": f"Service '{service_name}' not found"})


@st.tool(name="run_diagnostic", tags=["debugging", "technical"])
def run_diagnostic(account_id: str, check_type: str = "full") -> dict:
    """Run diagnostic checks on an account's configuration."""
    time.sleep(random.uniform(0.3, 0.8))  # Simulate diagnostic scan
    return {
        "account_id": account_id,
        "check_type": check_type,
        "results": [
            {"check": "api_key_valid", "status": "pass"},
            {"check": "webhook_reachable", "status": "pass"},
            {"check": "rate_limit_config", "status": "warning", "detail": "Using default rate limits (1000/min)"},
            {"check": "ssl_certificate", "status": "pass", "expires_in_days": 245},
            {"check": "dns_resolution", "status": "pass", "latency_ms": 12},
        ],
        "overall": "healthy",
        "recommendations": ["Consider customizing rate limits for production workload"],
    }


@st.tool(name="search_logs", tags=["debugging", "observability"])
def search_logs(account_id: str, query: str, minutes: int = 60) -> dict:
    """Search application logs for an account."""
    time.sleep(random.uniform(0.2, 0.5))
    return {
        "account_id": account_id,
        "query": query,
        "time_range_minutes": minutes,
        "total_results": 47,
        "sample_entries": [
            {"ts": "2026-03-20T11:58:32Z", "level": "ERROR", "msg": f"Rate limit exceeded for {account_id}"},
            {"ts": "2026-03-20T11:55:10Z", "level": "WARN", "msg": f"Slow query detected: {query} (1200ms)"},
            {"ts": "2026-03-20T11:50:01Z", "level": "INFO", "msg": f"API key rotated for {account_id}"},
        ],
    }


# ---------------------------------------------------------------------------
# 3. Chain / retriever decorators — for span kind diversity
# ---------------------------------------------------------------------------


@st.trace(name="retrieve_customer_context", kind="RETRIEVER", tags=["context"])
def retrieve_customer_context(account_id: str) -> dict:
    """Retrieve full customer context for the agent (retriever span)."""
    account = get_account_info(account_id)
    diagnostics = run_diagnostic(account_id, check_type="quick")
    return {
        "account": account,
        "diagnostics": diagnostics,
        "context_summary": f"Customer {account['name']} on {account['plan']} plan, status: {account['status']}",
    }


@st.trace(name="classify_intent", kind="CHAIN", tags=["nlp", "routing"])
def classify_intent(user_message: str) -> str:
    """Classify user intent to route to the right workflow (chain span)."""
    time.sleep(random.uniform(0.01, 0.05))
    keywords = {
        "billing": ["billing", "payment", "charge", "invoice", "price"],
        "technical": ["error", "bug", "broken", "not working", "api", "debug", "log"],
        "account": ["account", "plan", "upgrade", "cancel", "suspend"],
        "general": ["help", "how", "what", "when", "password", "reset"],
    }
    msg_lower = user_message.lower()
    for intent, words in keywords.items():
        if any(w in msg_lower for w in words):
            return intent
    return "general"


@st.trace(name="format_response", kind="CHAIN", tags=["formatting"])
def format_response(raw_response: str, style: str = "friendly") -> str:
    """Post-process and format the agent response (chain span)."""
    time.sleep(random.uniform(0.01, 0.02))
    # Simulated formatting
    if style == "concise":
        return raw_response[:500] if len(raw_response) > 500 else raw_response
    return raw_response


# ---------------------------------------------------------------------------
# 4. Tool definitions — canonical form, converted per provider
# ---------------------------------------------------------------------------

# Canonical tool definitions (provider-agnostic)
_TOOL_DEFS = [
    {
        "name": "search_knowledge_base",
        "description": "Search the knowledge base for relevant articles, FAQs, and documentation.",
        "parameters": {
            "type": "object",
            "properties": {
                "query": {"type": "string", "description": "Natural language search query"},
            },
            "required": ["query"],
        },
    },
    {
        "name": "create_ticket",
        "description": "Create a support ticket for issues requiring follow-up or escalation.",
        "parameters": {
            "type": "object",
            "properties": {
                "title": {"type": "string", "description": "Short descriptive ticket title"},
                "description": {"type": "string", "description": "Detailed description of the issue"},
                "priority": {
                    "type": "string",
                    "enum": ["low", "medium", "high"],
                    "description": "Ticket priority level",
                },
            },
            "required": ["title", "description"],
        },
    },
    {
        "name": "get_account_info",
        "description": "Look up detailed account information including plan, usage, and balance.",
        "parameters": {
            "type": "object",
            "properties": {
                "account_id": {"type": "string", "description": "The account ID (e.g. ACC-789)"},
            },
            "required": ["account_id"],
        },
    },
    {
        "name": "send_email",
        "description": "Send an email notification to a customer about their issue or resolution.",
        "parameters": {
            "type": "object",
            "properties": {
                "to": {"type": "string", "description": "Recipient email address"},
                "subject": {"type": "string", "description": "Email subject line"},
                "body": {"type": "string", "description": "Email body content"},
            },
            "required": ["to", "subject", "body"],
        },
    },
    {
        "name": "check_service_status",
        "description": "Check the operational status of a platform service (api, dashboard, billing, auth).",
        "parameters": {
            "type": "object",
            "properties": {
                "service_name": {"type": "string", "description": "Service to check (api, dashboard, billing, auth)"},
            },
            "required": ["service_name"],
        },
    },
    {
        "name": "run_diagnostic",
        "description": "Run diagnostic checks on an account to identify configuration issues.",
        "parameters": {
            "type": "object",
            "properties": {
                "account_id": {"type": "string", "description": "Account ID to diagnose"},
                "check_type": {
                    "type": "string",
                    "enum": ["quick", "full"],
                    "description": "Diagnostic depth (quick ~2s, full ~10s)",
                },
            },
            "required": ["account_id"],
        },
    },
    {
        "name": "search_logs",
        "description": "Search application logs for a specific account to debug issues.",
        "parameters": {
            "type": "object",
            "properties": {
                "account_id": {"type": "string", "description": "Account to search logs for"},
                "query": {"type": "string", "description": "Log search query"},
                "minutes": {"type": "integer", "description": "How far back to search (default: 60)"},
            },
            "required": ["account_id", "query"],
        },
    },
]


def _tools_for_provider() -> list[dict]:
    """Convert canonical tool defs to the active provider's format."""
    if PROVIDER == "openai":
        return [
            {"type": "function", "function": {"name": t["name"], "description": t["description"], "parameters": t["parameters"]}}
            for t in _TOOL_DEFS
        ]
    # Anthropic uses input_schema instead of parameters
    return [
        {"name": t["name"], "description": t["description"], "input_schema": t["parameters"]}
        for t in _TOOL_DEFS
    ]


TOOLS = _tools_for_provider()

TOOL_MAP: dict[str, callable] = {
    "search_knowledge_base": search_knowledge_base,
    "create_ticket": create_ticket,
    "get_account_info": get_account_info,
    "send_email": send_email,
    "check_service_status": check_service_status,
    "run_diagnostic": run_diagnostic,
    "search_logs": search_logs,
}


# ---------------------------------------------------------------------------
# 5. Provider-agnostic LLM call + tool dispatch
# ---------------------------------------------------------------------------


def _llm_chat(
    messages: list[dict],
    system: str,
    model: str | None = None,
    max_tokens: int = 1024,
) -> LLMResponse:
    """Call the active LLM provider and return a normalized response."""
    model = model or _default_model()

    if PROVIDER == "openai":
        oai_messages = [{"role": "system", "content": system}] + messages
        response = _llm_client.chat.completions.create(
            model=model,
            max_tokens=max_tokens,
            messages=oai_messages,
            tools=TOOLS,
        )
        choice = response.choices[0]
        text = choice.message.content
        tool_calls = []
        if choice.message.tool_calls:
            for tc in choice.message.tool_calls:
                tool_calls.append(ToolCall(
                    id=tc.id,
                    name=tc.function.name,
                    arguments=json.loads(tc.function.arguments),
                ))
        return LLMResponse(text=text, tool_calls=tool_calls, raw=response)
    else:
        response = _llm_client.messages.create(
            model=model,
            max_tokens=max_tokens,
            system=system,
            tools=TOOLS,
            messages=messages,
        )
        text = None
        tool_calls = []
        for block in response.content:
            if hasattr(block, "text"):
                text = (text or "") + block.text
            if getattr(block, "type", None) == "tool_use":
                tool_calls.append(ToolCall(
                    id=block.id,
                    name=block.name,
                    arguments=block.input,
                ))
        return LLMResponse(text=text, tool_calls=tool_calls, raw=response)


def _execute_tools(tool_calls: list[ToolCall]) -> list[dict]:
    """Execute tool calls and return results in provider-specific format."""
    results = []
    for tc in tool_calls:
        tool_fn = TOOL_MAP.get(tc.name)
        content = f"Error: unknown tool '{tc.name}'"
        is_error = True
        if tool_fn:
            try:
                result = tool_fn(**tc.arguments)
                content = json.dumps(result) if isinstance(result, dict) else str(result)
                is_error = False
            except Exception as e:
                content = f"Error: {e}"

        if PROVIDER == "openai":
            results.append({"role": "tool", "tool_call_id": tc.id, "content": content})
        else:
            entry = {"type": "tool_result", "tool_use_id": tc.id, "content": content}
            if is_error:
                entry["is_error"] = True
            results.append(entry)
    return results


def _append_assistant_and_tools(messages: list[dict], response: LLMResponse, tool_results: list[dict]) -> None:
    """Append assistant response + tool results to message history (provider-specific format)."""
    if PROVIDER == "openai":
        raw_choice = response.raw.choices[0]
        messages.append({"role": "assistant", "content": raw_choice.message.content, "tool_calls": [
            {"id": tc.id, "type": "function", "function": {"name": tc.function.name, "arguments": tc.function.arguments}}
            for tc in raw_choice.message.tool_calls
        ]})
        messages.extend(tool_results)
    else:
        messages.append({"role": "assistant", "content": response.raw.content})
        messages.append({"role": "user", "content": tool_results})


# ---------------------------------------------------------------------------
# 6. Agent implementations
# ---------------------------------------------------------------------------


@st.agent(name="customer-support-agent", tags=["support", "tier-1"])
def run_support_agent(user_message: str, model: str | None = None) -> str:
    """Tier-1 customer support agent using tool-use loop."""
    messages: list[dict] = [{"role": "user", "content": user_message}]
    system = (
        "You are a helpful customer support agent for a SaaS platform. "
        "Use the available tools to look up information, run diagnostics, "
        "and create tickets when needed. Be concise and friendly. "
        "Always check the knowledge base first before escalating."
    )

    for _turn in range(10):
        resp = _llm_chat(messages, system=system, model=model)

        if not resp.tool_calls:
            return format_response(resp.text or "", style="friendly")

        tool_results = _execute_tools(resp.tool_calls)
        _append_assistant_and_tools(messages, resp, tool_results)

    return "Agent reached maximum turns without completing."


@st.agent(name="escalation-agent", tags=["support", "tier-2", "escalation"])
def run_escalation_agent(
    user_message: str,
    account_id: str,
    original_ticket_id: str | None = None,
) -> str:
    """Tier-2 escalation agent with deeper diagnostic capabilities."""
    context = retrieve_customer_context(account_id)
    intent = classify_intent(user_message)

    context_str = (
        f"Customer: {context['account']['name']} ({context['account']['plan']} plan)\n"
        f"Status: {context['account']['status']}\n"
        f"Diagnostics: {context['diagnostics']['overall']}\n"
        f"Intent: {intent}\n"
    )
    if original_ticket_id:
        context_str += f"Escalated from ticket: {original_ticket_id}\n"

    messages: list[dict] = [{"role": "user", "content": user_message}]
    system = (
        f"You are a senior technical support engineer handling an escalated case.\n\n"
        f"Customer context:\n{context_str}\n\n"
        "You have access to advanced diagnostic tools. Investigate thoroughly, "
        "check logs, run diagnostics, and provide a detailed resolution. "
        "If the issue requires engineering intervention, create a high-priority ticket. "
        "Always send the customer a follow-up email with your findings."
    )

    for _turn in range(10):
        resp = _llm_chat(messages, system=system, max_tokens=2048)

        if not resp.tool_calls:
            return format_response(resp.text or "", style="concise")

        tool_results = _execute_tools(resp.tool_calls)
        _append_assistant_and_tools(messages, resp, tool_results)

    return "Escalation agent reached maximum turns."


@st.agent(name="triage-agent", tags=["triage", "routing"])
def run_triage_agent(user_message: str, account_id: str | None = None) -> str:
    """Triage agent that classifies and routes to the right tier.

    Demonstrates:
      - Chain span (intent classification)
      - Conditional routing to different sub-agents
      - Deep span nesting (triage → escalation → retriever → tools)
    """
    intent = classify_intent(user_message)

    if intent == "technical" and account_id:
        # Route to escalation agent for technical issues
        return run_escalation_agent(
            user_message=user_message,
            account_id=account_id,
        )
    else:
        # Route to standard support
        return run_support_agent(user_message)


# ---------------------------------------------------------------------------
# 7. Error scenario agent — exercises error status in the UI
# ---------------------------------------------------------------------------


@st.agent(name="flaky-agent", tags=["test", "error-scenarios"])
def run_flaky_agent(user_message: str) -> str:
    """Agent that intentionally encounters errors for testing error display."""
    try:
        get_account_info("ACC-NONEXISTENT")
    except Exception:
        pass  # Swallow, but the tool span will show error status

    check_service_status("billing")

    resp = _llm_chat(
        messages=[{"role": "user", "content": user_message}],
        system="You are a helpful assistant.",
        max_tokens=256,
    )
    return resp.text or ""


# ---------------------------------------------------------------------------
# 8. Run all scenarios
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    print("=" * 70)
    print("Staso SDK — Comprehensive UI Test Suite")
    print(f"  Provider:    {PROVIDER}")
    print(f"  Model:       {_default_model()}")
    print(f"  Environment: {ENVIRONMENT}")
    print(f"  Workspace:   {WORKSPACE_SLUG}")
    print(f"  API Key:     {API_KEY[:12]}...")
    print(f"  Backend:     {BASE_URL}")
    print("=" * 70)

    # ── Scenario 1: Multi-tool billing inquiry (tier-1) ────────────────
    # Exercises: session, agent span, LLM spans, tool spans,
    #            knowledge base + account lookup, input/output preview
    print("\n━━━ Scenario 1: Multi-tool billing inquiry ━━━")
    with st.conversation("customer-jane-doe-001"):
        result = run_support_agent(
            "I need help with billing. My account ID is ACC-789. "
            "Can you look up my account and find the billing FAQ? "
            "I want to change my payment method.",
        )
        print(f"  → {result[:200]}...")

    # ── Scenario 2: Simple FAQ query (tier-1, different session) ───────
    # Exercises: separate session, minimal tool usage, fast trace
    print("\n━━━ Scenario 2: Password reset FAQ ━━━")
    with st.conversation("customer-quick-002"):
        result = run_support_agent("How do I reset my password?")
        print(f"  → {result[:200]}...")

    # ── Scenario 3: Technical escalation (tier-2, sonnet model) ────────
    # Exercises: retriever span, chain spans, deep nesting,
    #            different model (cost comparison), diagnostic tools,
    #            email sending, log searching
    print("\n━━━ Scenario 3: Technical escalation (tier-2) ━━━")
    with st.conversation("enterprise-acme-003"):
        result = run_escalation_agent(
            user_message=(
                "Our API integration is returning 429 errors intermittently. "
                "We're seeing this across multiple endpoints since yesterday. "
                "Our account ID is ACC-456. This is blocking our production deployment."
            ),
            account_id="ACC-456",
            original_ticket_id="TKT-5678",
        )
        print(f"  → {result[:200]}...")

    # ── Scenario 4: Triage → auto-route to escalation ─────────────────
    # Exercises: triage agent → nested escalation agent,
    #            very deep span tree, agent-within-agent
    print("\n━━━ Scenario 4: Triage with auto-routing ━━━")
    with st.conversation("customer-debug-004"):
        result = run_triage_agent(
            user_message=(
                "I'm getting error 500 on the API when trying to upload files. "
                "The logs show a timeout but I can't figure out why."
            ),
            account_id="ACC-789",
        )
        print(f"  → {result[:200]}...")

    # ── Scenario 5: Triage → standard support ─────────────────────────
    # Exercises: triage routing to tier-1 instead
    print("\n━━━ Scenario 5: Triage routing to standard support ━━━")
    with st.conversation("customer-pricing-005"):
        result = run_triage_agent(
            user_message="What are your pricing plans? I'm considering upgrading.",
        )
        print(f"  → {result[:200]}...")

    # ── Scenario 6: Error scenarios ────────────────────────────────────
    # Exercises: error status spans, mixed OK/error within trace
    print("\n━━━ Scenario 6: Error scenario test ━━━")
    with st.conversation("test-errors-006"):
        result = run_flaky_agent(
            "I need help but my account seems to have issues."
        )
        print(f"  → {result[:200]}...")

    # ── Scenario 7: Multi-turn conversation in same session ────────────
    # Exercises: multiple traces in same session (session grouping)
    print("\n━━━ Scenario 7: Multi-turn conversation ━━━")
    with st.conversation("conversation-multi-007"):
        result1 = run_support_agent(
            "What's the status of the billing service?"
        )
        print(f"  Turn 1 → {result1[:150]}...")

        result2 = run_support_agent(
            "OK, can you create a ticket about the billing slowness? "
            "Make it high priority."
        )
        print(f"  Turn 2 → {result2[:150]}...")

        result3 = run_support_agent(
            "Thanks! Can you also check if the API service is healthy?"
        )
        print(f"  Turn 3 → {result3[:150]}...")

    # ── Scenario 8: Service status check (minimal, fast) ──────────────
    # Exercises: very short trace for waterfall contrast
    print("\n━━━ Scenario 8: Quick service health check ━━━")
    with st.conversation("health-check-008"):
        result = run_support_agent(
            "Is the dashboard service working right now?"
        )
        print(f"  → {result[:200]}...")

    # ── Scenario 9: SSO and data export queries ───────────────────────
    # Exercises: different knowledge base hits for search diversity
    print("\n━━━ Scenario 9: SSO setup help ━━━")
    with st.conversation("enterprise-sso-009"):
        result = run_support_agent(
            "We need to set up SSO with our Okta IdP. "
            "Can you walk me through the process and also tell me "
            "how to export our data in JSON format?"
        )
        print(f"  → {result[:200]}...")

    # ── Scenario 10: Suspended account (edge case) ────────────────────
    # Exercises: account with negative balance, suspended status
    print("\n━━━ Scenario 10: Suspended account inquiry ━━━")
    with st.conversation("suspended-user-010"):
        result = run_escalation_agent(
            user_message=(
                "My account seems to be suspended and I can't access anything. "
                "I thought my payment went through last week. Please help!"
            ),
            account_id="ACC-SUSPENDED",
        )
        print(f"  → {result[:200]}...")

    # ── Flush and exit ────────────────────────────────────────────────
    st.shutdown()
    print("\n" + "=" * 70)
    print("All scenarios complete. Trace events flushed to backend.")
    print(f"View traces at: {BASE_URL.replace(':6999', ':3000')}/observability/traces")
    print("=" * 70)
