from __future__ import annotations

from typing import TypeVar

import structlog
from pydantic import AliasChoices, BaseModel, Field

import mistralai
from mistralai import ResponseFormat, models
from mistralai.extra import response_format_from_pydantic_model
from mistralai_workflows.core.activity import activity
from mistralai_workflows.core.dependencies.dependency_injector import Depends
from mistralai_workflows.plugins.mistralai.models import AgentUpdateRequest, ConversationAppendRequest
from mistralai_workflows.plugins.mistralai.utils import (
    _get_agent_llm_rate_limit,
    get_mistral_client,
    handle_chat_stream,
    handle_conversation_stream,
)

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

logger = structlog.get_logger(__name__)


@activity(rate_limit=_get_agent_llm_rate_limit())
async def mistralai_create_agent(
    params: mistralai.AgentCreationRequest,
    mistral_client: mistralai.Mistral = Depends(get_mistral_client),
) -> mistralai.Agent:
    agent = mistral_client.beta.agents.create(**params.model_dump(by_alias=True))
    return agent


@activity(rate_limit=_get_agent_llm_rate_limit())
async def mistralai_start_conversation(
    params: mistralai.ConversationRequest,
    mistral_client: mistralai.Mistral = Depends(get_mistral_client),
) -> mistralai.ConversationResponse:
    return await mistral_client.beta.conversations.start_async(**params.model_dump(by_alias=True))


@activity(rate_limit=_get_agent_llm_rate_limit())
async def mistralai_append_conversation(
    params: ConversationAppendRequest,
    mistral_client: mistralai.Mistral = Depends(get_mistral_client),
) -> mistralai.ConversationResponse:
    return await mistral_client.beta.conversations.append_async(**params.model_dump(by_alias=True))


@activity(rate_limit=_get_agent_llm_rate_limit())
async def mistralai_update_agent(
    params: AgentUpdateRequest,
    mistral_client: mistralai.Mistral = Depends(get_mistral_client),
) -> mistralai.Agent:
    return await mistral_client.beta.agents.update_async(**params.model_dump(by_alias=True))


@activity(rate_limit=_get_agent_llm_rate_limit())
async def mistralai_chat_complete(
    params: mistralai.ChatCompletionRequest,
    mistral_client: mistralai.Mistral = Depends(get_mistral_client),
) -> mistralai.ChatCompletionResponse:
    return await mistral_client.chat.complete_async(**params.model_dump())


@activity(rate_limit=_get_agent_llm_rate_limit())
async def mistralai_chat_stream(
    params: mistralai.ChatCompletionRequest,
    mistral_client: mistralai.Mistral = Depends(get_mistral_client),
) -> mistralai.AssistantMessage:
    payload = params.model_copy(update={"stream": True})
    stream = await mistral_client.chat.stream_async(**payload.model_dump(by_alias=True, exclude_none=True))
    return await handle_chat_stream(stream=stream)


@activity(rate_limit=_get_agent_llm_rate_limit())
async def mistralai_start_conversation_stream(
    params: mistralai.ConversationRequest, mistral_client: mistralai.Mistral = Depends(get_mistral_client)
) -> mistralai.ConversationResponse:
    payload = params.model_copy(update={"stream": True})

    stream = await mistral_client.beta.conversations.start_stream_async(
        **payload.model_dump(by_alias=True, exclude_none=True)
    )

    return await handle_conversation_stream(stream=stream)


@activity(rate_limit=_get_agent_llm_rate_limit())
async def mistralai_append_conversation_stream(
    params: ConversationAppendRequest, mistral_client: mistralai.Mistral = Depends(get_mistral_client)
) -> mistralai.ConversationResponse:
    payload = params.model_copy(update={"stream": True})

    stream = await mistral_client.beta.conversations.append_stream_async(
        **payload.model_dump(by_alias=True, exclude_none=True)
    )

    return await handle_conversation_stream(stream=stream)


@activity(rate_limit=_get_agent_llm_rate_limit())
async def mistralai_chat_parse(
    params: mistralai.ChatCompletionRequest,
    response_format: ResponseFormat,
    mistral_client: mistralai.Mistral = Depends(get_mistral_client),
) -> mistralai.ChatCompletionResponse:
    """
    Chat completion with structured output parsing.

    This activity performs chat completion with structured output parsing.
    We use complete_async because Temporal doesn't support serializing Type[PydanticModel].

    Args:
        params: Chat completion request parameters
        response_format: Output response format
        mistral_client: Mistral client

    Returns:
        ChatCompletionResponse with parsed output in choices[0].message.content (as JSON string)
    """
    params_with_format = params.model_copy(update={"response_format": response_format})
    return await mistral_client.chat.complete_async(**params_with_format.model_dump(by_alias=True))


async def chat_parse_to_model(
    model_class: type[T],
    request: mistralai.ChatCompletionRequest,
) -> T:
    """
    Parse chat completion response into a Pydantic model.

    Wraps mistralai_chat_parse activity with schema conversion and response validation.
    """
    response_format = response_format_from_pydantic_model(model_class)
    response = await mistralai_chat_parse(request, response_format)

    if not response.choices or not response.choices[0].message or not response.choices[0].message.content:
        raise ValueError("No parsed response from Mistral")

    return model_class.model_validate_json(str(response.choices[0].message.content))


class MistralEmbeddingsParams(mistralai.EmbeddingRequest):
    # without this, the current `mistralai.EmbeddingRequest` doesn't keep the
    # `input` and `inputs` fields during model_dump(). This fixes it.
    # TODO(WFL-451): remove this once the `mistralai` SDK is updated.
    inputs: mistralai.EmbeddingRequestInputs = Field(validation_alias=AliasChoices("inputs", "input"))


@activity(rate_limit=_get_agent_llm_rate_limit())
async def mistralai_embeddings(
    params: MistralEmbeddingsParams,
    mistral_client: mistralai.Mistral = Depends(get_mistral_client),
) -> models.EmbeddingResponse:
    return await mistral_client.embeddings.create_async(**params.model_dump(by_alias=True))


@activity(rate_limit=_get_agent_llm_rate_limit())
async def mistralai_ocr(
    params: mistralai.OCRRequest,
    mistral_client: mistralai.Mistral = Depends(get_mistral_client),
) -> models.OCRResponse:
    return await mistral_client.ocr.process_async(**params.model_dump(by_alias=True))
