from datetime import datetime
from typing import Any
from unittest.mock import AsyncMock

import pytest
from mistralai_workflows.client import WorkflowsClient
from mistralai_workflows.core.events.event_context import EventContext
from mistralai_workflows.testing import create_test_worker_with_events, wait_for_pending_inputs
from pydantic import ValidationError
from temporalio.testing import WorkflowEnvironment

from mistralai_workflows import InteractiveWorkflow, workflow
from mistralai_workflows.plugins.mistralai.lechat import (
    DateTimeField,
    FormInput,
    NumberField,
    SingleChoice,
    TextField,
)


class BasicTextForm(FormInput):
    name: str = TextField(description="Enter your name")


class TextFormWithPattern(FormInput):
    email: str = TextField(
        description="Email address",
        pattern=r"^[\w.-]+@[\w.-]+\.\w+$",
    )


class NumberFormWithBounds(FormInput):
    amount: float = NumberField(
        description="Amount in USD",
        minimum=0,
        maximum=10000,
    )


class DateTimeForm(FormInput):
    scheduled_at: datetime = DateTimeField(description="Schedule date and time")


class SingleChoiceForm(FormInput):
    priority: str = SingleChoice(
        options=[
            ("low", "Low Priority"),
            ("medium", "Medium Priority"),
            ("high", "High Priority"),
        ],
        description="Select priority level",
    )


class CompleteRFCForm(FormInput):
    title: str = TextField(description="RFC title", pattern=r"^RFC.*")
    summary: str = TextField(description="Brief summary")
    priority: str = SingleChoice(
        options=[
            ("low", "Low Priority"),
            ("high", "High Priority"),
        ],
        description="Priority level",
    )
    estimated_effort: float = NumberField(
        description="Effort in person-days",
        minimum=1,
        maximum=365,
    )


@workflow.define(name="basic_text_form_workflow")
class BasicTextFormWorkflow(InteractiveWorkflow):
    @workflow.entrypoint
    async def run(self) -> dict[str, Any]:
        form_input = await self.wait_for_input(BasicTextForm, label="Basic Text Form")
        return {"name": form_input.name}


@workflow.define(name="text_pattern_form_workflow")
class TextPatternFormWorkflow(InteractiveWorkflow):
    @workflow.entrypoint
    async def run(self) -> dict[str, Any]:
        form_input = await self.wait_for_input(TextFormWithPattern, label="Email Form")
        return {"email": form_input.email}


@workflow.define(name="number_form_workflow")
class NumberFormWorkflow(InteractiveWorkflow):
    @workflow.entrypoint
    async def run(self) -> dict[str, Any]:
        form_input = await self.wait_for_input(NumberFormWithBounds, label="Amount Form")
        return {"amount": form_input.amount}


@workflow.define(name="datetime_form_workflow")
class DateTimeFormWorkflow(InteractiveWorkflow):
    @workflow.entrypoint
    async def run(self) -> dict[str, Any]:
        form_input = await self.wait_for_input(DateTimeForm, label="Schedule Form")
        return {"scheduled_at": form_input.scheduled_at.isoformat()}


@workflow.define(name="single_choice_form_workflow")
class SingleChoiceFormWorkflow(InteractiveWorkflow):
    @workflow.entrypoint
    async def run(self) -> dict[str, Any]:
        form_input = await self.wait_for_input(SingleChoiceForm, label="Priority Form")
        return {"priority": form_input.priority}


@workflow.define(name="complete_rfc_form_workflow")
class CompleteRFCFormWorkflow(InteractiveWorkflow):
    @workflow.entrypoint
    async def run(self) -> dict[str, Any]:
        form_input = await self.wait_for_input(CompleteRFCForm, label="RFC Feedback Form")
        return {
            "title": form_input.title,
            "summary": form_input.summary,
            "priority": form_input.priority,
            "estimated_effort": form_input.estimated_effort,
        }


def extract_input_schema_from_events(captured_events: list[Any]) -> dict[str, Any] | None:
    for event in captured_events:
        if (
            hasattr(event, "event_type")
            and event.event_type.value == "CUSTOM_TASK_STARTED"
            and hasattr(event, "attributes")
            and event.attributes.custom_task_type == "wait_for_input"
        ):
            return event.attributes.payload.value.get("input_schema")
    return None


