import os

import pytest
from pydantic import BaseModel

from mistralai.workflows.client import WorkflowsClient
from mistralai.workflows.constants import MAX_INPUT_SIZE_BYTES
from mistralai.workflows.exceptions import ErrorCode, WorkflowsException
from mistralai.workflows.testing import (
    WORKFLOW_EXAMPLE_HELLO_WORLD,
    WORKFLOW_EXAMPLE_INTERACTIVE_GAME,
    WORKFLOW_SIMPLE_CHATBOT,
)


class LargeInput(BaseModel):
    data: str


class HelloWorldInput(BaseModel):
    document_title: str


class InteractiveGameInput(BaseModel):
    initial_player_name: str
    max_riddles: int


class SignalInput(BaseModel):
    answer: str


class UpdateInput(BaseModel):
    role: str
    content: str


def _generate_large_string(size_bytes: int) -> str:
    return "x" * size_bytes


@pytest.fixture
def workflows_client() -> WorkflowsClient:
    server_url = os.getenv("SERVER_URL", "http://localhost:7444")
    api_key = os.getenv("MISTRAL_API_KEY")
    return WorkflowsClient(base_url=server_url, api_key=api_key)


@pytest.mark.integration
@pytest.mark.asyncio
async def test_execute_rejects_input_exceeding_limit(workflows_client: WorkflowsClient) -> None:
    async with workflows_client:
        large_input = LargeInput(data=_generate_large_string(MAX_INPUT_SIZE_BYTES + 1000))

        with pytest.raises(WorkflowsException) as exc_info:
            await workflows_client.execute_workflow(
                WORKFLOW_EXAMPLE_HELLO_WORLD,
                input_data=large_input,
            )

        assert exc_info.value.code == ErrorCode.INPUT_SIZE_EXCEEDED_ERROR
        assert exc_info.value.status.value == 413


@pytest.mark.integration
@pytest.mark.asyncio
async def test_execute_accepts_input_at_limit(workflows_client: WorkflowsClient) -> None:
    async with workflows_client:
        result = await workflows_client.execute_workflow(
            WORKFLOW_EXAMPLE_HELLO_WORLD,
            input_data=HelloWorldInput(document_title="test"),
        )

        assert result.execution_id is not None


@pytest.mark.integration
@pytest.mark.asyncio
async def test_signal_rejects_input_exceeding_limit(workflows_client: WorkflowsClient) -> None:
    async with workflows_client:
        result = await workflows_client.execute_workflow(
            WORKFLOW_EXAMPLE_INTERACTIVE_GAME,
            input_data=InteractiveGameInput(initial_player_name="test", max_riddles=1),
        )
        execution_id = result.execution_id

        large_input = LargeInput(data=_generate_large_string(MAX_INPUT_SIZE_BYTES + 1000))

        with pytest.raises(WorkflowsException) as exc_info:
            await workflows_client.signal_workflow(
                execution_id,
                signal_name="submit_answer",
                input_data=large_input,
            )

        assert exc_info.value.code == ErrorCode.INPUT_SIZE_EXCEEDED_ERROR
        assert exc_info.value.status.value == 413


@pytest.mark.integration
@pytest.mark.asyncio
async def test_query_rejects_input_exceeding_limit(workflows_client: WorkflowsClient) -> None:
    async with workflows_client:
        result = await workflows_client.execute_workflow(
            WORKFLOW_EXAMPLE_INTERACTIVE_GAME,
            input_data=InteractiveGameInput(initial_player_name="test", max_riddles=1),
        )
        execution_id = result.execution_id

        large_input = LargeInput(data=_generate_large_string(MAX_INPUT_SIZE_BYTES + 1000))

        with pytest.raises(WorkflowsException) as exc_info:
            await workflows_client.query_workflow(
                execution_id,
                query_name="get_status",
                input_data=large_input,
            )

        assert exc_info.value.code == ErrorCode.INPUT_SIZE_EXCEEDED_ERROR
        assert exc_info.value.status.value == 413


@pytest.mark.integration
@pytest.mark.asyncio
async def test_update_rejects_input_exceeding_limit(workflows_client: WorkflowsClient) -> None:
    async with workflows_client:
        result = await workflows_client.execute_workflow(
            WORKFLOW_SIMPLE_CHATBOT,
            input_data=None,
        )
        execution_id = result.execution_id

        large_content = _generate_large_string(MAX_INPUT_SIZE_BYTES + 1000)
        large_input = UpdateInput(role="user", content=large_content)

        with pytest.raises(WorkflowsException) as exc_info:
            await workflows_client.update_workflow(
                execution_id,
                update_name="discuss",
                input_data=large_input,
            )

        assert exc_info.value.code == ErrorCode.INPUT_SIZE_EXCEEDED_ERROR
        assert exc_info.value.status.value == 413
