"""
Mistral AI SDK instrumentation.

Patches Mistral SDK methods to capture:
- Chat completions (sync and async)
- Streaming responses

Based on ARCHITECTURE_V2.md Section 6.3.
"""

from __future__ import annotations

import logging
import time
from typing import Any, AsyncIterator, Callable, Dict, Iterator, Optional, TypeVar

import wrapt

from risicare.integrations._base import scrub_sensitive

logger = logging.getLogger(__name__)

T = TypeVar("T")

_instrumented = False


def instrument_mistral(module: Any) -> None:
    """
    Apply instrumentation to Mistral AI module.

    This patches:
    - mistralai.Mistral.chat.complete
    - mistralai.Mistral.chat.stream
    """
    global _instrumented
    if _instrumented:
        return

    try:
        # Try new SDK structure (mistralai >= 1.0)
        try:
            wrapt.wrap_function_wrapper(
                module,
                "Mistral.chat.complete",
                _wrap_chat_complete,
            )
            wrapt.wrap_function_wrapper(
                module,
                "Mistral.chat.complete_async",
                _wrap_async_chat_complete,
            )
        except AttributeError:
            # Try old SDK structure
            wrapt.wrap_function_wrapper(
                module,
                "MistralClient.chat",
                _wrap_legacy_chat,
            )
            wrapt.wrap_function_wrapper(
                module,
                "MistralAsyncClient.chat",
                _wrap_async_legacy_chat,
            )

        _instrumented = True
        logger.debug("Instrumented Mistral AI SDK")

    except Exception as e:
        logger.warning(f"Failed to instrument Mistral AI SDK: {e}")


def _get_tracer() -> Optional[Any]:
    """Get the Risicare tracer if available."""
    try:
        from risicare.tracer import get_tracer

        return get_tracer()
    except ImportError:
        return None


def _get_config() -> Optional[Any]:
    """Get the Risicare client config if available."""
    try:
        from risicare.client import get_client

        client = get_client()
        return client.config if client else None
    except ImportError:
        return None


def _should_trace_content() -> bool:
    """Check if content should be traced."""
    config = _get_config()
    return config.trace_content if config else False


def _extract_model_info(kwargs: Dict[str, Any]) -> Dict[str, Any]:
    """Extract model information from request kwargs.

    Only JSON-primitive values are included — SDK sentinels are skipped.
    """
    from risicare.integrations._base import is_json_primitive

    temperature = kwargs.get("temperature")
    max_tokens = kwargs.get("max_tokens")

    info: Dict[str, Any] = {
        "gen_ai.system": "mistral",
        "gen_ai.request.model": kwargs.get("model", "unknown"),
    }

    if is_json_primitive(temperature):
        info["gen_ai.request.temperature"] = temperature
    if is_json_primitive(max_tokens):
        info["gen_ai.request.max_tokens"] = max_tokens

    return info


def _capture_messages(span: Any, messages: list, trace_content: bool) -> None:
    """Capture input messages to span."""
    if not trace_content:
        span.set_attribute("gen_ai.prompt.count", len(messages))
        return

    for i, msg in enumerate(messages[:10]):
        role = msg.get("role", "") if isinstance(msg, dict) else getattr(msg, "role", "")
        span.set_attribute(f"gen_ai.prompt.{i}.role", role)

        content = msg.get("content", "") if isinstance(msg, dict) else getattr(msg, "content", "")
        if isinstance(content, str):
            span.set_attribute(
                f"gen_ai.prompt.{i}.content",
                content[:10000] if len(content) > 10000 else content,
            )


def _capture_response(span: Any, response: Any, latency_ms: float, trace_content: bool) -> None:
    """Capture response data into span."""
    span.set_attribute("gen_ai.latency_ms", latency_ms)
    span.set_attribute("gen_ai.response.id", getattr(response, "id", ""))
    span.set_attribute("gen_ai.response.model", getattr(response, "model", "unknown"))

    # Token usage
    usage = getattr(response, "usage", None)
    if usage:
        span.set_attribute("gen_ai.usage.prompt_tokens", getattr(usage, "prompt_tokens", 0))
        span.set_attribute("gen_ai.usage.completion_tokens", getattr(usage, "completion_tokens", 0))
        span.set_attribute("gen_ai.usage.total_tokens", getattr(usage, "total_tokens", 0))

    # Choices
    choices = getattr(response, "choices", [])
    span.set_attribute("gen_ai.response.choices", len(choices))

    if choices and trace_content:
        first = choices[0]
        message = getattr(first, "message", None)
        if message:
            content = getattr(message, "content", "")
            if content:
                span.set_attribute(
                    "gen_ai.completion.content",
                    content[:10000] if len(content) > 10000 else content,
                )

            # Tool calls
            tool_calls = getattr(message, "tool_calls", None)
            if tool_calls:
                span.set_attribute("gen_ai.completion.tool_calls", len(tool_calls))

        span.set_attribute("gen_ai.completion.finish_reason", getattr(first, "finish_reason", ""))