def extract_label_from_events(captured_events: list[Any]) -> str | None:
    for event in captured_events:
        if (
            hasattr(event, "event_type")
            and event.event_type.value == "CUSTOM_TASK_STARTED"
            and hasattr(event, "attributes")
            and event.attributes.custom_task_type == "wait_for_input"
        ):
            return event.attributes.payload.value.get("label")
    return None


@pytest.mark.asyncio
async def test_text_field_schema_in_event(temporal_env: WorkflowEnvironment) -> None:
    captured_events: list[Any] = []
    mock_client = AsyncMock(spec=WorkflowsClient)
    mock_client.send_event = AsyncMock(side_effect=lambda e: captured_events.append(e))

    async with EventContext(mock_client):
        async with create_test_worker_with_events(
            temporal_env,
            workflows=[BasicTextFormWorkflow],
        ):
            handle = await temporal_env.client.start_workflow(
                "basic_text_form_workflow",
                id="test-text-field-schema",
                task_queue="test-task-queue",
            )

            pending_inputs = await wait_for_pending_inputs(handle, expected_count=1)
            task_id = pending_inputs[0]["task_id"]

            await handle.execute_update(
                "__submit_input",
                {"task_id": task_id, "input": {"name": "John Doe"}},
            )
            await handle.result()

    input_schema = extract_input_schema_from_events(captured_events)
    assert input_schema is not None

    assert input_schema["type"] == "object"
    assert "properties" in input_schema

    name_schema = input_schema["properties"]["name"]
    assert name_schema["type"] == "string"
    assert name_schema["description"] == "Enter your name"


@pytest.mark.asyncio
async def test_text_field_with_pattern_in_event(temporal_env: WorkflowEnvironment) -> None:
    captured_events: list[Any] = []
    mock_client = AsyncMock(spec=WorkflowsClient)
    mock_client.send_event = AsyncMock(side_effect=lambda e: captured_events.append(e))

    async with EventContext(mock_client):
        async with create_test_worker_with_events(
            temporal_env,
            workflows=[TextPatternFormWorkflow],
        ):
            handle = await temporal_env.client.start_workflow(
                "text_pattern_form_workflow",
                id="test-text-pattern-schema",
                task_queue="test-task-queue",
            )

            pending_inputs = await wait_for_pending_inputs(handle, expected_count=1)
            task_id = pending_inputs[0]["task_id"]

            await handle.execute_update(
                "__submit_input",
                {"task_id": task_id, "input": {"email": "test@example.com"}},
            )
            await handle.result()

    input_schema = extract_input_schema_from_events(captured_events)
    assert input_schema is not None

    email_schema = input_schema["properties"]["email"]
    assert email_schema["type"] == "string"
    assert email_schema["description"] == "Email address"
    assert email_schema["pattern"] == r"^[\w.-]+@[\w.-]+\.\w+$"


@pytest.mark.asyncio
async def test_number_field_with_bounds_in_event(temporal_env: WorkflowEnvironment) -> None:
    captured_events: list[Any] = []
    mock_client = AsyncMock(spec=WorkflowsClient)
    mock_client.send_event = AsyncMock(side_effect=lambda e: captured_events.append(e))

    async with EventContext(mock_client):
        async with create_test_worker_with_events(
            temporal_env,
            workflows=[NumberFormWorkflow],
        ):
            handle = await temporal_env.client.start_workflow(
                "number_form_workflow",
                id="test-number-field-schema",
                task_queue="test-task-queue",
            )

            pending_inputs = await wait_for_pending_inputs(handle, expected_count=1)
            task_id = pending_inputs[0]["task_id"]

            await handle.execute_update(
                "__submit_input",
                {"task_id": task_id, "input": {"amount": 500.0}},
            )
            await handle.result()

    input_schema = extract_input_schema_from_events(captured_events)
    assert input_schema is not None

    amount_schema = input_schema["properties"]["amount"]
    assert amount_schema["type"] == "number"
    assert amount_schema["description"] == "Amount in USD"
    assert amount_schema["minimum"] == 0
    assert amount_schema["maximum"] == 10000


