"""Mistral provider implementation."""

from __future__ import annotations

import asyncio
import time
from typing import Any

from pi_agents.ai.env_api_keys import get_env_api_key
from pi_agents.ai.models import calculate_cost
from pi_agents.ai.types import (
    AssistantMessage,
    Context,
    Cost,
    Model,
    SimpleStreamOptions,
    TextContent,
    ToolCall,
    Usage,
)
from pi_agents.ai.utils.event_stream import AssistantMessageEventStream
from pi_agents.ai.utils.json_parse import parse_streaming_json


def stream_mistral(
    model: Model,
    context: Context,
    options: dict[str, Any] | None = None,
) -> AssistantMessageEventStream:
    """Stream responses from the Mistral API."""
    event_stream = AssistantMessageEventStream()
    opts = options or {}

    async def _run():
        try:
            from mistralai import Mistral

            api_key = opts.get("api_key") or opts.get("apiKey") or get_env_api_key("mistral")
            if not api_key:
                raise ValueError("Mistral API key not found. Set MISTRAL_API_KEY environment variable.")

            client = Mistral(api_key=api_key)

            # Build messages
            messages: list[dict[str, Any]] = []
            if context.system_prompt:
                messages.append({"role": "system", "content": context.system_prompt})

            for msg in context.messages:
                role = msg.role if hasattr(msg, "role") else msg.get("role", "user")
                if role == "user":
                    content = msg.content if hasattr(msg, "content") else msg.get("content", "")
                    messages.append({"role": "user", "content": content if isinstance(content, str) else str(content)})
                elif role == "assistant":
                    content_list = msg.content if hasattr(msg, "content") else msg.get("content", [])
                    text = ""
                    for block in content_list:
                        btype = block.type if hasattr(block, "type") else block.get("type", "text")
                        if btype == "text":
                            text += block.text if hasattr(block, "text") else block.get("text", "")
                    messages.append({"role": "assistant", "content": text})

            params: dict[str, Any] = {
                "model": model.id,
                "messages": messages,
            }

            max_tokens = opts.get("max_tokens") or opts.get("maxTokens") or model.max_tokens
            if max_tokens:
                params["max_tokens"] = max_tokens

            if opts.get("temperature") is not None:
                params["temperature"] = opts["temperature"]

            partial = AssistantMessage(
                role="assistant",
                content=[],
                api=model.api,
                provider=model.provider,
                model=model.id,
                usage=Usage(cost=Cost()),
                stop_reason="stop",
                timestamp=int(time.time() * 1000),
            )

            event_stream.push({"type": "start", "partial": partial})

            text_buffer = ""
            content_index = 0

            stream = await client.chat.stream_async(**params)
            async for chunk in stream:
                data = chunk.data
                if data.choices:
                    choice = data.choices[0]
                    delta = choice.delta
                    if delta and delta.content:
                        if not text_buffer:
                            event_stream.push({"type": "text_start", "contentIndex": content_index, "partial": partial})
                        text_buffer += delta.content
                        event_stream.push({
                            "type": "text_delta",
                            "contentIndex": content_index,
                            "delta": delta.content,
                            "partial": partial,
                        })

                if data.usage:
                    partial.usage.input = data.usage.prompt_tokens or 0
                    partial.usage.output = data.usage.completion_tokens or 0
                    partial.usage.total_tokens = data.usage.total_tokens or 0

            if text_buffer:
                event_stream.push({
                    "type": "text_end",
                    "contentIndex": content_index,
                    "content": text_buffer,
                    "partial": partial,
                })
                partial.content.append(TextContent(type="text", text=text_buffer))

            calculate_cost(model, partial.usage)
            event_stream.push({"type": "done", "reason": "stop", "message": partial})

        except Exception as e:
            error_msg = AssistantMessage(
                role="assistant",
                content=[],
                api=model.api,
                provider=model.provider,
                model=model.id,
                usage=Usage(cost=Cost()),
                stop_reason="error",
                error_message=str(e),
                timestamp=int(time.time() * 1000),
            )
            event_stream.push({"type": "error", "reason": "error", "error": error_msg})

        event_stream.end()

    asyncio.ensure_future(_run())
    return event_stream


def stream_simple_mistral(
    model: Model,
    context: Context,
    options: SimpleStreamOptions | None = None,
) -> AssistantMessageEventStream:
    """Stream with simplified options."""
    opts: dict[str, Any] = {}
    if options:
        if options.temperature is not None:
            opts["temperature"] = options.temperature
        if options.max_tokens is not None:
            opts["maxTokens"] = options.max_tokens
        if options.api_key:
            opts["apiKey"] = options.api_key
    return stream_mistral(model, context, opts)
