"""
Manual smoke-test runner for agent_meter.

Usage:
  OPENAI_API_KEY=... python -m agent_meter
"""

import asyncio
import os

from a2a_llm_tracker  import Meter, PricingRegistry, meter_context
from a2a_llm_tracker.integrations.litellm import LiteLLM
from a2a_llm_tracker.sinks.jsonl import JSONLSink
from dotenv import load_dotenv
load_dotenv()

def setup_meter() -> LiteLLM:
    pricing = PricingRegistry()

    # Example prices — adjust as needed
    pricing.set_price(
        provider="openai",
        model="openai/gpt-4.1",
        input_per_million=2.0,
        output_per_million=8.0,
    )

    meter = Meter(
        pricing=pricing,
        sinks=[JSONLSink("debug_usage.jsonl")],
        project="agent-meter-smoke-test",
    )

    return LiteLLM(meter=meter)


def run_sync(llm: LiteLLM) -> None:
    print("\n=== SYNC NON-STREAMING ===")
    resp = llm.completion(
        model="openai/gpt-4.1",
        messages=[{"role": "user", "content": "Say hello in one sentence."}],
    )
    print(resp)


def run_stream(llm: LiteLLM) -> None:
    print("\n=== SYNC STREAMING ===")
    for chunk in llm.completion(
        model="openai/gpt-4.1",
        messages=[{"role": "user", "content": "Write a short haiku about testing."}],
        stream=True,
    ):
        print(chunk, end="", flush=True)
    print()


async def run_async(llm: LiteLLM) -> None:
    print("\n=== ASYNC NON-STREAMING ===")
    resp = await llm.acompletion(
        model="openai/gpt-4.1",
        messages=[{"role": "user", "content": "Async hello!"}],
    )
    print(resp)


async def run_async_stream(llm: LiteLLM) -> None:
    print("\n=== ASYNC STREAMING ===")
    stream = await llm.acompletion(
        model="openai/gpt-4.1",
        messages=[{"role": "user", "content": "Stream something async."}],
        stream=True,
    )
    async for chunk in stream:
        print(chunk, end="", flush=True)
    print()


def main() -> None:
    if not os.getenv("OPENAI_API_KEY"):
        print("⚠️  OPENAI_API_KEY is not set. Exiting.")
        return

    llm = setup_meter()

    with meter_context(agent_id="smoke-test-agent", session_id="local"):
        run_sync(llm)
        run_stream(llm)
        asyncio.run(run_async(llm))
        asyncio.run(run_async_stream(llm))

    print("\n✔ Smoke test complete. Check debug_usage.jsonl")


if __name__ == "__main__":
    main()