@pytest.mark.asyncio
async def test_datetime_field_in_event(temporal_env: WorkflowEnvironment) -> None:
    captured_events: list[Any] = []
    mock_client = AsyncMock(spec=WorkflowsClient)
    mock_client.send_event = AsyncMock(side_effect=lambda e: captured_events.append(e))

    async with EventContext(mock_client):
        async with create_test_worker_with_events(
            temporal_env,
            workflows=[DateTimeFormWorkflow],
        ):
            handle = await temporal_env.client.start_workflow(
                "datetime_form_workflow",
                id="test-datetime-field-schema",
                task_queue="test-task-queue",
            )

            pending_inputs = await wait_for_pending_inputs(handle, expected_count=1)
            task_id = pending_inputs[0]["task_id"]

            await handle.execute_update(
                "__submit_input",
                {"task_id": task_id, "input": {"scheduled_at": "2025-01-15T10:30:00"}},
            )
            await handle.result()

    input_schema = extract_input_schema_from_events(captured_events)
    assert input_schema is not None

    scheduled_schema = input_schema["properties"]["scheduled_at"]
    assert scheduled_schema["type"] == "string"
    assert scheduled_schema["format"] == "date-time"
    assert scheduled_schema["description"] == "Schedule date and time"


@pytest.mark.asyncio
async def test_single_choice_field_in_event(temporal_env: WorkflowEnvironment) -> None:
    """Test that SingleChoice generates oneOf schema with const values."""
    captured_events: list[Any] = []
    mock_client = AsyncMock(spec=WorkflowsClient)
    mock_client.send_event = AsyncMock(side_effect=lambda e: captured_events.append(e))

    async with EventContext(mock_client):
        async with create_test_worker_with_events(
            temporal_env,
            workflows=[SingleChoiceFormWorkflow],
        ):
            handle = await temporal_env.client.start_workflow(
                "single_choice_form_workflow",
                id="test-single-choice-schema",
                task_queue="test-task-queue",
            )

            pending_inputs = await wait_for_pending_inputs(handle, expected_count=1)
            task_id = pending_inputs[0]["task_id"]

            await handle.execute_update(
                "__submit_input",
                {"task_id": task_id, "input": {"priority": "high"}},
            )
            await handle.result()

    input_schema = extract_input_schema_from_events(captured_events)
    assert input_schema is not None

    priority_schema = input_schema["properties"]["priority"]

    assert "oneOf" in priority_schema
    assert "type" not in priority_schema
    assert priority_schema["description"] == "Select priority level"

    one_of = priority_schema["oneOf"]
    assert len(one_of) == 3

    assert one_of[0] == {"type": "string", "const": "low", "description": "Low Priority"}
    assert one_of[1] == {"type": "string", "const": "medium", "description": "Medium Priority"}
    assert one_of[2] == {"type": "string", "const": "high", "description": "High Priority"}


@pytest.mark.asyncio
async def test_complete_rfc_form_schema_in_event(temporal_env: WorkflowEnvironment) -> None:
    captured_events: list[Any] = []
    mock_client = AsyncMock(spec=WorkflowsClient)
    mock_client.send_event = AsyncMock(side_effect=lambda e: captured_events.append(e))

    async with EventContext(mock_client):
        async with create_test_worker_with_events(
            temporal_env,
            workflows=[CompleteRFCFormWorkflow],
        ):
            handle = await temporal_env.client.start_workflow(
                "complete_rfc_form_workflow",
                id="test-complete-form-schema",
                task_queue="test-task-queue",
            )

            pending_inputs = await wait_for_pending_inputs(handle, expected_count=1)
            task_id = pending_inputs[0]["task_id"]

            await handle.execute_update(
                "__submit_input",
                {
                    "task_id": task_id,
                    "input": {
                        "title": "RFC: New Feature",
                        "summary": "A brief summary",
                        "priority": "high",
                        "estimated_effort": 10.0,
                    },
                },
            )
            await handle.result()

    input_schema = extract_input_schema_from_events(captured_events)
    assert input_schema is not None

    assert input_schema["type"] == "object"
    assert input_schema["title"] == "FormInput"

    props = input_schema["properties"]

    assert props["title"]["type"] == "string"
    assert props["title"]["description"] == "RFC title"
    assert props["title"]["pattern"] == r"^RFC.*"

    assert props["summary"]["type"] == "string"
    assert props["summary"]["description"] == "Brief summary"

    assert "oneOf" in props["priority"]
    assert len(props["priority"]["oneOf"]) == 2
    assert props["priority"]["description"] == "Priority level"

    assert props["estimated_effort"]["type"] == "number"
    assert props["estimated_effort"]["minimum"] == 1
    assert props["estimated_effort"]["maximum"] == 365


