"""
Test script for the new Context Manager approach

This tests the decide/report pattern where:
1. Latch provides routing decisions
2. You use your own LLM code (native streaming, tools, etc.)
3. Latch tracks analytics

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_context_manager.py
"""

import asyncio
import os
import time

from latch import latch, AsyncLatch


async def test_context_manager_basic():
    """Test basic context manager usage"""
    print("\n" + "=" * 50)
    print("Test: Basic Context Manager")
    print("=" * 50)

    try:
        async with latch(task="qa", budget=0.05) as ctx:
            print(f"Decision ID: {ctx.decision_id}")
            print(f"Recommended Model: {ctx.model_name}")
            print(f"Should Call LLM: {ctx.should_call_llm}")
            print(f"Estimated Cost: ${ctx.estimated_cost:.4f}")

            await asyncio.sleep(0.1)

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

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


async def test_context_manager_with_metadata():
    """Test context manager with custom metadata"""
    print("\n" + "=" * 50)
    print("Test: Context Manager with Metadata")
    print("=" * 50)

    try:
        async with latch(
            task="code-generation",
            budget=0.10,
            max_latency=5000,
            metadata={"user_id": "test-123", "feature": "code-assist"},
        ) as ctx:
            print(f"Decision ID: {ctx.decision_id}")
            print(f"Model: {ctx.model_name}")

            ctx.record_usage(input_tokens=500, output_tokens=200)

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


async def test_report_endpoint_directly():
    """Test the report endpoint using AsyncLatch client"""
    print("\n" + "=" * 50)
    print("Test: Report Endpoint (Direct)")
    print("=" * 50)

    api_key = os.environ.get("LATCH_API_KEY", "test_key")

    async with AsyncLatch(
        api_key=api_key,
        base_url="http://localhost:8081",
    ) as client:
        try:
            decision = await client.decide(
                messages=[{"role": "user", "content": "Test message"}],
                task="testing",
            )
            print(f"Decision ID: {decision.decision_id}")
            print(f"Model: {decision.model}")

            await asyncio.sleep(0.05)

            report = await client.report(
                decision_id=decision.decision_id,
                input_tokens=150,
                output_tokens=75,
                llm_latency_ms=500,
                model=decision.model,
            )
            print(f"Reported Cost: ${report.actual_cost:.4f}")
            print(f"Total Latency: {report.total_latency_ms}ms")

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


async def test_context_manager_error_handling():
    """Test that context manager reports errors correctly"""
    print("\n" + "=" * 50)
    print("Test: Context Manager Error Handling")
    print("=" * 50)

    try:
        try:
            async with latch(task="error-test") as ctx:
                print(f"Decision ID: {ctx.decision_id}")
                ctx.record_usage(input_tokens=50, output_tokens=0)
                raise ValueError("Simulated LLM error")
        except ValueError:
            pass

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


async def test_multiple_sequential_contexts():
    """Test multiple sequential context usages"""
    print("\n" + "=" * 50)
    print("Test: Multiple Sequential Contexts")
    print("=" * 50)

    try:
        tasks = ["research", "summarize", "translate"]

        for task in tasks:
            async with latch(task=task, budget=0.05) as ctx:
                print(f"  Task '{task}': Model={ctx.model_name}")
                ctx.record_usage(input_tokens=100, output_tokens=50)

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


async def test_context_without_usage_recording():
    """Test context manager without calling record_usage"""
    print("\n" + "=" * 50)
    print("Test: Context Without Usage Recording")
    print("=" * 50)

    try:
        async with latch(task="quick-test") as ctx:
            print(f"Decision ID: {ctx.decision_id}")

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


async def main():
    print("=" * 60)
    print("Latch SDK - Context Manager Tests")
    print("=" * 60)

    api_key = os.environ.get("LATCH_API_KEY")
    if not api_key:
        print("\nWarning: LATCH_API_KEY not set")
        print("Set it or create one via the dashboard.\n")

    results = []

    results.append(("Basic Context Manager", await test_context_manager_basic()))
    results.append(("Context with Metadata", await test_context_manager_with_metadata()))
    results.append(("Report Endpoint Direct", await test_report_endpoint_directly()))
    results.append(("Error Handling", await test_context_manager_error_handling()))
    results.append(("Multiple Sequential", await test_multiple_sequential_contexts()))
    results.append(("No Usage Recording", await test_context_without_usage_recording()))

    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. Check output above.")


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