"""
Mistral AI provider integration for Amplitude AI SDK.

Wraps the ``mistralai`` Python client to automatically track LLM
usage in Amplitude.

Usage::

    from amplitude import Amplitude
    from amplitude_ai.providers.mistral import Mistral

    amplitude = Amplitude("YOUR_AMP_KEY")
    client = Mistral(
        amplitude=amplitude,
        api_key="MISTRAL_KEY",
    )

    response = client.chat.complete(
        model="mistral-large-latest",
        messages=[{"role": "user", "content": "Hello!"}],
        amplitude_user_id="user-1",
    )
"""
from __future__ import annotations

import time
import uuid
from typing import Any, Dict, Iterator, List, Optional

try:
    from mistralai import Mistral as MistralClient

    MISTRAL_AVAILABLE = True
except ImportError:
    MISTRAL_AVAILABLE = False
    MistralClient = None  # type: ignore[assignment,misc]

from amplitude import Amplitude

from ..core.privacy import PrivacyConfig
from ..core.tracking import track_ai_message, track_user_message
from ..exceptions import ProviderError
from ..utils.costs import calculate_cost
from ..utils.logger import get_logger
from .base import _apply_session_context, AmplitudeOrAI, resolve_amplitude


class Mistral:
    """
    Mistral AI client wrapper with automatic Amplitude tracking.

    Uses composition (rather than inheritance) because the ``mistralai``
    SDK's client layout varies across versions.

    .. todo:: Refactor to inherit from :class:`BaseAIProvider` for shared
       tracking logic (privacy config, metadata helpers, etc.).

    The wrapper exposes a
    ``chat`` namespace with a ``complete`` / ``stream`` interface that
    mirrors the upstream SDK.
    """

    def __init__(
        self,
        amplitude: AmplitudeOrAI,
        *,
        api_key: Optional[str] = None,
        privacy_config: Optional[PrivacyConfig] = None,
        privacy_mode: bool = False,
        content_mode: Optional[str] = None,
        **mistral_kwargs: Any,
    ) -> None:
        if not MISTRAL_AVAILABLE:
            raise ProviderError(
                "Mistral SDK not installed. Install with: pip install mistralai"
            )

        self.amplitude = resolve_amplitude(amplitude)
        self.privacy_config = privacy_config or PrivacyConfig(
            privacy_mode=privacy_mode, content_mode=content_mode
        )

        init_kwargs: Dict[str, Any] = {**mistral_kwargs}
        if api_key is not None:
            init_kwargs["api_key"] = api_key

        self._client = MistralClient(**init_kwargs)  # type: ignore[misc]
        self.chat = _WrappedChat(self)

    # ------------------------------------------------------------------
    # Passthrough
    # ------------------------------------------------------------------

    @property
    def client(self) -> Any:
        """Access the underlying ``mistralai.Mistral`` client."""
        return self._client

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _extract_response_content(response: Any) -> str:
        """Extract the final text response from a Mistral completion.

        Handles both simple string ``content`` and the list-of-chunks
        format used by magistral reasoning models (``type="text"``
        chunks only).
        """
        if not (hasattr(response, "choices") and response.choices):
            return ""
        msg = getattr(response.choices[0], "message", None)
        if not msg or not hasattr(msg, "content"):
            return ""
        content = msg.content
        if isinstance(content, str):
            return content
        if isinstance(content, list):
            # Magistral models: list of typed chunks
            return "".join(
                chunk.get("text", "")
                for chunk in content
                if isinstance(chunk, dict) and chunk.get("type") == "text"
            )
        return str(content) if content else ""

    @staticmethod
    def _extract_reasoning_content(response: Any) -> Optional[str]:
        """Extract reasoning/thinking content from a Mistral completion.

        Magistral models return content as a list of typed chunks.
        Chunks with ``type="thinking"`` contain the reasoning text
        under ``thinking[].text``.
        """
        if not (hasattr(response, "choices") and response.choices):
            return None
        msg = getattr(response.choices[0], "message", None)
        if not msg or not hasattr(msg, "content"):
            return None
        content = msg.content
        if not isinstance(content, list):
            return None
        parts: List[str] = []
        for chunk in content:
            if not isinstance(chunk, dict) or chunk.get("type") != "thinking":
                continue
            for t in chunk.get("thinking", []):
                if isinstance(t, dict) and t.get("type") == "text":
                    parts.append(t.get("text", ""))
        return "\n".join(parts) if parts else None

    @staticmethod
    def _extract_finish_reason(response: Any) -> Optional[str]:
        if hasattr(response, "choices") and response.choices:
            return getattr(response.choices[0], "finish_reason", None)
        return None

    @staticmethod
    def _extract_tool_calls(response: Any) -> Optional[List[Dict[str, Any]]]:
        if hasattr(response, "choices") and response.choices:
            msg = getattr(response.choices[0], "message", None)
            if msg and hasattr(msg, "tool_calls") and msg.tool_calls:
                return [
                    {
                        "type": getattr(tc, "type", "function"),
                        "id": getattr(tc, "id", ""),
                        "function": {
                            "name": getattr(tc.function, "name", ""),
                            "arguments": getattr(tc.function, "arguments", ""),
                        },
                    }
                    for tc in msg.tool_calls
                    if hasattr(tc, "function")
                ]
        return None

    @staticmethod
    def _extract_usage(response: Any) -> Dict[str, Optional[int]]:
        result: Dict[str, Optional[int]] = {
            "input_tokens": None,
            "output_tokens": None,
            "total_tokens": None,
        }
        usage = getattr(response, "usage", None)
        if usage is None:
            return result
        result["input_tokens"] = getattr(usage, "prompt_tokens", None)
        result["output_tokens"] = getattr(usage, "completion_tokens", None)
        result["total_tokens"] = getattr(usage, "total_tokens", None)
        if (
            result["total_tokens"] is None
            and result["input_tokens"] is not None
            and result["output_tokens"] is not None
        ):
            result["total_tokens"] = result["input_tokens"] + result["output_tokens"]
        return result

    def _extract_system_prompt(self, messages: List[Dict[str, Any]]) -> Optional[str]:
        """Extract system prompt from Mistral messages."""
        if messages and isinstance(messages, list) and len(messages) > 0:
            first = messages[0]
            if isinstance(first, dict) and first.get("role") == "system":
                content = first.get("content", "")
                return content if isinstance(content, str) else str(content)
        return None

    def _track_completion(
        self,
        *,
        user_id: str,
        model: str,
        messages: List[Dict[str, Any]],
        response: Any,
        latency_ms: float,
        session_id: Optional[str],
        trace_id: Optional[str],
        turn_id: Optional[int],
        is_error: bool = False,
        error_message: Optional[str] = None,
        total_cost_usd: Optional[float] = None,
        provider_ttfb_ms: Optional[float] = None,
        event_properties: Optional[Dict[str, Any]] = None,
        user_properties: Optional[Dict[str, Any]] = None,
        groups: Optional[Dict[str, Any]] = None,
        api_kwargs: Optional[Dict[str, Any]] = None,
    ) -> None:
        try:
            current_turn = turn_id or 1

            # Extract agent-level fields that _apply_session_context may
            # have injected into event_properties.
            _ep = dict(event_properties) if event_properties else {}
            _ctx_agent_id = _ep.pop("agent_id", None)
            _ctx_parent_agent_id = _ep.pop("parent_agent_id", None)
            _ctx_env = _ep.pop("env", None)
            _ctx_customer_org_id = _ep.pop("customer_org_id", None)
            _ctx_agent_version = _ep.pop("agent_version", None)
            _ctx_context = _ep.pop("context", None)
            if not _ep and event_properties is not None:
                event_properties = None
            else:
                event_properties = _ep or event_properties

            # In agent loops each create() re-sends the full conversation.
            # Only track user messages appended after the last assistant/tool
            # reply so the same message isn't autotracked on every iteration.
            last_reply = -1
            for idx, m in enumerate(messages):
                if m.get("role") in ("assistant", "tool"):
                    last_reply = idx
            for msg in messages[last_reply + 1:]:
                if msg.get("role") == "user":
                    track_user_message(
                        amplitude=self.amplitude,
                        user_id=user_id,
                        message_content=msg.get("content", ""),
                        session_id=session_id,
                        trace_id=trace_id,
                        turn_id=current_turn,
                        agent_id=_ctx_agent_id,
                        parent_agent_id=_ctx_parent_agent_id,
                        env=_ctx_env,
                        customer_org_id=_ctx_customer_org_id,
                        agent_version=_ctx_agent_version,
                        context=_ctx_context,
                        message_source="agent" if _ctx_parent_agent_id else "user",
                        event_properties=event_properties,
                        user_properties=user_properties,
                        groups=groups,
                        privacy_config=self.privacy_config,
                    )
                    current_turn += 1

            response_model = getattr(response, "model", None) if response else None
            effective_model = response_model or model

            response_content = self._extract_response_content(response) if not is_error else ""
            usage = self._extract_usage(response) if not is_error else {}
            finish_reason = self._extract_finish_reason(response) if not is_error else None
            tool_calls = self._extract_tool_calls(response) if not is_error else None
            reasoning_content = (
                self._extract_reasoning_content(response) if not is_error else None
            )

            if total_cost_usd is None and not is_error:
                inp = usage.get("input_tokens")
                out = usage.get("output_tokens")
                if inp is not None and out is not None:
                    try:
                        total_cost_usd = calculate_cost(
                            model_name=model,
                            input_tokens=inp,
                            output_tokens=out,
                            default_provider='mistral',
                        )
                    except Exception:
                        pass

            # Extract model configuration
            system_prompt = self._extract_system_prompt(messages)
            _kw = api_kwargs or {}
            temperature = _kw.get("temperature")
            top_p = _kw.get("top_p")
            max_output_tokens = _kw.get("max_tokens")
            is_streaming = _kw.get("stream")

            track_ai_message(
                amplitude=self.amplitude,
                user_id=user_id,
                model_name=effective_model,
                provider="mistral",
                response_content=response_content,
                latency_ms=latency_ms,
                session_id=session_id,
                trace_id=trace_id,
                turn_id=current_turn,
                agent_id=_ctx_agent_id,
                parent_agent_id=_ctx_parent_agent_id,
                env=_ctx_env,
                customer_org_id=_ctx_customer_org_id,
                agent_version=_ctx_agent_version,
                context=_ctx_context,
                provider_ttfb_ms=provider_ttfb_ms,
                is_error=is_error,
                error_message=error_message,
                finish_reason=finish_reason,
                tool_calls=tool_calls,
                total_cost_usd=total_cost_usd,
                reasoning_content=reasoning_content,
                system_prompt=system_prompt,
                temperature=temperature,
                max_output_tokens=max_output_tokens,
                top_p=top_p,
                is_streaming=is_streaming,
                event_properties=event_properties,
                user_properties=user_properties,
                groups=groups,
                privacy_config=self.privacy_config,
                input_tokens=usage.get("input_tokens"),
                output_tokens=usage.get("output_tokens"),
                total_tokens=usage.get("total_tokens"),
            )
        except Exception as exc:
            get_logger(self.amplitude).error(f"Failed to track Mistral completion: {exc}")


