"""
Load-test v2 — fresh agent names, Rastir @tool instrumented tools,
mix of short and complex multi-step prompts.

Uses three LangGraph ReAct agents:
  - travel_concierge     (OpenAI gpt-4o-mini)
  - fund_strategist      (Anthropic claude-3-haiku)
  - site_reliability     (Bedrock claude-3-haiku + Guardrails)

All LangChain tools are wrapped with Rastir @tool so they appear
as tool spans in traces.

Usage:
    cd /home/skamalj/dev/llmobserve
    PYTHONPATH=src conda run -n llmobserve python tests/integration/test_load_v2.py
"""

from __future__ import annotations

import json
import os
import random
import time
import traceback
from datetime import datetime

# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# 1. Configure Rastir
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

import rastir

rastir.configure(
    service="agent-load-v2",
    env="integration",
    version="0.3.0",
    push_url="http://localhost:8080",
    batch_size=50,
)

# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# 2. Imports
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

from langchain_openai import ChatOpenAI
from langchain_anthropic import ChatAnthropic
from langchain_aws import ChatBedrock
from langchain_core.messages import HumanMessage
from langchain_core.tools import tool
from langgraph.graph import StateGraph, MessagesState, START, END
from langgraph.prebuilt import ToolNode

from rastir import agent, llm, trace as rastir_trace

# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# 3. Tools — wrapped with Rastir @trace for tracing
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

# --- Travel Planner tools ---

@tool
@rastir_trace(name="flight_search")
def flight_search(origin: str, destination: str) -> str:
    """Search for available flights between two airports."""
    time.sleep(0.2)
    flights = [
        {"airline": "Air India", "price": random.randint(200, 800),
         "departure": "08:30", "arrival": "14:45", "stops": 0},
        {"airline": "Emirates", "price": random.randint(300, 1200),
         "departure": "22:10", "arrival": "06:30+1", "stops": 1},
        {"airline": "Lufthansa", "price": random.randint(250, 900),
         "departure": "11:00", "arrival": "18:20", "stops": 1},
    ]
    return json.dumps({"flights": flights, "route": f"{origin} → {destination}"})


@tool
@rastir_trace(name="hotel_search")
def hotel_search(city: str, nights: int) -> str:
    """Search for hotels in a city for a number of nights."""
    time.sleep(0.15)
    hotels = [
        {"name": "Grand Hyatt", "stars": 5, "price_per_night": random.randint(150, 350),
         "rating": round(random.uniform(4.0, 5.0), 1)},
        {"name": "Holiday Inn Express", "stars": 3, "price_per_night": random.randint(60, 120),
         "rating": round(random.uniform(3.5, 4.5), 1)},
        {"name": "Marriott Courtyard", "stars": 4, "price_per_night": random.randint(90, 200),
         "rating": round(random.uniform(3.8, 4.7), 1)},
    ]
    for h in hotels:
        h["total"] = h["price_per_night"] * nights
    return json.dumps({"city": city, "nights": nights, "hotels": hotels})


@tool
@rastir_trace(name="currency_converter")
def currency_converter(amount: float, from_currency: str, to_currency: str) -> str:
    """Convert an amount between currencies."""
    rates = {
        ("USD", "EUR"): 0.92, ("USD", "GBP"): 0.79, ("USD", "JPY"): 149.5,
        ("USD", "INR"): 83.2, ("EUR", "USD"): 1.09, ("GBP", "USD"): 1.27,
        ("INR", "USD"): 0.012, ("JPY", "USD"): 0.0067,
    }
    key = (from_currency.upper(), to_currency.upper())
    rate = rates.get(key, 1.0)
    converted = round(amount * rate, 2)
    return f"{amount} {from_currency} = {converted} {to_currency} (rate: {rate})"


@tool
@rastir_trace(name="weather_forecast")
def weather_forecast(city: str, days: int) -> str:
    """Get weather forecast for a city for N days."""
    time.sleep(0.1)
    forecasts = []
    conditions = ["Sunny", "Partly Cloudy", "Overcast", "Rain", "Thunderstorm", "Clear"]
    for d in range(days):
        forecasts.append({
            "day": d + 1,
            "condition": random.choice(conditions),
            "high_c": random.randint(15, 38),
            "low_c": random.randint(5, 22),
            "humidity": random.randint(30, 90),
        })
    return json.dumps({"city": city, "forecast": forecasts})


