"""Полный тест всех операций SDK на проде.

Запуск:
    export RUNBLOB_API_KEY="rb_your_key"
    python examples/test_full_flow.py

Тестирует:
    1. generate()          — создание задачи
    2. get_status()        — проверка статуса
    3. wait_for_result()   — поллинг до завершения
    4. generate_and_wait() — всё в одном вызове
    5. list()              — список генераций
    6. generate_batch()    — пакетная генерация
"""
import asyncio
import time
import sys

from runblob import (
    RunBlob,
    ImageStatus,
    RunBlobError,
    AuthenticationError,
    InsufficientCreditsError,
    RateLimitError,
)


# ─── Настройки ────────────────────────────────────────────────────────────────
API_KEY = None  # None = берёт из RUNBLOB_API_KEY env var
BASE_URL = None  # None = https://platform.runblob.io (по умолчанию)


class Colors:
    GREEN = "\033[92m"
    RED = "\033[91m"
    YELLOW = "\033[93m"
    CYAN = "\033[96m"
    BOLD = "\033[1m"
    END = "\033[0m"


def ok(msg: str):
    print(f"  {Colors.GREEN}✓{Colors.END} {msg}")


def fail(msg: str):
    print(f"  {Colors.RED}✗{Colors.END} {msg}")


def info(msg: str):
    print(f"  {Colors.CYAN}ℹ{Colors.END} {msg}")


def header(msg: str):
    print(f"\n{Colors.BOLD}{Colors.YELLOW}{'═' * 60}{Colors.END}")
    print(f"{Colors.BOLD}{Colors.YELLOW}  {msg}{Colors.END}")
    print(f"{Colors.BOLD}{Colors.YELLOW}{'═' * 60}{Colors.END}")


passed = 0
failed = 0


async def run_test(name: str, coro):
    global passed, failed
    print(f"\n  {Colors.BOLD}▸ {name}{Colors.END}")
    t0 = time.monotonic()
    try:
        await coro
        elapsed = time.monotonic() - t0
        ok(f"Прошёл за {elapsed:.2f}s")
        passed += 1
    except Exception as e:
        elapsed = time.monotonic() - t0
        fail(f"Упал за {elapsed:.2f}s — {type(e).__name__}: {e}")
        failed += 1


async def main():
    global passed, failed

    kwargs = {}
    if API_KEY:
        kwargs["api_key"] = API_KEY
    if BASE_URL:
        kwargs["base_url"] = BASE_URL

    header("RunBlob SDK — Тест на проде")

    try:
        client = RunBlob(**kwargs)
    except ValueError as e:
        fail(f"Не удалось создать клиент: {e}")
        fail("Установи RUNBLOB_API_KEY или укажи API_KEY в скрипте")
        sys.exit(1)

    async with client:
        info(f"Подключено к {client.config.base_url}")
        info(f"SDK версия: {client.config.version}")

        # ── 1. generate() ─────────────────────────────────────────────
        task_uuid = None

        async def test_generate():
            nonlocal task_uuid
            result = await client.nanobanana.generate(
                prompt="A red apple on a white table, studio lighting",
                model="standard",
            )
            assert result.task_uuid, "task_uuid пустой"
            assert result.status == ImageStatus.PENDING, f"Ожидал pending, получил {result.status}"
            assert result.price, "price пустой"
            task_uuid = result.task_uuid
            info(f"task_uuid: {task_uuid}")
            info(f"price: {result.price}")

        await run_test("generate() — создание задачи", test_generate())

        # ── 2. get_status() ───────────────────────────────────────────
        async def test_get_status():
            assert task_uuid, "Нет task_uuid из предыдущего теста"
            result = await client.nanobanana.get_status(task_uuid)
            assert result.task_uuid == task_uuid
            assert result.status in (
                ImageStatus.PENDING,
                ImageStatus.PROCESSING,
                ImageStatus.COMPLETED,
                ImageStatus.FAILED,
            )
            info(f"Статус: {result.status.value}")

        await run_test("get_status() — проверка статуса", test_get_status())

        # ── 3. wait_for_result() ──────────────────────────────────────
        completed_url = None

        async def test_wait():
            nonlocal completed_url
            assert task_uuid, "Нет task_uuid"
            result = await client.nanobanana.wait_for_result(
                task_uuid,
                timeout=120.0,
            )
            assert result.is_terminal, f"Не терминальный статус: {result.status}"
            if result.is_completed:
                assert result.result_image_url, "URL картинки пустой"
                completed_url = result.result_image_url
                info(f"URL: {completed_url}")
            else:
                info(f"Завершился с ошибкой: {result.message}")

        await run_test("wait_for_result() — поллинг до завершения", test_wait())

        # ── 4. generate_and_wait() ────────────────────────────────────
        async def test_generate_and_wait():
            result = await client.nanobanana.generate_and_wait(
                prompt="A blue ocean wave, cinematic style",
                model="standard",
                aspect_ratio="16:9",
                timeout=120.0,
            )
            assert result.is_terminal
            if result.is_completed:
                info(f"URL: {result.result_image_url}")
            else:
                info(f"Статус: {result.status.value}, сообщение: {result.message}")

        await run_test("generate_and_wait() — генерация + ожидание", test_generate_and_wait())

        # ── 5. list() ─────────────────────────────────────────────────
        async def test_list():
            results = await client.nanobanana.list(limit=5)
            assert isinstance(results, list)
            info(f"Получено {len(results)} генераций")
            for r in results[:3]:
                info(f"  {r.task_uuid[:8]}... — {r.status.value}")

        await run_test("list() — список генераций", test_list())

        # ── 6. generate_batch() ───────────────────────────────────────
        async def test_batch():
            prompts = [
                "A green forest at dawn",
                "A mountain peak covered in snow",
            ]
            results = await client.nanobanana.generate_batch(
                prompts,
                model="standard",
                concurrency=2,
            )
            assert len(results) == len(prompts), f"Ожидал {len(prompts)} результатов, получил {len(results)}"
            for i, r in enumerate(results):
                assert r.task_uuid, f"Пустой task_uuid для промпта #{i}"
                info(f"  #{i}: {r.task_uuid[:8]}... — {r.status.value}")

        await run_test("generate_batch() — пакетная генерация", test_batch())

        # ── 7. generate_and_wait() с параметрами ──────────────────────
        async def test_params():
            result = await client.nanobanana.generate_and_wait(
                prompt="Anime style girl with blue hair",
                model="standard",
                quality="standard",
                aspect_ratio="9:16",
                temperature=0.9,
                timeout=120.0,
            )
            assert result.is_terminal
            if result.is_completed:
                info(f"URL: {result.result_image_url}")
            else:
                info(f"Ошибка: {result.message}")

        await run_test("generate_and_wait() — с кастомными параметрами", test_params())

    # ── Итоги ─────────────────────────────────────────────────────────
    header("Результаты")
    total = passed + failed
    print(f"\n  Всего:   {total}")
    print(f"  {Colors.GREEN}Прошло:  {passed}{Colors.END}")
    if failed:
        print(f"  {Colors.RED}Упало:   {failed}{Colors.END}")
    print()

    if failed:
        sys.exit(1)


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