def _wrap_chat_complete(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for Mistral.chat.complete (new SDK)."""
    tracer = _get_tracer()
    if tracer is None or not tracer.is_enabled:
        return wrapped(*args, **kwargs)

    from risicare.integrations._dedup import is_provider_suppressed

    if is_provider_suppressed():
        return wrapped(*args, **kwargs)

    from risicare_core import SpanKind

    model = kwargs.get("model", "mistral-large")
    trace_content = _should_trace_content()

    with tracer.start_span(
        name=f"mistral.chat.complete/{model}",
        kind=SpanKind.LLM_CALL,
        attributes=_extract_model_info(kwargs),
    ) as span:
        messages = kwargs.get("messages", [])
        _capture_messages(span, messages, trace_content)

        try:
            start_time = time.perf_counter()
            response = wrapped(*args, **kwargs)
            latency_ms = (time.perf_counter() - start_time) * 1000

            _capture_response(span, response, latency_ms, trace_content)

            return response

        except Exception as e:
            span.record_exception(e)
            span.set_attribute("error", True)
            span.set_attribute("error.type", type(e).__name__)
            span.set_attribute("error.message", scrub_sensitive(str(e)[:2000]))
            raise


async def _wrap_async_chat_complete(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for Mistral.chat.complete_async (new SDK)."""
    tracer = _get_tracer()
    if tracer is None or not tracer.is_enabled:
        return await wrapped(*args, **kwargs)

    from risicare.integrations._dedup import is_provider_suppressed

    if is_provider_suppressed():
        return await wrapped(*args, **kwargs)

    from risicare_core import SpanKind

    model = kwargs.get("model", "mistral-large")
    trace_content = _should_trace_content()

    with tracer.start_span(
        name=f"mistral.chat.complete/{model}",
        kind=SpanKind.LLM_CALL,
        attributes=_extract_model_info(kwargs),
    ) as span:
        messages = kwargs.get("messages", [])
        _capture_messages(span, messages, trace_content)

        try:
            start_time = time.perf_counter()
            response = await wrapped(*args, **kwargs)
            latency_ms = (time.perf_counter() - start_time) * 1000

            _capture_response(span, response, latency_ms, trace_content)

            return response

        except Exception as e:
            span.record_exception(e)
            span.set_attribute("error", True)
            span.set_attribute("error.type", type(e).__name__)
            span.set_attribute("error.message", scrub_sensitive(str(e)[:2000]))
            raise


def _wrap_legacy_chat(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for MistralClient.chat (legacy SDK)."""
    tracer = _get_tracer()
    if tracer is None or not tracer.is_enabled:
        return wrapped(*args, **kwargs)

    from risicare.integrations._dedup import is_provider_suppressed

    if is_provider_suppressed():
        return wrapped(*args, **kwargs)

    from risicare_core import SpanKind

    model = kwargs.get("model", "mistral-large")
    trace_content = _should_trace_content()

    with tracer.start_span(
        name=f"mistral.chat/{model}",
        kind=SpanKind.LLM_CALL,
        attributes=_extract_model_info(kwargs),
    ) as span:
        messages = kwargs.get("messages", [])
        _capture_messages(span, messages, trace_content)

        try:
            start_time = time.perf_counter()
            response = wrapped(*args, **kwargs)
            latency_ms = (time.perf_counter() - start_time) * 1000

            _capture_response(span, response, latency_ms, trace_content)

            return response

        except Exception as e:
            span.record_exception(e)
            span.set_attribute("error", True)
            span.set_attribute("error.type", type(e).__name__)
            span.set_attribute("error.message", scrub_sensitive(str(e)[:2000]))
            raise


async def _wrap_async_legacy_chat(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for MistralAsyncClient.chat (legacy SDK)."""
    tracer = _get_tracer()
    if tracer is None or not tracer.is_enabled:
        return await wrapped(*args, **kwargs)

    from risicare.integrations._dedup import is_provider_suppressed

    if is_provider_suppressed():
        return await wrapped(*args, **kwargs)

    from risicare_core import SpanKind

    model = kwargs.get("model", "mistral-large")
    trace_content = _should_trace_content()

    with tracer.start_span(
        name=f"mistral.chat/{model}",
        kind=SpanKind.LLM_CALL,
        attributes=_extract_model_info(kwargs),
    ) as span:
        messages = kwargs.get("messages", [])
        _capture_messages(span, messages, trace_content)

        try:
            start_time = time.perf_counter()
            response = await wrapped(*args, **kwargs)
            latency_ms = (time.perf_counter() - start_time) * 1000

            _capture_response(span, response, latency_ms, trace_content)

            return response

        except Exception as e:
            span.record_exception(e)
            span.set_attribute("error", True)
            span.set_attribute("error.type", type(e).__name__)
            span.set_attribute("error.message", scrub_sensitive(str(e)[:2000]))
            raise
