"""
Test script for pydantic-ai with NATIVE models via Latch context manager

This is the recommended approach:
1. Use Latch context manager to get routing decision
2. Use native pydantic-ai models (OpenAI, Anthropic) for actual LLM calls
3. Native streaming, tools, and all features work!

Prerequisites:
1. Start the Latch API: task dev:api
2. Set environment variables:
   - LATCH_API_KEY: Your Latch API key
   - OPENAI_API_KEY: For OpenAI models
   - ANTHROPIC_API_KEY: For Anthropic models

Run:
    cd packages/sdk-python
    uv run python examples/test_pydantic_ai_native.py
"""

import asyncio
import os

from pydantic_ai import Agent

from latch import latch
from latch.integrations.pydantic_ai import get_model


async def test_non_streaming():
    """Test non-streaming with native model"""
    print("\n" + "=" * 50)
    print("Test: Non-Streaming with Native Model")
    print("=" * 50)

    try:
        async with latch(task="qa", budget=0.05) as ctx:
            print(f"Latch Decision: {ctx.model_name}")

            model = get_model(ctx.model_name)
            agent = Agent(
                model=model,
                system_prompt="You are a helpful assistant. Keep responses very brief.",
            )

            result = await agent.run("What is 2 + 2? Answer with just the number.")

            print(f"Response: {result.output}")
            print(f"Input tokens: {result.usage().input_tokens}")
            print(f"Output tokens: {result.usage().output_tokens}")

            ctx.record_usage(
                input_tokens=result.usage().input_tokens or 0,
                output_tokens=result.usage().output_tokens or 0,
            )

        print("Non-streaming test PASSED")
        return True
    except Exception as e:
        print(f"Non-streaming test FAILED: {e}")
        import traceback
        traceback.print_exc()
        return False


async def test_streaming():
    """Test streaming with native model"""
    print("\n" + "=" * 50)
    print("Test: Streaming with Native Model")
    print("=" * 50)

    try:
        async with latch(task="creative-writing", budget=0.10) as ctx:
            print(f"Latch Decision: {ctx.model_name}")

            model = get_model(ctx.model_name)
            agent = Agent(
                model=model,
                system_prompt="You are a creative writer.",
            )

            print("Streaming response: ", end="", flush=True)

            async with agent.run_stream("Write a haiku about coding.") as stream:
                async for chunk in stream.stream_text():
                    print(chunk, end="", flush=True)

                print()

                usage = stream.usage()
                ctx.record_usage(
                    input_tokens=usage.input_tokens or 0,
                    output_tokens=usage.output_tokens or 0,
                )

                print(f"Total usage: {usage}")

        print("Streaming test PASSED")
        return True
    except Exception as e:
        print(f"Streaming test FAILED: {e}")
        import traceback
        traceback.print_exc()
        return False


async def test_different_models():
    """Test that Latch returns different models for different tasks"""
    print("\n" + "=" * 50)
    print("Test: Different Models for Different Tasks")
    print("=" * 50)

    try:
        tasks = [
            ("qa", 0.001),
            ("chat", 0.005),
            ("summarization", 0.02),
            ("code-generation", 0.10),
            ("reasoning", 0.20),
            ("creative-writing", 0.15),
        ]

        models_used = []

        for task, budget in tasks:
            async with latch(task=task, budget=budget) as ctx:
                print(f"  Task '{task}' (budget=${budget:.3f}): {ctx.model_name}")
                models_used.append(ctx.model_name)
                ctx.record_usage(input_tokens=50, output_tokens=25)

        unique_models = set(models_used)
        print(f"\nModels selected: {unique_models}")
        print(f"Model diversity: {len(unique_models)} unique models from {len(tasks)} tasks")
        print("Different models test PASSED")
        return True
    except Exception as e:
        print(f"Different models test FAILED: {e}")
        return False


