"""Integration tests for streaming text generation - all registered models."""

import warnings

# Suppress deprecation warnings from legacy capability packages
warnings.filterwarnings(
    "ignore",
    message=".*capability parameter is deprecated.*",
    category=DeprecationWarning,
)

import pytest  # noqa: E402

from celeste import (  # noqa: E402
    Modality,
    Model,
    Operation,
    create_client,
    list_models,
)
from celeste.modalities.text import TextChunk, TextUsage  # noqa: E402
from celeste.providers.google.auth import GoogleADC  # noqa: E402

TEST_MAX_TOKENS = 200


@pytest.mark.parametrize(
    "model",
    [
        m
        for m in list_models(modality=Modality.TEXT, operation=Operation.GENERATE)
        if m.streaming
        and not m.optional_input_types  # Media-capable models tested in test_stream_analyze_*
    ],
    ids=lambda m: f"{m.provider}-{m.id}",
)
@pytest.mark.integration
@pytest.mark.asyncio
async def test_stream_generate(model: Model) -> None:
    """Test streaming text generation for all streaming-capable models.

    Dynamically discovers all streaming models and verifies each can stream text.
    Failures indicate deprecated or misconfigured models.
    """
    client = create_client(
        modality=Modality.TEXT,
        model=model,
    )

    chunks: list[TextChunk] = []
    async for chunk in client.stream.generate(prompt="Hi", max_tokens=TEST_MAX_TOKENS):
        chunks.append(chunk)

    # Assert - empty stream is valid for reasoning models that use all tokens for thinking
    if not chunks:
        return

    # Assert - received chunks are valid
    assert all(isinstance(c, TextChunk) for c in chunks), (
        "All chunks should be TextChunk"
    )

    # Assert - usage in final chunks (provider-dependent)
    usage_chunks = [c for c in chunks if c.usage is not None]
    if usage_chunks:
        usage = usage_chunks[-1].usage
        assert isinstance(usage, TextUsage), f"Expected TextUsage, got {type(usage)}"
        if usage.output_tokens is not None and usage.output_tokens > TEST_MAX_TOKENS:
            warnings.warn(
                f"Model {model.provider}/{model.id} exceeded max_tokens: {usage.output_tokens} > {TEST_MAX_TOKENS}",
                stacklevel=1,
            )


@pytest.mark.integration
def test_sync_stream_generate() -> None:
    """Test sync streaming wrapper works correctly.

    Single model smoke test - sync stream iteration bridges async internally.
    """
    models = [
        m
        for m in list_models(modality=Modality.TEXT, operation=Operation.GENERATE)
        if m.streaming and not m.optional_input_types
    ]
    model = models[0]

    client = create_client(
        modality=Modality.TEXT,
        model=model,
    )

    for _chunk in client.sync.stream.generate(prompt="Hi", max_tokens=TEST_MAX_TOKENS):
        pass  # Just exhaust the stream


@pytest.mark.parametrize(
    ("provider", "model", "location"),
    [
        ("google", "gemini-2.5-flash", "global"),
        ("anthropic", "claude-haiku-4-5", "us-east5"),
        ("mistral", "mistral-small-2503", "us-central1"),
        ("deepseek", "deepseek-ai/deepseek-v3.2-maas", "global"),
    ],
)
@pytest.mark.integration
@pytest.mark.asyncio
async def test_vertex_stream_generate(provider: str, model: str, location: str) -> None:
    """Test streaming text generation via Vertex AI - one model per provider."""
    client = create_client(
        modality=Modality.TEXT,
        provider=provider,
        model=model,
        auth=GoogleADC(location=location),
    )

    chunks: list[TextChunk] = []
    async for chunk in client.stream.generate(prompt="Hi", max_tokens=TEST_MAX_TOKENS):
        chunks.append(chunk)

    if not chunks:
        return

    assert all(isinstance(c, TextChunk) for c in chunks)
