"""
Mistral AI integration — wraps Mistral client to auto-trace chat completions.

Supports ``chat.complete()``, ``chat.complete_async()``, and streaming
via ``chat.stream()``.

Usage:
    from mistralai import Mistral
    from syntropy import Syntropy
    from syntropy.integrations.mistral import wrap_mistral

    syn = Syntropy(api_key="syn_...", agent_id="my-agent")
    client = wrap_mistral(Mistral(api_key="..."), syntropy=syn)

    # Non-streaming (auto-traced):
    response = client.chat.complete(
        model="mistral-large-latest",
        messages=[{"role": "user", "content": "Hello!"}],
    )

    # Streaming (auto-traced on completion):
    for event in client.chat.stream(
        model="mistral-large-latest",
        messages=[{"role": "user", "content": "Tell me a story"}],
    ):
        print(event.data.choices[0].delta.content, end="")
"""

from __future__ import annotations

import time
from typing import Any, Optional

from syntropy.client import Syntropy
from syntropy.trace_context import is_inside_trace


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

def _extract_prompt(kwargs: dict) -> str:
    messages = kwargs.get("messages", [])
    if not messages:
        return ""
    last = messages[-1]
    return last.get("content", "") if isinstance(last, dict) else str(last)


def _wrap_stream(
    stream: Any,
    prompt: str,
    start: float,
    model_kwarg: str,
    syntropy: Syntropy,
    agent_id: str,
) -> Any:
    """Yield SSE-like events from Mistral stream and trace on end."""
    text_parts: list[str] = []
    model_name = model_kwarg
    tokens_in = 0
    tokens_out = 0
    finish_reason: Optional[str] = None
    stream_error: BaseException | None = None

    try:
        for event in stream:
            # Mistral streams emit CompletionEvent objects with a `.data` attr
            data = getattr(event, "data", event)

            # Model name
            m = getattr(data, "model", None)
            if m:
                model_name = m

            # Usage (often on the last event)
            usage = getattr(data, "usage", None)
            if usage:
                tokens_in = getattr(usage, "prompt_tokens", tokens_in)
                tokens_out = getattr(usage, "completion_tokens", tokens_out)

            # Content delta
            choices = getattr(data, "choices", None) or []
            if choices:
                delta = getattr(choices[0], "delta", None)
                content = getattr(delta, "content", None) or ""
                if content:
                    text_parts.append(content)
                fr = getattr(choices[0], "finish_reason", None)
                if fr:
                    finish_reason = str(fr)

            yield event
    except BaseException as exc:
        stream_error = exc
        raise
    finally:
        elapsed_ms = int((time.perf_counter() - start) * 1000)
        syntropy.trace(
            prompt=prompt,
            response="".join(text_parts),
            model=model_name,
            agent_id=agent_id,
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            latency_ms=elapsed_ms,
            status="error" if stream_error else "success",
            metadata={
                "provider": "mistral",
                "stream": True,
                "finish_reason": finish_reason,
                **({
                    "error": str(stream_error),
                    "partialResponse": True,
                } if stream_error else {}),
            },
        )


# ----------------------------------------------------------------
# Traced classes
# ----------------------------------------------------------------

class _TracedChatComplete:
    """Proxy around mistral.chat that traces .complete(), .complete_async(), and .stream()."""

    def __init__(self, original: Any, syntropy: Syntropy, agent_id: str):
        self._original = original
        self._syntropy = syntropy
        self._agent_id = agent_id

    # ---- Non-streaming sync ----
    def complete(self, **kwargs: Any) -> Any:
        # Skip if already inside an orchestration-level trace (e.g. LangChain)
        if is_inside_trace():
            return self._original.complete(**kwargs)

        start = time.perf_counter()
        response = self._original.complete(**kwargs)
        elapsed_ms = int((time.perf_counter() - start) * 1000)

        prompt = _extract_prompt(kwargs)

        output = ""
        finish_reason = None
        if hasattr(response, "choices") and response.choices:
            choice = response.choices[0]
            msg = getattr(choice, "message", None)
            output = getattr(msg, "content", "") or ""
            finish_reason = getattr(choice, "finish_reason", None)

        usage = getattr(response, "usage", None)
        tokens_in = getattr(usage, "prompt_tokens", 0) if usage else 0
        tokens_out = getattr(usage, "completion_tokens", 0) if usage else 0

        self._syntropy.trace(
            prompt=prompt,
            response=output,
            model=getattr(response, "model", kwargs.get("model", "unknown")),
            agent_id=self._agent_id,
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            latency_ms=elapsed_ms,
            metadata={
                "provider": "mistral",
                "stream": False,
                "finish_reason": str(finish_reason) if finish_reason else None,
            },
        )

        return response

    # ---- Non-streaming async ----
    async def complete_async(self, **kwargs: Any) -> Any:
        start = time.perf_counter()
        response = await self._original.complete_async(**kwargs)
        elapsed_ms = int((time.perf_counter() - start) * 1000)

        prompt = _extract_prompt(kwargs)

        output = ""
        if hasattr(response, "choices") and response.choices:
            msg = getattr(response.choices[0], "message", None)
            output = getattr(msg, "content", "") or ""

        usage = getattr(response, "usage", None)
        tokens_in = getattr(usage, "prompt_tokens", 0) if usage else 0
        tokens_out = getattr(usage, "completion_tokens", 0) if usage else 0

        self._syntropy.trace(
            prompt=prompt,
            response=output,
            model=getattr(response, "model", kwargs.get("model", "unknown")),
            agent_id=self._agent_id,
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            latency_ms=elapsed_ms,
            metadata={"provider": "mistral", "stream": False},
        )

        return response

    # ---- Streaming ----
    def stream(self, **kwargs: Any) -> Any:
        start = time.perf_counter()
        prompt = _extract_prompt(kwargs)
        model_kwarg = kwargs.get("model", "unknown")
        raw_stream = self._original.stream(**kwargs)

        return _wrap_stream(
            raw_stream, prompt, start, model_kwarg,
            self._syntropy, self._agent_id,
        )

    def __getattr__(self, name: str) -> Any:
        return getattr(self._original, name)


class _TracedMistralClient:
    """Proxy around Mistral client that injects traced chat namespace."""

    def __init__(self, client: Any, syntropy: Syntropy, agent_id: str):
        self._client = client
        self._syntropy = syntropy
        self._agent_id = agent_id
        self.chat = _TracedChatComplete(client.chat, syntropy, agent_id)

    def __getattr__(self, name: str) -> Any:
        return getattr(self._client, name)


def wrap_mistral(
    client: Any,
    *,
    syntropy: Syntropy,
    agent_id: str = "",
) -> Any:
    """
    Wrap a Mistral client to auto-trace chat.complete() and chat.stream() calls.

    Args:
        client: A mistralai.Mistral() instance.
        syntropy: A Syntropy client instance.
        agent_id: Override agent_id.

    Returns:
        A drop-in replacement with tracing enabled.
    """
    return _TracedMistralClient(
        client, syntropy, agent_id or syntropy._config.agent_id
    )