async def test_actual_llm_call():
    """Test a real LLM call through native model"""
    print("\n" + "=" * 50)
    print("Test: Actual LLM Call with Native Model")
    print("=" * 50)

    try:
        async with latch(task="reasoning", budget=0.15) as ctx:
            print(f"Using model: {ctx.model_name}")

            if not ctx.should_call_llm:
                print("Latch blocked this request (budget/policy)")
                return True

            model = get_model(ctx.model_name)
            agent = Agent(
                model=model,
                system_prompt="You are a math tutor. Be concise.",
            )

            result = await agent.run(
                "If I have 3 apples and buy 5 more, how many do I have? Just the number."
            )

            print(f"Question: 3 + 5 = ?")
            print(f"Response: {result.output}")

            ctx.record_usage(
                input_tokens=result.usage().input_tokens or 0,
                output_tokens=result.usage().output_tokens or 0,
            )

        print("Actual LLM call test PASSED")
        return True
    except Exception as e:
        print(f"Actual LLM call test FAILED: {e}")
        import traceback
        traceback.print_exc()
        return False


async def test_get_model_function():
    """Test the get_model helper function"""
    print("\n" + "=" * 50)
    print("Test: get_model() Helper Function")
    print("=" * 50)

    try:
        test_cases = [
            ("gpt-4o-mini", "OpenAI"),
            ("gpt-4o", "OpenAI"),
            ("gpt-4.1", "OpenAI"),
            ("gpt-4.1-mini", "OpenAI"),
            ("gpt-5", "OpenAI"),
            ("gpt-5-mini", "OpenAI"),
            ("o1", "OpenAI"),
            ("o3", "OpenAI"),
            ("o4-mini", "OpenAI"),
            ("claude-3-sonnet", "Anthropic"),
            ("claude-sonnet-4-5-20250514", "Anthropic"),
            ("claude-opus-4-5-20250514", "Anthropic"),
            ("claude-haiku-4-5-20250514", "Anthropic"),
        ]

        for model_name, expected_provider in test_cases:
            model = get_model(model_name)
            provider = type(model).__name__
            print(f"  {model_name} -> {provider}")
            assert expected_provider.lower() in provider.lower(), f"Expected {expected_provider} provider"

        try:
            get_model("unknown-model-xyz")
            print("Should have raised ValueError!")
            return False
        except ValueError as e:
            print(f"  Correctly raised ValueError for unknown model")

        print("get_model() test PASSED")
        return True
    except Exception as e:
        print(f"get_model() test FAILED: {e}")
        return False


async def main():
    print("=" * 60)
    print("Latch SDK - Pydantic-AI Native Model Tests")
    print("=" * 60)

    api_key = os.environ.get("LATCH_API_KEY")
    openai_key = os.environ.get("OPENAI_API_KEY")
    anthropic_key = os.environ.get("ANTHROPIC_API_KEY")

    print("\nEnvironment Check:")
    print(f"  LATCH_API_KEY: {'Set' if api_key else 'NOT SET'}")
    print(f"  OPENAI_API_KEY: {'Set' if openai_key else 'NOT SET'}")
    print(f"  ANTHROPIC_API_KEY: {'Set' if anthropic_key else 'NOT SET'}")

    if not api_key:
        print("\nWarning: LATCH_API_KEY not set. Tests may fail.")
    if not openai_key and not anthropic_key:
        print("\nWarning: No LLM API keys set. Actual LLM tests will fail.")

    results = []

    results.append(("get_model() Helper", await test_get_model_function()))
    results.append(("Different Models", await test_different_models()))

    if openai_key or anthropic_key:
        results.append(("Non-Streaming", await test_non_streaming()))
        results.append(("Streaming", await test_streaming()))
        results.append(("Actual LLM Call", await test_actual_llm_call()))
    else:
        print("\nSkipping LLM tests - no API keys set")

    print("\n" + "=" * 60)
    print("Test Results Summary")
    print("=" * 60)

    passed = 0
    failed = 0
    for name, result in results:
        status = "PASSED" if result else "FAILED"
        symbol = "[+]" if result else "[-]"
        print(f"  {symbol} {name}: {status}")
        if result:
            passed += 1
        else:
            failed += 1

    print(f"\nTotal: {passed} passed, {failed} failed")

    if failed == 0:
        print("\nAll tests passed!")
    else:
        print(f"\n{failed} test(s) failed.")


if __name__ == "__main__":
    asyncio.run(main())