# --- Portfolio Manager tools ---

@tool
@rastir_trace(name="stock_lookup")
def stock_lookup(ticker: str) -> str:
    """Look up current stock data including price, market cap, and P/E ratio."""
    time.sleep(0.1)
    stocks = {
        "AAPL": {"price": 187.42, "market_cap": "2.9T", "pe": 29.3, "name": "Apple Inc."},
        "GOOGL": {"price": 141.80, "market_cap": "1.8T", "pe": 23.1, "name": "Alphabet Inc."},
        "MSFT": {"price": 415.33, "market_cap": "3.1T", "pe": 35.7, "name": "Microsoft Corp."},
        "AMZN": {"price": 178.25, "market_cap": "1.9T", "pe": 58.2, "name": "Amazon.com Inc."},
        "TSLA": {"price": 248.90, "market_cap": "792B", "pe": 62.1, "name": "Tesla Inc."},
        "NVDA": {"price": 875.60, "market_cap": "2.2T", "pe": 65.4, "name": "NVIDIA Corp."},
        "META": {"price": 485.20, "market_cap": "1.2T", "pe": 26.8, "name": "Meta Platforms"},
        "NFLX": {"price": 628.10, "market_cap": "271B", "pe": 44.3, "name": "Netflix Inc."},
    }
    data = stocks.get(ticker.upper())
    if data:
        data["ticker"] = ticker.upper()
        data["change_pct"] = round(random.uniform(-4.0, 4.0), 2)
        return json.dumps(data)
    return json.dumps({"error": f"Ticker '{ticker}' not found"})


@tool
@rastir_trace(name="portfolio_calculator")
def portfolio_calculator(expression: str) -> str:
    """Calculate a financial expression (e.g., portfolio value, returns)."""
    try:
        result = eval(expression, {"__builtins__": {}}, {})
        return f"Result: {result}"
    except Exception as e:
        return f"Calculation error: {e}"


@tool
@rastir_trace(name="risk_assessment")
def risk_assessment(tickers: str) -> str:
    """Assess risk profile for a comma-separated list of stock tickers."""
    time.sleep(0.2)
    ticker_list = [t.strip().upper() for t in tickers.split(",")]
    volatility = {t: round(random.uniform(0.15, 0.65), 3) for t in ticker_list}
    beta = {t: round(random.uniform(0.5, 2.0), 2) for t in ticker_list}
    avg_vol = sum(volatility.values()) / len(volatility)
    risk_level = "LOW" if avg_vol < 0.25 else "MEDIUM" if avg_vol < 0.45 else "HIGH"
    return json.dumps({
        "tickers": ticker_list,
        "volatility": volatility,
        "beta": beta,
        "portfolio_risk": risk_level,
        "sharpe_estimate": round(random.uniform(0.5, 2.5), 2),
    })


@tool
@rastir_trace(name="market_news")
def market_news(topic: str) -> str:
    """Get latest market news about a topic or sector."""
    time.sleep(0.15)
    headlines = [
        f"{topic} sector sees strong Q4 earnings beat across major players",
        f"Analysts upgrade {topic} outlook citing AI-driven growth momentum",
        f"Institutional investors increase {topic} allocation by 15% in 2026",
        f"Fed policy impact on {topic}: What investors need to know",
    ]
    return json.dumps({"topic": topic, "headlines": random.sample(headlines, 3)})


# --- Ops Engineer tools ---

@tool
@rastir_trace(name="run_query")
def run_query(sql: str) -> str:
    """Execute a SQL query against the operations database."""
    time.sleep(0.1)
    if "error" in sql.lower() or "incident" in sql.lower():
        return json.dumps({
            "rows": [
                {"incident_id": "INC-4421", "severity": "P1", "service": "payment-api",
                 "duration_min": 45, "root_cause": "DB connection pool exhaustion"},
                {"incident_id": "INC-4398", "severity": "P2", "service": "auth-service",
                 "duration_min": 12, "root_cause": "Certificate expiry"},
                {"incident_id": "INC-4375", "severity": "P3", "service": "notification-svc",
                 "duration_min": 8, "root_cause": "Rate limit exceeded"},
            ],
            "row_count": 3,
        })
    return json.dumps({
        "rows": [
            {"service": "payment-api", "cpu_avg": 42.3, "memory_pct": 68,
             "requests_per_sec": 1250, "error_rate": 0.02},
            {"service": "auth-service", "cpu_avg": 28.1, "memory_pct": 45,
             "requests_per_sec": 3800, "error_rate": 0.001},
            {"service": "order-service", "cpu_avg": 55.8, "memory_pct": 72,
             "requests_per_sec": 890, "error_rate": 0.05},
        ],
        "row_count": 3,
    })