@pytest.mark.asyncio
async def test_form_input_label_in_event(temporal_env: WorkflowEnvironment) -> None:
    captured_events: list[Any] = []
    mock_client = AsyncMock(spec=WorkflowsClient)
    mock_client.send_event = AsyncMock(side_effect=lambda e: captured_events.append(e))

    async with EventContext(mock_client):
        async with create_test_worker_with_events(
            temporal_env,
            workflows=[CompleteRFCFormWorkflow],
        ):
            handle = await temporal_env.client.start_workflow(
                "complete_rfc_form_workflow",
                id="test-form-label",
                task_queue="test-task-queue",
            )

            pending_inputs = await wait_for_pending_inputs(handle, expected_count=1)
            task_id = pending_inputs[0]["task_id"]

            await handle.execute_update(
                "__submit_input",
                {
                    "task_id": task_id,
                    "input": {
                        "title": "RFC: Test",
                        "summary": "Test",
                        "priority": "low",
                        "estimated_effort": 5.0,
                    },
                },
            )
            await handle.result()

    label = extract_label_from_events(captured_events)
    assert label == "RFC Feedback Form"


def test_text_field_pattern_validation() -> None:
    class PatternForm(FormInput):
        email: str = TextField(description="Email", pattern=r"^[\w.-]+@[\w.-]+\.\w+$")

    # Valid input should work
    valid = PatternForm(email="test@example.com")
    assert valid.email == "test@example.com"

    # Invalid input should raise ValidationError
    with pytest.raises(ValidationError):
        PatternForm(email="not-an-email")


def test_number_field_minimum_validation() -> None:
    class MinForm(FormInput):
        amount: float = NumberField(description="Amount", minimum=0)

    # Valid input at boundary
    valid = MinForm(amount=0)
    assert valid.amount == 0

    # Valid input above minimum
    valid = MinForm(amount=100)
    assert valid.amount == 100

    # Invalid input below minimum
    with pytest.raises(ValidationError):
        MinForm(amount=-1)


def test_number_field_maximum_validation() -> None:
    class MaxForm(FormInput):
        amount: float = NumberField(description="Amount", maximum=100)

    # Valid input at boundary
    valid = MaxForm(amount=100)
    assert valid.amount == 100

    # Valid input below maximum
    valid = MaxForm(amount=50)
    assert valid.amount == 50

    # Invalid input above maximum
    with pytest.raises(ValidationError):
        MaxForm(amount=101)


def test_number_field_exclusive_bounds_validation() -> None:
    class ExclusiveForm(FormInput):
        value: float = NumberField(
            description="Value",
            exclusive_minimum=0,
            exclusive_maximum=100,
        )

    valid = ExclusiveForm(value=50)
    assert valid.value == 50

    with pytest.raises(ValidationError):
        ExclusiveForm(value=0)

    with pytest.raises(ValidationError):
        ExclusiveForm(value=100)


def test_single_choice_valid_option_validation() -> None:
    class ChoiceForm(FormInput):
        priority: str = SingleChoice(
            options=[("low", "Low"), ("medium", "Medium"), ("high", "High")],
            description="Priority",
        )

    low = ChoiceForm(priority="low")
    assert low.priority == "low"

    medium = ChoiceForm(priority="medium")
    assert medium.priority == "medium"

    high = ChoiceForm(priority="high")
    assert high.priority == "high"


def test_single_choice_invalid_option_validation() -> None:
    class ChoiceForm(FormInput):
        priority: str = SingleChoice(
            options=[("low", "Low"), ("medium", "Medium"), ("high", "High")],
            description="Priority",
        )

    with pytest.raises(ValidationError):
        ChoiceForm(priority="invalid")


def test_single_choice_simple_strings_validation() -> None:
    class SimpleChoiceForm(FormInput):
        color: str = SingleChoice(
            options=["red", "green", "blue"],
            description="Color",
        )

    valid = SimpleChoiceForm(color="red")
    assert valid.color == "red"

    with pytest.raises(ValidationError):
        SimpleChoiceForm(color="yellow")
