#!/usr/bin/env python3
import asyncio
import os

from gemini import GeminiClient, GeminiOptions

PROMPT = os.environ.get("TEST_PROMPT", "Say hi in exactly 3 words")

CONFIGS = [
    {"label": "flash-lite", "model": "gemini-2.5-flash-lite"},
    {"label": "flash", "model": "gemini-2.5-flash"},
    {"label": "pro", "model": "gemini-2.5-pro"},
]


async def run_one(cfg):
    opts = GeminiOptions(model=cfg["model"])
    async with GeminiClient(options=opts) as client:
        text = ""
        async for chunk in client.send_message_stream(PROMPT):
            text += chunk.text_delta
        print(f"  [{cfg['label']}] {text.strip()[:80]}")


async def main():
    mode = os.environ.get("TEST_MODE", "single")
    configs = CONFIGS if mode == "all" else [CONFIGS[0]]

    for cfg in configs:
        print(f">>> {cfg['label']}: model={cfg['model']}")
        await run_one(cfg)
        print()


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