"""
Comprehensive test script for all Latch SDK features

This runs all test scenarios:
1. Basic client (sync/async)
2. Context manager approach (new)
3. Report endpoint (new)
4. Pydantic-AI with native models (new)
5. Pydantic-AI with proxy model (legacy)
6. Streaming

Prerequisites:
1. Start the Latch API: task dev:api
2. Start the Latch Engine: task dev:engine (for proxy tests)
3. 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_all.py
"""

import asyncio
import os
import sys
import time

from latch import Latch, AsyncLatch, latch
from latch.integrations.pydantic_ai import LatchModel, LatchModelSettings, get_model


class TestRunner:
    def __init__(self):
        self.results = []
        self.api_key = os.environ.get("LATCH_API_KEY", "test_key")
        self.base_url = os.environ.get("LATCH_BASE_URL", "http://localhost:8081")
        self.has_openai = bool(os.environ.get("OPENAI_API_KEY"))
        self.has_anthropic = bool(os.environ.get("ANTHROPIC_API_KEY"))

    def add_result(self, name: str, passed: bool, duration_ms: float = 0):
        self.results.append((name, passed, duration_ms))

    def print_summary(self):
        print("\n" + "=" * 70)
        print("TEST RESULTS SUMMARY")
        print("=" * 70)

        passed = sum(1 for _, p, _ in self.results if p)
        failed = sum(1 for _, p, _ in self.results if not p)

        for name, result, duration in self.results:
            status = "[+] PASSED" if result else "[-] FAILED"
            duration_str = f"({duration:.0f}ms)" if duration > 0 else ""
            print(f"  {status} {name} {duration_str}")

        print("-" * 70)
        print(f"Total: {passed} passed, {failed} failed out of {len(self.results)} tests")

        if failed == 0:
            print("\n SUCCESS - All tests passed!")
            return 0
        else:
            print(f"\n FAILURE - {failed} test(s) failed")
            return 1


async def test_sync_client(runner: TestRunner):
    """Test synchronous Latch client"""
    print("\n--- Test: Sync Client ---")
    start = time.perf_counter()

    try:
        with Latch(api_key=runner.api_key, base_url=runner.base_url) as client:
            response = client.decide(
                messages=[{"role": "user", "content": "Hello"}],
                task="test",
            )
            assert response.decision_id
            assert response.model
            print(f"  Decision: {response.decision_id}, Model: {response.model}")

        duration = (time.perf_counter() - start) * 1000
        runner.add_result("Sync Client - decide()", True, duration)
    except Exception as e:
        print(f"  Error: {e}")
        runner.add_result("Sync Client - decide()", False)


async def test_async_client(runner: TestRunner):
    """Test asynchronous Latch client"""
    print("\n--- Test: Async Client ---")
    start = time.perf_counter()

    try:
        async with AsyncLatch(api_key=runner.api_key, base_url=runner.base_url) as client:
            response = await client.decide(
                messages=[{"role": "user", "content": "Hello"}],
                task="test",
            )
            assert response.decision_id
            print(f"  Decision: {response.decision_id}, Model: {response.model}")

        duration = (time.perf_counter() - start) * 1000
        runner.add_result("Async Client - decide()", True, duration)
    except Exception as e:
        print(f"  Error: {e}")
        runner.add_result("Async Client - decide()", False)


async def test_report_endpoint(runner: TestRunner):
    """Test report endpoint"""
    print("\n--- Test: Report Endpoint ---")
    start = time.perf_counter()

    try:
        async with AsyncLatch(api_key=runner.api_key, base_url=runner.base_url) as client:
            decision = await client.decide(
                messages=[{"role": "user", "content": "Test"}],
                task="report-test",
            )

            report = await client.report(
                decision_id=decision.decision_id,
                input_tokens=100,
                output_tokens=50,
                llm_latency_ms=250,
            )
            assert report.decision_id == decision.decision_id
            print(f"  Reported cost: ${report.actual_cost:.4f}")

        duration = (time.perf_counter() - start) * 1000
        runner.add_result("Report Endpoint", True, duration)
    except Exception as e:
        print(f"  Error: {e}")
        runner.add_result("Report Endpoint", False)


async def test_context_manager(runner: TestRunner):
    """Test context manager approach"""
    print("\n--- Test: Context Manager ---")
    start = time.perf_counter()

    try:
        os.environ.setdefault("LATCH_API_KEY", runner.api_key)

        async with latch(
            api_key=runner.api_key,
            base_url=runner.base_url,
            task="context-test",
            budget=0.05,
        ) as ctx:
            assert ctx.decision_id
            assert ctx.model_name
            print(f"  Decision: {ctx.decision_id}")
            print(f"  Model: {ctx.model_name}")
            print(f"  Should call LLM: {ctx.should_call_llm}")

            ctx.record_usage(input_tokens=100, output_tokens=50)

        duration = (time.perf_counter() - start) * 1000
        runner.add_result("Context Manager", True, duration)
    except Exception as e:
        print(f"  Error: {e}")
        import traceback
        traceback.print_exc()
        runner.add_result("Context Manager", False)


async def test_get_model_helper(runner: TestRunner):
    """Test get_model helper function"""
    print("\n--- Test: get_model() Helper ---")
    start = time.perf_counter()

    try:
        openai_model = get_model("gpt-4o-mini")
        assert "openai" in type(openai_model).__name__.lower()
        print(f"  gpt-4o-mini -> {type(openai_model).__name__}")

        anthropic_model = get_model("claude-3-sonnet")
        assert "anthropic" in type(anthropic_model).__name__.lower()
        print(f"  claude-3-sonnet -> {type(anthropic_model).__name__}")

        try:
            get_model("invalid-model")
            runner.add_result("get_model() Helper", False)
            return
        except ValueError:
            print("  Correctly raises ValueError for unknown model")

        duration = (time.perf_counter() - start) * 1000
        runner.add_result("get_model() Helper", True, duration)
    except Exception as e:
        print(f"  Error: {e}")
        runner.add_result("get_model() Helper", False)