class _WrappedChat:
    """Namespace wrapper for ``client.chat.complete(...)``."""

    def __init__(self, wrapper: Mistral) -> None:
        self._wrapper = wrapper

    def complete(
        self,
        *,
        model: str,
        messages: List[Dict[str, Any]],
        amplitude_user_id: Optional[str] = None,
        amplitude_session_id: Optional[str] = None,
        amplitude_trace_id: Optional[str] = None,
        amplitude_turn_id: Optional[int] = None,
        amplitude_event_properties: Optional[Dict[str, Any]] = None,
        amplitude_user_properties: Optional[Dict[str, Any]] = None,
        amplitude_groups: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> Any:
        """Non-streaming chat completion with tracking."""
        # Auto-inherit session/agent context when inside agent.session()
        (
            amplitude_user_id, amplitude_session_id,
            amplitude_trace_id, amplitude_turn_id,
            amplitude_event_properties, amplitude_groups,
        ) = _apply_session_context(
            user_id=amplitude_user_id,
            session_id=amplitude_session_id,
            trace_id=amplitude_trace_id,
            turn_id=amplitude_turn_id,
            event_properties=amplitude_event_properties,
            groups=amplitude_groups,
        )
        trace_id = amplitude_trace_id or str(uuid.uuid4())
        session_id = amplitude_session_id or str(uuid.uuid4())
        start_time = time.time()

        try:
            response = self._wrapper._client.chat.complete(
                model=model, messages=messages, **kwargs
            )
        except Exception as exc:
            latency_ms = (time.time() - start_time) * 1000
            if amplitude_user_id:
                self._wrapper._track_completion(
                    user_id=amplitude_user_id,
                    model=model,
                    messages=messages,
                    response=None,
                    latency_ms=latency_ms,
                    session_id=session_id,
                    trace_id=trace_id,
                    turn_id=amplitude_turn_id,
                    is_error=True,
                    error_message=str(exc),
                    event_properties=amplitude_event_properties,
                    user_properties=amplitude_user_properties,
                    groups=amplitude_groups,
                    api_kwargs=kwargs,
                )
            raise

        latency_ms = (time.time() - start_time) * 1000
        if amplitude_user_id:
            self._wrapper._track_completion(
                user_id=amplitude_user_id,
                model=model,
                messages=messages,
                response=response,
                latency_ms=latency_ms,
                session_id=session_id,
                trace_id=trace_id,
                turn_id=amplitude_turn_id,
                event_properties=amplitude_event_properties,
                user_properties=amplitude_user_properties,
                groups=amplitude_groups,
                api_kwargs=kwargs,
            )
        return response

    def stream(
        self,
        *,
        model: str,
        messages: List[Dict[str, Any]],
        amplitude_user_id: Optional[str] = None,
        amplitude_session_id: Optional[str] = None,
        amplitude_trace_id: Optional[str] = None,
        amplitude_turn_id: Optional[int] = None,
        amplitude_event_properties: Optional[Dict[str, Any]] = None,
        amplitude_user_properties: Optional[Dict[str, Any]] = None,
        amplitude_groups: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> Iterator[Any]:
        """Streaming chat completion with tracking."""
        # Auto-inherit session/agent context when inside agent.session()
        (
            amplitude_user_id, amplitude_session_id,
            amplitude_trace_id, amplitude_turn_id,
            amplitude_event_properties, amplitude_groups,
        ) = _apply_session_context(
            user_id=amplitude_user_id,
            session_id=amplitude_session_id,
            trace_id=amplitude_trace_id,
            turn_id=amplitude_turn_id,
            event_properties=amplitude_event_properties,
            groups=amplitude_groups,
        )
        trace_id = amplitude_trace_id or str(uuid.uuid4())
        session_id = amplitude_session_id or str(uuid.uuid4())
        start_time = time.time()
        first_byte_time: Optional[float] = None

        content_chunks: List[str] = []
        reasoning_chunks: List[str] = []
        streaming_tool_calls: List[Dict[str, Any]] = []
        usage_data: Dict[str, Optional[int]] = {}
        finish_reason: Optional[str] = None
        is_error = False
        error_message: Optional[str] = None

        try:
            stream_resp = self._wrapper._client.chat.stream(
                model=model, messages=messages, **kwargs
            )
        except Exception as exc:
            is_error = True
            error_message = str(exc)
            if amplitude_user_id:
                self._wrapper._track_completion(
                    user_id=amplitude_user_id,
                    model=model,
                    messages=messages,
                    response=None,
                    latency_ms=(time.time() - start_time) * 1000,
                    session_id=session_id,
                    trace_id=trace_id,
                    turn_id=amplitude_turn_id,
                    is_error=True,
                    error_message=error_message,
                    event_properties=amplitude_event_properties,
                    user_properties=amplitude_user_properties,
                    groups=amplitude_groups,
                    api_kwargs={**kwargs, "stream": True},
                )
            raise

        try:
            for event in stream_resp:
                try:
                    chunk = getattr(event, "data", event)
                    if hasattr(chunk, "choices") and chunk.choices:
                        delta = getattr(chunk.choices[0], "delta", None)
                        if delta and hasattr(delta, "content") and delta.content:
                            delta_content = delta.content
                            if first_byte_time is None:
                                first_byte_time = time.time()
                            # Magistral streaming: content may be a string
                            # or a list of typed chunks (text vs thinking).
                            if isinstance(delta_content, str):
                                content_chunks.append(delta_content)
                            elif isinstance(delta_content, list):
                                for dc in delta_content:
                                    if not isinstance(dc, dict):
                                        continue
                                    dc_type = dc.get("type")
                                    if dc_type == "text":
                                        content_chunks.append(dc.get("text", ""))
                                    elif dc_type == "thinking":
                                        for t in dc.get("thinking", []):
                                            if isinstance(t, dict) and t.get("type") == "text":
                                                reasoning_chunks.append(t.get("text", ""))
                            else:
                                content_chunks.append(str(delta_content))
                        # Reasoning via a dedicated attribute (future SDK versions)
                        reasoning_attr = getattr(delta, "reasoning_content", None)
                        if reasoning_attr:
                            reasoning_chunks.append(str(reasoning_attr))
                        # Tool calls
                        tool_calls_delta = getattr(delta, "tool_calls", None)
                        if tool_calls_delta:
                            for tc in tool_calls_delta:
                                func = getattr(tc, "function", None)
                                tc_id = getattr(tc, "id", None)
                                tc_name = getattr(func, "name", None) if func else None
                                tc_args = getattr(func, "arguments", None) if func else None
                                tc_index = getattr(tc, "index", None)
                                if tc_id and tc_name:
                                    streaming_tool_calls.append({
                                        "type": "function",
                                        "id": tc_id,
                                        "function": {"name": tc_name, "arguments": tc_args or ""},
                                    })
                                elif tc_args and streaming_tool_calls:
                                    if tc_index is not None and tc_index < len(streaming_tool_calls):
                                        streaming_tool_calls[tc_index]["function"]["arguments"] += tc_args
                                    else:
                                        streaming_tool_calls[-1]["function"]["arguments"] += tc_args
                        fr = getattr(chunk.choices[0], "finish_reason", None)
                        if fr:
                            finish_reason = fr

                    usage = getattr(chunk, "usage", None)
                    if usage:
                        if hasattr(usage, "prompt_tokens"):
                            usage_data["input_tokens"] = usage.prompt_tokens
                        if hasattr(usage, "completion_tokens"):
                            usage_data["output_tokens"] = usage.completion_tokens
                        if hasattr(usage, "total_tokens"):
                            usage_data["total_tokens"] = usage.total_tokens
                except Exception as tracking_exc:
                    get_logger(self._wrapper.amplitude).warning(
                        f"Tracking failed for Mistral chunk: {tracking_exc}"
                    )
                yield event

        except Exception as exc:
            is_error = True
            error_message = str(exc)
            raise
        finally:
            if amplitude_user_id:
                latency_ms = (time.time() - start_time) * 1000
                ttfb_ms = (
                    (first_byte_time - start_time) * 1000
                    if first_byte_time
                    else None
                )

                # Build a minimal response-like object for tracking
                class _SynthResponse:
                    pass

                class _SynthChoice:
                    pass

                class _SynthMessage:
                    pass

                class _SynthUsage:
                    pass

                synth = _SynthResponse()

                # Build a synthetic message with accumulated content
                full_content = "".join(content_chunks)
                full_reasoning = "".join(reasoning_chunks) if reasoning_chunks else None

                synth_msg = _SynthMessage()
                # If we have reasoning, use magistral list-of-chunks format
                # so _extract_response_content / _extract_reasoning_content work.
                if full_reasoning:
                    msg_content: Any = []
                    if full_reasoning:
                        msg_content.append({
                            "type": "thinking",
                            "thinking": [{"type": "text", "text": full_reasoning}],
                        })
                    if full_content:
                        msg_content.append({"type": "text", "text": full_content})
                    synth_msg.content = msg_content  # type: ignore[attr-defined]
                else:
                    synth_msg.content = full_content  # type: ignore[attr-defined]
                if streaming_tool_calls:
                    class _SynthToolCall:
                        pass
                    class _SynthFunction:
                        pass
                    synth_tcs = []
                    for tc in streaming_tool_calls:
                        stc = _SynthToolCall()
                        stc.type = tc.get("type", "function")  # type: ignore[attr-defined]
                        stc.id = tc.get("id", "")  # type: ignore[attr-defined]
                        sf = _SynthFunction()
                        sf.name = tc["function"]["name"]  # type: ignore[attr-defined]
                        sf.arguments = tc["function"]["arguments"]  # type: ignore[attr-defined]
                        stc.function = sf  # type: ignore[attr-defined]
                        synth_tcs.append(stc)
                    synth_msg.tool_calls = synth_tcs  # type: ignore[attr-defined]
                else:
                    synth_msg.tool_calls = None  # type: ignore[attr-defined]

                synth_choice = _SynthChoice()
                synth_choice.message = synth_msg  # type: ignore[attr-defined]
                synth_choice.finish_reason = finish_reason  # type: ignore[attr-defined]
                synth.choices = [synth_choice]  # type: ignore[attr-defined]

                synth_usage = _SynthUsage()
                synth_usage.prompt_tokens = usage_data.get("input_tokens")  # type: ignore[attr-defined]
                synth_usage.completion_tokens = usage_data.get("output_tokens")  # type: ignore[attr-defined]
                synth_usage.total_tokens = usage_data.get("total_tokens")  # type: ignore[attr-defined]
                synth.usage = synth_usage  # type: ignore[attr-defined]

                self._wrapper._track_completion(
                    user_id=amplitude_user_id,
                    model=model,
                    messages=messages,
                    response=synth if not is_error else None,
                    latency_ms=latency_ms,
                    session_id=session_id,
                    trace_id=trace_id,
                    turn_id=amplitude_turn_id,
                    is_error=is_error,
                    error_message=error_message,
                    provider_ttfb_ms=ttfb_ms,
                    event_properties=amplitude_event_properties,
                    user_properties=amplitude_user_properties,
                    groups=amplitude_groups,
                    api_kwargs={**kwargs, "stream": True},
                )
