import uuid
from typing import Any, TypeVar

import temporalio.workflow
from pydantic import BaseModel, ValidationError

from mistralai_workflows.core.task import task
from mistralai_workflows.core.workflow import workflow

T = TypeVar("T", bound=BaseModel)


class _WaitingForInputTaskState(BaseModel):
    input_schema: dict
    label: str | None


class _SubmitInputParams(BaseModel):
    task_id: str
    input: Any


class _SubmitInputResult(BaseModel):
    error: str | None = None


class _PendingInputRequest(BaseModel):
    task_id: str
    input_schema: type[BaseModel]
    input_schema_dict: dict
    label: str | None = None
    has_received_input: bool = False
    input: Any | None = None


class _PendingInputInfo(BaseModel):
    task_id: str
    input_schema: dict
    label: str | None = None


class _PendingInputsResult(BaseModel):
    pending_inputs: list[_PendingInputInfo]


class InteractiveWorkflow:
    """
    Base class for workflows that need to wait for external input.

    Workflows inheriting from this class gain the ability to pause execution
    and wait for external input via the wait_for_input() method.

    Example usage:
        @workflow.define(name="my-workflow")
        class MyWorkflow(InteractiveWorkflow):
            @workflow.entrypoint
            async def run(self, params: MyInput):
                # Ask human for approval
                approval = await self.wait_for_input(ApprovalSchema)
                if approval.approved:
                    return "Approved!"
                return "Denied"

    How to submit input (from external client):
        await handle.execute_update(
            "__submit_input",
            {"task_id": "task-456", "input": {"approved": True}}
        )
    """

    def __init__(self) -> None:
        # In-memory storage: tracks which tasks are waiting for input
        self._pending_inputs: dict[str, _PendingInputRequest] = {}

    async def wait_for_input(self, schema: type[T], label: str | None = None) -> T:
        """
        Pause workflow and wait for external input matching the provided schema.

        The workflow suspends execution until input is submitted via the workflow
        update handler. The task becomes visible in streaming events, allowing
        external systems to detect and respond to the input request.

        Args:
            schema: Pydantic model class defining the expected input structure.
                    Input will be validated against this schema.
            label: Optional human-readable label to identify this input request.
                   Useful when multiple inputs are pending to help clients distinguish
                   between them (e.g., "Manager Approval", "File Upload", "Review Step 1").

        Returns:
            Validated input data as an instance of the schema class.

        Raises:
            ValidationError: If submitted input fails validation against the schema.

        Example:
            ```python
            from pydantic import BaseModel
            import workflows

            class ApprovalRequest(BaseModel):
                approved: bool
                reviewer: str
                comments: str

            @workflows.define()
            class ApprovalWorkflow(InteractiveWorkflow):
                @workflows.entrypoint
                async def run(self, document_id: str) -> str:
                    # Workflow pauses here until input is submitted
                    approval = await self.wait_for_input(ApprovalRequest)

                    if approval.approved:
                        return f"Document approved by {approval.reviewer}"
                    return f"Document rejected: {approval.comments}"
            ```

        Submitting Input (Client Side):
            ```python
            # Get pending inputs with metadata
            result = await handle.query("__get_pending_inputs")
            # result["pending_inputs"] = [
            #   {
            #     "task_id": "abc-123",
            #     "label": "Manager Approval",
            #     "input_schema": {"type": "object", "properties": {...}}
            #   }
            # ]

            # Find the specific task you want to respond to
            approval_task = next(
                inp for inp in result["pending_inputs"]
                if inp["label"] == "Manager Approval"
            )

            # Submit input for that specific task
            await handle.execute_update(
                "__submit_input",
                {
                    "task_id": approval_task["task_id"],
                    "input": {"approved": True, "reviewer": "Alice", "comments": "LGTM"}
                }
            )
            ```
        """
        task_id = str(temporalio.workflow.uuid4() if temporalio.workflow.in_workflow() else uuid.uuid4())
        schema_dict = schema.model_json_schema()

        pending_request = _PendingInputRequest(
            task_id=task_id,
            input_schema=schema,
            input_schema_dict=schema_dict,
            label=label,
        )

        self._pending_inputs[task_id] = pending_request

        with task("wait_for_input", state=_WaitingForInputTaskState(input_schema=schema_dict, label=label), id=task_id):

            def check_received() -> bool:
                return self._pending_inputs.get(task_id, pending_request).has_received_input

            try:
                await temporalio.workflow.wait_condition(check_received)
                validated_input = schema.model_validate(pending_request.input)
                return validated_input
            finally:
                self._pending_inputs.pop(task_id, None)

    @workflow.update(name="__submit_input", _internal=True)
    async def _handle_input_submission(self, params: _SubmitInputParams) -> _SubmitInputResult:
        """
        Called by external API when external input is submitted.

        Flow:
        1. External client calls: client.update_workflow(execution_id, "__submit_input", {...})
        2. Temporal routes update to this handler
        3. This handler sets: has_received_input = True
        4. Setting that flag is an EVENT → Temporal re-evaluates wait_condition()
        5. wait_condition() sees flag is True → workflow resumes!

        Args:
            params: Contains task_id and input data from external source

        Returns:
            Result with error message if validation fails
        """
        pending_request = self._pending_inputs.get(params.task_id)
        if not pending_request:
            return _SubmitInputResult(error=f"No pending input request found for task {params.task_id}")

        try:
            pending_request.input_schema.model_validate(params.input)
        except ValidationError as e:
            return _SubmitInputResult(error=f"Invalid input for task {params.task_id}: {str(e)}")

        pending_request.input = params.input
        pending_request.has_received_input = True
        return _SubmitInputResult()

    @workflow.query(name="__get_pending_inputs", _internal=True)
    def _get_pending_inputs(self) -> _PendingInputsResult:
        """
        Query handler to get all pending input requests.

        Returns:
            List of pending input requests with their schemas and labels.
        """
        pending_input_list = [
            _PendingInputInfo(
                task_id=task_id,
                input_schema=request.input_schema_dict,
                label=request.label,
            )
            for task_id, request in self._pending_inputs.items()
        ]
        return _PendingInputsResult(pending_inputs=pending_input_list)