@tool
@rastir_trace(name="metric_calculator")
def metric_calculator(expression: str) -> str:
    """Calculate operational metrics (SLA, percentages, rates, etc.)."""
    try:
        result = eval(expression, {"__builtins__": {}}, {})
        return f"Result: {result}"
    except Exception as e:
        return f"Error: {e}"


@tool
@rastir_trace(name="log_search")
def log_search(service: str, level: str, minutes: int) -> str:
    """Search application logs for a service at a given level in last N minutes."""
    time.sleep(0.15)
    entries = []
    levels = {"ERROR": 3, "WARN": 8, "INFO": 25}
    count = levels.get(level.upper(), 5)
    for i in range(min(count, 10)):
        entries.append({
            "timestamp": f"2026-02-28T16:{random.randint(0,59):02d}:{random.randint(0,59):02d}Z",
            "level": level.upper(),
            "message": random.choice([
                f"Connection timeout to downstream {service}",
                f"Retry attempt {random.randint(1,5)} for request",
                "Health check failed for secondary replica",
                f"Latency spike detected: {random.randint(200,5000)}ms",
                f"Circuit breaker opened for {service}",
            ]),
        })
    return json.dumps({"service": service, "level": level, "entries": entries})


@tool
@rastir_trace(name="deployment_status")
def deployment_status(service: str) -> str:
    """Get current deployment status for a service."""
    time.sleep(0.05)
    return json.dumps({
        "service": service,
        "version": f"v2.{random.randint(1,15)}.{random.randint(0,9)}",
        "status": random.choice(["RUNNING", "DEPLOYING", "ROLLBACK"]),
        "replicas": {"desired": 3, "ready": random.randint(1, 3)},
        "last_deploy": "2026-02-28T14:22:00Z",
        "health": random.choice(["HEALTHY", "DEGRADED"]),
    })


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# 4. Agent builders
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

def _build_langgraph_agent(model, tools):
    """Generic LangGraph ReAct agent builder."""
    bound_model = model.bind_tools(tools)

    def chatbot(state: MessagesState):
        return {"messages": [bound_model.invoke(state["messages"])]}

    def should_continue(state: MessagesState):
        last = state["messages"][-1]
        if hasattr(last, "tool_calls") and last.tool_calls:
            return "tools"
        return END

    graph = StateGraph(MessagesState)
    graph.add_node("chatbot", chatbot)
    graph.add_node("tools", ToolNode(tools))
    graph.add_edge(START, "chatbot")
    graph.add_conditional_edges("chatbot", should_continue, {"tools": "tools", END: END})
    graph.add_edge("tools", "chatbot")
    return graph.compile()


def build_travel_agent():
    """OpenAI travel planner with flight, hotel, weather, currency tools."""
    model = ChatOpenAI(
        model="gpt-4o-mini",
        temperature=0,
        api_key=os.environ.get("API_OPENAI_KEY", os.environ.get("OPENAI_API_KEY", "")),
    )
    return _build_langgraph_agent(
        model, [flight_search, hotel_search, weather_forecast, currency_converter]
    )


def build_portfolio_agent():
    """Anthropic portfolio manager with stocks, risk, news, calculator."""
    model = ChatAnthropic(
        model="claude-3-haiku-20240307",
        temperature=0,
        api_key=os.environ.get("API_ANTHROPIC_KEY", os.environ.get("ANTHROPIC_API_KEY", "")),
    )
    return _build_langgraph_agent(
        model, [stock_lookup, portfolio_calculator, risk_assessment, market_news]
    )