async def test_native_pydantic_ai_nonstreaming(runner: TestRunner):
    """Test native pydantic-ai without streaming"""
    print("\n--- Test: Native Pydantic-AI (Non-Streaming) ---")

    if not runner.has_openai and not runner.has_anthropic:
        print("  Skipped - no LLM API keys set")
        runner.add_result("Native Pydantic-AI Non-Streaming", True, 0)
        return

    start = time.perf_counter()

    try:
        from pydantic_ai import Agent

        async with latch(
            api_key=runner.api_key,
            base_url=runner.base_url,
            task="native-test",
            budget=0.05,
        ) as ctx:
            model = get_model(ctx.model_name)
            agent = Agent(model=model, system_prompt="Be very brief.")

            result = await agent.run("What is 1+1? Just the number.")
            print(f"  Model: {ctx.model_name}")
            print(f"  Response: {result.output}")

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

        duration = (time.perf_counter() - start) * 1000
        runner.add_result("Native Pydantic-AI Non-Streaming", True, duration)
    except Exception as e:
        print(f"  Error: {e}")
        import traceback
        traceback.print_exc()
        runner.add_result("Native Pydantic-AI Non-Streaming", False)


async def test_native_pydantic_ai_streaming(runner: TestRunner):
    """Test native pydantic-ai with streaming"""
    print("\n--- Test: Native Pydantic-AI (Streaming) ---")

    if not runner.has_openai and not runner.has_anthropic:
        print("  Skipped - no LLM API keys set")
        runner.add_result("Native Pydantic-AI Streaming", True, 0)
        return

    start = time.perf_counter()

    try:
        from pydantic_ai import Agent

        async with latch(
            api_key=runner.api_key,
            base_url=runner.base_url,
            task="streaming-test",
            budget=0.05,
        ) as ctx:
            model = get_model(ctx.model_name)
            agent = Agent(model=model, system_prompt="Be brief.")

            print(f"  Model: {ctx.model_name}")
            print("  Streaming: ", end="", flush=True)

            async with agent.run_stream("Count from 1 to 3.") 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,
                )

        duration = (time.perf_counter() - start) * 1000
        runner.add_result("Native Pydantic-AI Streaming", True, duration)
    except Exception as e:
        print(f"  Error: {e}")
        import traceback
        traceback.print_exc()
        runner.add_result("Native Pydantic-AI Streaming", False)


async def test_proxy_model_legacy(runner: TestRunner):
    """Test legacy proxy model approach (LatchModel)"""
    print("\n--- Test: Proxy Model (Legacy) ---")

    if not runner.has_openai and not runner.has_anthropic:
        print("  Skipped - no LLM API keys set (engine needs them)")
        runner.add_result("Proxy Model (Legacy)", True, 0)
        return

    start = time.perf_counter()

    try:
        from pydantic_ai import Agent

        model = LatchModel(
            api_key=runner.api_key,
            base_url=runner.base_url,
            task="proxy-test",
            budget=0.05,
        )

        agent = Agent(model=model, system_prompt="Be very brief.")

        result = await agent.run("What is 2+2? Just the number.")
        print(f"  Response: {result.output}")

        await model.close()

        duration = (time.perf_counter() - start) * 1000
        runner.add_result("Proxy Model (Legacy)", True, duration)
    except Exception as e:
        print(f"  Error: {e}")
        runner.add_result("Proxy Model (Legacy)", False)


async def test_multiple_tasks(runner: TestRunner):
    """Test multiple sequential tasks with different routing"""
    print("\n--- Test: Multiple Tasks ---")
    start = time.perf_counter()

    try:
        tasks = ["qa", "code-generation", "summarization"]
        models_seen = set()

        for task in tasks:
            async with latch(
                api_key=runner.api_key,
                base_url=runner.base_url,
                task=task,
            ) as ctx:
                models_seen.add(ctx.model_name)
                print(f"  Task '{task}' -> {ctx.model_name}")
                ctx.record_usage(input_tokens=50, output_tokens=25)

        print(f"  Unique models: {len(models_seen)}")

        duration = (time.perf_counter() - start) * 1000
        runner.add_result("Multiple Tasks", True, duration)
    except Exception as e:
        print(f"  Error: {e}")
        runner.add_result("Multiple Tasks", False)


async def main():
    print("=" * 70)
    print("LATCH SDK - COMPREHENSIVE TEST SUITE")
    print("=" * 70)

    runner = TestRunner()

    print("\nConfiguration:")
    print(f"  Base URL: {runner.base_url}")
    print(f"  API Key: {'Set' if runner.api_key != 'test_key' else 'Using test_key'}")
    print(f"  OpenAI: {'Available' if runner.has_openai else 'Not set'}")
    print(f"  Anthropic: {'Available' if runner.has_anthropic else 'Not set'}")

    await test_sync_client(runner)
    await test_async_client(runner)
    await test_report_endpoint(runner)
    await test_context_manager(runner)
    await test_get_model_helper(runner)
    await test_multiple_tasks(runner)

    await test_native_pydantic_ai_nonstreaming(runner)
    await test_native_pydantic_ai_streaming(runner)
    await test_proxy_model_legacy(runner)

    exit_code = runner.print_summary()
    sys.exit(exit_code)


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