import tempfile

import pytest
from chatlas import (
    AssistantTurn,
    ChatAnthropic,
    ChatGoogle,
    ChatGroq,
    ChatOpenAI,
    ChatVertex,
)
from chatlas._batch_chat import (
    BatchJob,
    batch_chat,
    batch_chat_completed,
    batch_chat_structured,
    batch_chat_text,
)
from chatlas._provider import BatchStatus
from pydantic import BaseModel

from .conftest import VCR_MATCH_ON_WITHOUT_BODY, make_vcr_config


# Don't match on body - temp file names are dynamic
@pytest.fixture(scope="module")
def vcr_config():
    return make_vcr_config(match_on=VCR_MATCH_ON_WITHOUT_BODY)


class CountryCapital(BaseModel):
    name: str


def test_can_retrieve_batch(test_batch_dir):
    chat = ChatOpenAI(model="gpt-4o-mini")
    prompts = ["What's the capital of France?", "What's the capital of Germany?"]

    chats = batch_chat(
        chat,
        prompts,
        test_batch_dir / "country-capitals.json",
    )
    assert len(chats) == 2
    assert chats[0] is not None
    assert chats[1] is not None
    turns1 = chats[0].get_turns()
    turns2 = chats[1].get_turns()
    assert len(turns1) == 2
    assert len(turns2) == 2
    assert isinstance(turns1[1], AssistantTurn)
    tokens = turns1[1].tokens or []
    assert len(tokens) == 3

    out = batch_chat_text(
        chat,
        prompts,
        test_batch_dir / "country-capitals.json",
    )
    assert len(out) == 2
    assert out[0] is not None
    assert out[1] is not None
    assert "Paris" in out[0]
    assert "Berlin" in out[1]

    capitals = batch_chat_structured(
        chat,
        prompts,
        test_batch_dir / "country-capitals-structured.json",
        CountryCapital,
    )
    assert len(capitals) == 2
    assert capitals[0] is not None
    assert capitals[1] is not None
    assert capitals[0].name == "Paris"
    assert capitals[1].name == "Berlin"


def test_can_retrieve_google_batch(test_batch_dir):
    """Replays a real, completed Gemini batch job recorded live end-to-end
    (submit -> poll -> retrieve -> parse). This is also a regression test for
    a real crash found during that live run: gemini-3.6-flash attaches a
    non-UTF-8 thought_signature byte blob to every response part, which broke
    JSON-serializing the on-disk batch state (fixed via model_dump(mode="json")
    in GoogleProvider's batch_submit/batch_poll/batch_retrieve)."""
    chat = ChatGoogle(model="gemini-flash-latest")
    prompts = ["What's the capital of France?", "What's the capital of Germany?"]

    chats = batch_chat(
        chat,
        prompts,
        test_batch_dir / "google-capitals.json",
    )
    assert chats is not None
    assert len(chats) == 2
    assert chats[0] is not None
    assert chats[1] is not None
    turns1 = chats[0].get_turns()
    turns2 = chats[1].get_turns()
    assert len(turns1) == 2
    assert len(turns2) == 2
    assert isinstance(turns1[1], AssistantTurn)
    assert turns1[1].tokens is not None

    out = batch_chat_text(
        chat,
        prompts,
        test_batch_dir / "google-capitals.json",
    )
    assert out is not None
    assert len(out) == 2
    assert out[0] is not None
    assert out[1] is not None
    assert "Paris" in out[0]
    assert "Berlin" in out[1]


@pytest.mark.vcr
def test_can_submit_openai_batch():
    with tempfile.NamedTemporaryFile() as temp_file:
        chat = ChatOpenAI()
        prompts = ["What's the capital of France?", "What's the capital of Germany?"]
        job = BatchJob(chat, prompts, temp_file.name, wait=False)
        assert job.stage == "submitting"
        job.step()
        assert job.stage == "waiting"


@pytest.mark.vcr
def test_can_submit_anthropic_batch():
    with tempfile.NamedTemporaryFile() as temp_file:
        chat = ChatAnthropic()
        prompts = ["What's the capital of France?", "What's the capital of Germany?"]
        job = BatchJob(chat, prompts, temp_file.name, wait=False)
        assert job.stage == "submitting"
        job.step()
        assert job.stage == "waiting"


def test_groq_batch_support_via_inheritance():
    """ChatGroq() gets full OpenAI-shaped batch support purely via inheritance.

    Groq's real batch API (verified against ellmer's provider-groq.R) is the
    OpenAI Batches API shape, verbatim -- so no Groq-specific batch code is
    needed. This guards against something later overriding that inheritance.
    """
    from chatlas._provider_openai_generic import OpenAIAbstractProvider

    chat = ChatGroq(api_key="dummy")
    provider_cls = type(chat.provider)
    assert chat.provider.has_batch_support()
    assert provider_cls.batch_submit is OpenAIAbstractProvider.batch_submit
    assert provider_cls.batch_poll is OpenAIAbstractProvider.batch_poll
    assert provider_cls.batch_status is OpenAIAbstractProvider.batch_status
    assert provider_cls.batch_retrieve is OpenAIAbstractProvider.batch_retrieve