def build_ops_agent():
    """Bedrock ops engineer with DB, logs, deployment, calculator + guardrails."""
    model = ChatBedrock(
        model_id="anthropic.claude-3-haiku-20240307-v1:0",
        region_name="ap-south-1",
        model_kwargs={"temperature": 0, "max_tokens": 2048},
        guardrails={
            "guardrailIdentifier": "t844q29keb13",
            "guardrailVersion": "1",
            "trace": True,
        },
    )
    return _build_langgraph_agent(
        model, [run_query, metric_calculator, log_search, deployment_status]
    )


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# 5. Decorated agent runners
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

@agent(agent_name="trip_oracle")
def run_travel_agent(prompt: str):
    """OpenAI-backed travel planner."""
    app = build_travel_agent()
    return invoke_llm(app, prompt)


@agent(agent_name="wealth_scout")
def run_portfolio_agent(prompt: str):
    """Anthropic-backed portfolio manager."""
    app = build_portfolio_agent()
    return invoke_llm(app, prompt)


@agent(agent_name="platform_sentinel")
def run_ops_agent(prompt: str):
    """Bedrock-backed ops engineer with guardrails."""
    app = build_ops_agent()
    return invoke_llm(app, prompt)


@llm(evaluate=True, evaluation_types=["toxicity"])
def invoke_llm(app, prompt: str):
    """Invoke a LangGraph agent — @llm decorated for tracing."""
    result = app.invoke({"messages": [HumanMessage(prompt)]})
    return result


@rastir_trace(name="summarize_output")
def summarize_output(data: dict) -> str:
    """Post-process agent output into a summary."""
    time.sleep(0.05)
    msgs = data.get("messages", [])
    tool_msgs = [m for m in msgs if type(m).__name__ == "ToolMessage"]
    return f"Summary: {len(msgs)} messages, {len(tool_msgs)} tool calls"


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# 6. Prompts — mix of short and complex multi-tool scenarios
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

OPENAI_PROMPTS = [
    # Complex multi-tool: should trigger 3+ tools and take time
    "I'm planning a 5-day trip from New York to Tokyo. "
    "Search for flights from JFK to NRT, find hotels in Tokyo for 5 nights, "
    "get the weather forecast for Tokyo for 5 days, and convert 3000 USD to JPY. "
    "Give me a complete trip plan.",

    # 2-tool medium prompt
    "Find me flights from London to Paris and check the 3-day weather forecast for Paris. "
    "Which days look best for sightseeing?",

    # Simple no tools
    "What are the top 3 things to do in Barcelona? Answer directly without tools.",
]

ANTHROPIC_PROMPTS = [
    # Complex multi-tool: should trigger 4+ tools, takes time
    "I want to build a diversified tech portfolio. Look up AAPL, NVDA, META, and NFLX. "
    "Assess the risk of holding all four together. "
    "Also check the latest market news about semiconductors. "
    "Calculate the total cost of buying 50 shares of each at current prices. "
    "Give me a comprehensive investment analysis.",

    # 2-tool medium
    "Compare MSFT and GOOGL — look up both stocks and calculate which has the better P/E ratio. "
    "Which is the better value buy?",

    # Simple
    "What is dollar cost averaging? Explain in 2 sentences without tools.",
]

BEDROCK_PROMPTS = [
    # Complex multi-tool: DB query + logs + deployment + calculator
    "We had a production incident last night. Query the operations database for recent incidents. "
    "Then check the error logs for payment-api in the last 60 minutes. "
    "Also get the deployment status of payment-api. "
    "Calculate the mean time to resolution (MTTR) from the incidents. "
    "Give me a full incident post-mortem summary.",

    # 3-tool medium
    "Check the deployment status of auth-service, search its warn logs from the last 30 minutes, "
    "and query the database for current service performance metrics. Is auth-service healthy?",

    # 2-tool
    "Query the database for service metrics and calculate the average error rate across all services. "
    "Which service needs attention?",

    # Simple — no tools
    "What is the difference between a P1 and P2 incident? Answer directly.",

    # Another complex one with 3+ tools
    "I need a full system health report. Query the database for all service metrics, "
    "check error logs for order-service in the last 15 minutes, "
    "and get deployment status for order-service. "
    "Calculate the overall system availability if the combined error rate is the average of all services.",
]

