"""
Local testing script for Latch Python SDK

Prerequisites:
1. Start the Latch API service: task dev:api
2. Start the Latch Engine service: task dev:engine
3. Create an API key via dashboard or database
4. Set environment variables:
   - LATCH_API_KEY: Your API key
   - OPENAI_API_KEY: For OpenAI models
   - ANTHROPIC_API_KEY: For Anthropic models

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

import os
import asyncio
from latch import Latch, AsyncLatch


def test_sync_decide():
    """Test synchronous decision endpoint"""
    api_key = os.environ.get("LATCH_API_KEY", "test_key")

    client = Latch(
        api_key=api_key,
        base_url="http://localhost:8081",
    )

    try:
        print("Testing decide() endpoint...")
        response = client.decide(
            messages=[
                {"role": "user", "content": "What is 2 + 2?"}
            ],
            task="math",
            budget=0.01,
        )
        print(f"Decision ID: {response.decision_id}")
        print(f"Recommended Model: {response.model}")
        print(f"Estimated Cost: ${response.estimated_cost:.4f}")
        print(f"Decision Latency: {response.metadata.decision_latency_ms}ms")
        print("decide() test passed!")
        return True
    except Exception as e:
        print(f"decide() test failed: {e}")
        return False
    finally:
        client.close()


def test_sync_run():
    """Test synchronous execute endpoint"""
    api_key = os.environ.get("LATCH_API_KEY", "test_key")

    client = Latch(
        api_key=api_key,
        base_url="http://localhost:8081",
    )

    try:
        print("\nTesting run() endpoint...")
        response = client.run(
            messages=[
                {"role": "user", "content": "What is the capital of France? Answer in one word."}
            ],
            task="qa",
            budget=0.05,
        )
        print(f"Decision ID: {response.decision_id}")
        print(f"Model Used: {response.model}")
        print(f"Response: {response.content[:100]}...")
        print(f"Actual Cost: ${response.cost:.4f}")
        print(f"Input Tokens: {response.usage.input_tokens}")
        print(f"Output Tokens: {response.usage.output_tokens}")
        print("run() test passed!")
        return True
    except Exception as e:
        print(f"run() test failed: {e}")
        return False
    finally:
        client.close()


async def test_async_decide():
    """Test asynchronous decision endpoint"""
    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:
            print("\nTesting async decide() endpoint...")
            response = await client.decide(
                messages=[
                    {"role": "system", "content": "You are a helpful assistant."},
                    {"role": "user", "content": "Explain Python in one sentence."}
                ],
                task="explanation",
            )
            print(f"Decision ID: {response.decision_id}")
            print(f"Recommended Model: {response.model}")
            print("async decide() test passed!")
            return True
        except Exception as e:
            print(f"async decide() test failed: {e}")
            return False


async def test_async_run():
    """Test asynchronous execute endpoint"""
    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:
            print("\nTesting async run() endpoint...")
            response = await client.run(
                messages=[
                    {"role": "user", "content": "Write a haiku about coding."}
                ],
                task="creative_writing",
                budget=0.10,
            )
            print(f"Decision ID: {response.decision_id}")
            print(f"Model Used: {response.model}")
            print(f"Response:\n{response.content}")
            print("async run() test passed!")
            return True
        except Exception as e:
            print(f"async run() test failed: {e}")
            return False


def main():
    print("=" * 50)
    print("Latch Python SDK - Local Test")
    print("=" * 50)

    api_key = os.environ.get("LATCH_API_KEY")
    if not api_key:
        print("\nWarning: LATCH_API_KEY not set, using 'test_key'")
        print("Create an API key via the dashboard first.\n")

    results = []

    results.append(("sync decide()", test_sync_decide()))
    results.append(("sync run()", test_sync_run()))
    results.append(("async decide()", asyncio.run(test_async_decide())))
    results.append(("async run()", asyncio.run(test_async_run())))

    print("\n" + "=" * 50)
    print("Test Results:")
    print("=" * 50)
    for name, passed in results:
        status = "PASSED" if passed else "FAILED"
        print(f"  {name}: {status}")

    all_passed = all(passed for _, passed in results)
    if all_passed:
        print("\nAll tests passed!")
    else:
        print("\nSome tests failed. Check the output above.")


if __name__ == "__main__":
    main()