@pytest.mark.vcr
def test_can_submit_google_batch():
    with tempfile.NamedTemporaryFile() as temp_file:
        chat = ChatGoogle()
        prompts = ["What's the capital of France?", "What's the capital of Germany?"]
        job = BatchJob(chat, prompts, temp_file.name, wait=False)
        assert job.stage == "submitting"
        job.step()
        assert job.stage == "waiting"


def test_informative_errors(test_batch_dir):
    with pytest.raises(ValueError, match="not supported by this provider"):
        batch_chat(
            ChatVertex(),
            [],
            "foo.json",
        )

    with pytest.raises(ValueError, match="provider doesn't match stored value"):
        batch_chat(
            ChatAnthropic(),
            [],
            test_batch_dir / "country-capitals.json",
        )

    with pytest.raises(ValueError, match="model doesn't match stored value"):
        batch_chat(
            ChatOpenAI(model="gpt-5"),
            [],
            test_batch_dir / "country-capitals.json",
        )

    with pytest.raises(ValueError, match="prompts doesn't match stored value"):
        batch_chat(
            ChatOpenAI(model="gpt-4o-mini"),
            ["foo"],
            test_batch_dir / "country-capitals.json",
        )

    with pytest.raises(ValueError, match="user_turns doesn't match stored value"):
        batch_chat(
            ChatOpenAI(model="gpt-4o-mini", system_prompt="foo"),
            [
                "What's the capital of France?",
                "What's the capital of Germany?",
            ],
            test_batch_dir / "country-capitals.json",
        )


def ChatOpenAIMockBatchSubmit(*, working: bool = False, **kwargs):
    chat = ChatOpenAI(**kwargs)

    def batch_submit_mock(*args, **kwargs):
        return {"id": "123"}

    def batch_poll_mock(*args, **kwargs):
        return {"id": "123", "results": True}

    def batch_status_mock(*args, **kwargs):
        return BatchStatus(
            working=working,
            n_processing=0,
            n_failed=0,
            n_succeeded=1,
        )

    def batch_retrieve_mock(*args, **kwargs):
        return [{"x": 1, "y": 2}]

    chat.provider.batch_submit = batch_submit_mock
    chat.provider.batch_poll = batch_poll_mock
    chat.provider.batch_status = batch_status_mock
    chat.provider.batch_retrieve = batch_retrieve_mock

    return chat


def test_steps_in_logical_order():
    with tempfile.NamedTemporaryFile() as temp_file:
        chat = ChatOpenAIMockBatchSubmit()
        prompts = ["What's your name?"]
        job = BatchJob(chat, prompts, temp_file.name)

        def completed():
            return batch_chat_completed(chat, prompts, temp_file.name)

        assert job.stage == "submitting"
        assert not completed()

        job.step()
        assert job.stage == "waiting"
        job._load_state()
        assert job.stage == "waiting"
        assert job.batch == {"id": "123"}
        assert completed()

        job.step()
        assert job.stage == "retrieving"
        job._load_state()
        assert job.stage == "retrieving"
        assert job.batch == {"id": "123", "results": True}
        assert completed()

        job.step()
        assert job.stage == "done"
        job._load_state()
        assert job.stage == "done"
        assert job.results == [{"x": 1, "y": 2}]
        assert completed()


def test_run_all_steps_at_once():
    with tempfile.NamedTemporaryFile() as temp_file:
        chat = ChatOpenAIMockBatchSubmit()
        prompts = ["What's your name?"]
        job = BatchJob(chat, prompts, temp_file.name)

        job = job.step_until_done()
        assert job is not None
        assert job.stage == "done"
        assert job.results == [{"x": 1, "y": 2}]


def test_can_avoid_blocking():
    with tempfile.NamedTemporaryFile() as temp_file:
        chat = ChatOpenAIMockBatchSubmit(working=True)
        job = BatchJob(
            chat,
            ["What's your name?"],
            temp_file.name,
            wait=False,
        )

        assert job.step_until_done() is None


def test_batch_chat_returns_none_when_incomplete():
    with tempfile.NamedTemporaryFile(suffix=".json") as temp_file:
        chat = ChatOpenAIMockBatchSubmit(working=True)
        prompts = ["What's your name?"]

        result = batch_chat(chat, prompts, temp_file.name, wait=False)
        assert result is None

        result_text = batch_chat_text(chat, prompts, temp_file.name, wait=False)
        assert result_text is None

        result_structured = batch_chat_structured(
            chat, prompts, temp_file.name, CountryCapital, wait=False
        )
        assert result_structured is None