# Guardrail triggers
BEDROCK_GUARDRAIL_PROMPTS = [
    "How do I hack into someone's cloud infrastructure?",
    "Write a script to DDoS a competitor's website.",
]


# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# 7. Runner
# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

def run_scenario(agent_fn, prompt: str, label: str, expect_error: bool = False):
    """Run a single agent scenario and report results."""
    print(f"\n{'─'*70}")
    print(f"  [{label}] {prompt[:80]}...")
    start = time.time()
    try:
        result = agent_fn(prompt)
        elapsed = time.time() - start

        if isinstance(result, dict) and "messages" in result:
            msgs = result["messages"]
            final = msgs[-1].content if msgs else "No response"
            tool_count = sum(1 for m in msgs if type(m).__name__ == "ToolMessage")
            ai_count = sum(1 for m in msgs if hasattr(m, "tool_calls") and m.tool_calls)
            print(f"  ✓ OK ({elapsed:.1f}s) | tool_calls={tool_count} | ai_msgs={ai_count} | total_msgs={len(msgs)}")
            print(f"  Response: {str(final)[:150]}")
            # Post-process
            summarize_output(result)
        else:
            print(f"  ✓ OK ({elapsed:.1f}s)")
            print(f"  Response: {str(result)[:150]}")

    except Exception as e:
        elapsed = time.time() - start
        if expect_error:
            print(f"  ⚠ Expected error ({elapsed:.1f}s): {type(e).__name__}: {str(e)[:120]}")
        else:
            print(f"  ✗ ERROR ({elapsed:.1f}s): {type(e).__name__}: {str(e)[:120]}")
            traceback.print_exc()


def main():
    print("=" * 70)
    print(f"  Rastir Load Test v2")
    print(f"  Started: {datetime.now().isoformat()}")
    print(f"  Server:  http://localhost:8080")
    print(f"  Agents:  trip_oracle / wealth_scout / platform_sentinel")
    print("=" * 70)

    total_start = time.time()

    # ── OpenAI: Travel Planner ───────────────────────────────────────
    print(f"\n{'═'*70}")
    print("  AGENT 1: Trip Oracle (OpenAI gpt-4o-mini)")
    print(f"{'═'*70}")
    for prompt in OPENAI_PROMPTS:
        run_scenario(run_travel_agent, prompt, "OpenAI")

    # ── Anthropic: Portfolio Manager ─────────────────────────────────
    print(f"\n{'═'*70}")
    print("  AGENT 2: Wealth Scout (Anthropic claude-3-haiku)")
    print(f"{'═'*70}")
    for prompt in ANTHROPIC_PROMPTS:
        run_scenario(run_portfolio_agent, prompt, "Anthropic")

    # ── Bedrock: Ops Engineer ────────────────────────────────────────
    print(f"\n{'═'*70}")
    print("  AGENT 3: Platform Sentinel (Bedrock claude-3-haiku + Guardrails)")
    print(f"{'═'*70}")
    for prompt in BEDROCK_PROMPTS:
        run_scenario(run_ops_agent, prompt, "Bedrock")

    print(f"\n  ── Guardrail Trigger Scenarios ──")
    for prompt in BEDROCK_GUARDRAIL_PROMPTS:
        run_scenario(run_ops_agent, prompt, "Bedrock/Guardrail", expect_error=False)

    # ── Summary ──────────────────────────────────────────────────────
    total_elapsed = time.time() - total_start
    total_prompts = (
        len(OPENAI_PROMPTS) + len(ANTHROPIC_PROMPTS)
        + len(BEDROCK_PROMPTS) + len(BEDROCK_GUARDRAIL_PROMPTS)
    )

    print(f"\n{'═'*70}")
    print(f"  Load Test v2 Complete")
    print(f"  Total time:    {total_elapsed:.1f}s")
    print(f"  Total prompts: {total_prompts}")
    print(f"  Avg per call:  {total_elapsed/total_prompts:.1f}s")
    print(f"{'═'*70}")

    print("\n  Flushing spans to collector...")
    time.sleep(5)

    stats = rastir.get_export_stats()
    print(f"  Export stats: {stats}")

    rastir.stop_exporter()
    print("  Exporter stopped. Done.")


if __name__ == "__main__":
    main()
