"""
Cohere SDK instrumentation.

Patches Cohere SDK methods to capture:
- Chat (sync and async)
- Generate (sync and async)
- Embed (sync and async)

Based on ARCHITECTURE_V2.md Section 6.3.
"""

from __future__ import annotations

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

import wrapt

from risicare.integrations._base import scrub_sensitive

logger = logging.getLogger(__name__)

T = TypeVar("T")

_instrumented = False


def instrument_cohere(module: Any) -> None:
    """
    Apply instrumentation to Cohere module.

    This patches:
    - cohere.Client.chat
    - cohere.AsyncClient.chat
    - cohere.Client.generate
    - cohere.AsyncClient.generate
    - cohere.Client.embed
    - cohere.AsyncClient.embed
    """
    global _instrumented
    if _instrumented:
        return

    try:
        # Patch chat methods
        wrapt.wrap_function_wrapper(module, "Client.chat", _wrap_chat)
        wrapt.wrap_function_wrapper(module, "AsyncClient.chat", _wrap_async_chat)

        # Patch generate methods
        wrapt.wrap_function_wrapper(module, "Client.generate", _wrap_generate)
        wrapt.wrap_function_wrapper(module, "AsyncClient.generate", _wrap_async_generate)

        # Patch embed methods
        wrapt.wrap_function_wrapper(module, "Client.embed", _wrap_embed)
        wrapt.wrap_function_wrapper(module, "AsyncClient.embed", _wrap_async_embed)

        _instrumented = True
        logger.debug("Instrumented Cohere SDK")

    except Exception as e:
        logger.warning(f"Failed to instrument Cohere 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 _wrap_chat(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for cohere.Client.chat."""
    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", "command-r-plus")
    trace_content = _should_trace_content()

    with tracer.start_span(
        name=f"cohere.chat/{model}",
        kind=SpanKind.LLM_CALL,
        attributes={
            "gen_ai.system": "cohere",
            "gen_ai.request.model": model,
            "gen_ai.operation": "chat",
        },
    ) as span:
        # Capture message
        message = kwargs.get("message", args[0] if args else "")
        if trace_content and message:
            span.set_attribute(
                "gen_ai.prompt.message",
                message[:10000] if len(message) > 10000 else message,
            )

        # Capture chat history
        chat_history = kwargs.get("chat_history", [])
        span.set_attribute("gen_ai.prompt.history_count", len(chat_history))

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

            span.set_attribute("gen_ai.latency_ms", latency_ms)
            span.set_attribute("gen_ai.response.id", getattr(response, "generation_id", ""))

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

            # Response text
            if trace_content:
                text = getattr(response, "text", "")
                if text:
                    span.set_attribute(
                        "gen_ai.completion.content",
                        text[:10000] if len(text) > 10000 else text,
                    )

            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(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for cohere.AsyncClient.chat."""
    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", "command-r-plus")
    trace_content = _should_trace_content()

    with tracer.start_span(
        name=f"cohere.chat/{model}",
        kind=SpanKind.LLM_CALL,
        attributes={
            "gen_ai.system": "cohere",
            "gen_ai.request.model": model,
            "gen_ai.operation": "chat",
        },
    ) as span:
        message = kwargs.get("message", args[0] if args else "")
        if trace_content and message:
            span.set_attribute(
                "gen_ai.prompt.message",
                message[:10000] if len(message) > 10000 else message,
            )

        chat_history = kwargs.get("chat_history", [])
        span.set_attribute("gen_ai.prompt.history_count", len(chat_history))

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

            span.set_attribute("gen_ai.latency_ms", latency_ms)
            span.set_attribute("gen_ai.response.id", getattr(response, "generation_id", ""))

            meta = getattr(response, "meta", None)
            if meta:
                tokens = getattr(meta, "tokens", None)
                if tokens:
                    span.set_attribute(
                        "gen_ai.usage.prompt_tokens",
                        getattr(tokens, "input_tokens", 0),
                    )
                    span.set_attribute(
                        "gen_ai.usage.completion_tokens",
                        getattr(tokens, "output_tokens", 0),
                    )

            if trace_content:
                text = getattr(response, "text", "")
                if text:
                    span.set_attribute(
                        "gen_ai.completion.content",
                        text[:10000] if len(text) > 10000 else text,
                    )

            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_generate(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for cohere.Client.generate."""
    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", "command")

    with tracer.start_span(
        name=f"cohere.generate/{model}",
        kind=SpanKind.LLM_CALL,
        attributes={
            "gen_ai.system": "cohere",
            "gen_ai.request.model": model,
            "gen_ai.operation": "generate",
        },
    ) as span:
        try:
            start_time = time.perf_counter()
            response = wrapped(*args, **kwargs)
            latency_ms = (time.perf_counter() - start_time) * 1000

            span.set_attribute("gen_ai.latency_ms", latency_ms)
            span.set_attribute("gen_ai.response.generations", len(getattr(response, "generations", [])))

            return response

        except Exception as e:
            span.record_exception(e)
            span.set_attribute("error", True)
            raise


async def _wrap_async_generate(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for cohere.AsyncClient.generate."""
    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", "command")

    with tracer.start_span(
        name=f"cohere.generate/{model}",
        kind=SpanKind.LLM_CALL,
        attributes={
            "gen_ai.system": "cohere",
            "gen_ai.request.model": model,
            "gen_ai.operation": "generate",
        },
    ) as span:
        try:
            start_time = time.perf_counter()
            response = await wrapped(*args, **kwargs)
            latency_ms = (time.perf_counter() - start_time) * 1000

            span.set_attribute("gen_ai.latency_ms", latency_ms)
            span.set_attribute("gen_ai.response.generations", len(getattr(response, "generations", [])))

            return response

        except Exception as e:
            span.record_exception(e)
            span.set_attribute("error", True)
            raise


def _wrap_embed(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for cohere.Client.embed."""
    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", "embed-english-v3.0")

    with tracer.start_span(
        name=f"cohere.embed/{model}",
        kind=SpanKind.LLM_CALL,
        attributes={
            "gen_ai.system": "cohere",
            "gen_ai.request.model": model,
            "gen_ai.operation": "embed",
        },
    ) as span:
        texts = kwargs.get("texts", args[0] if args else [])
        span.set_attribute("gen_ai.input.count", len(texts) if isinstance(texts, list) else 1)

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

            span.set_attribute("gen_ai.latency_ms", latency_ms)
            embeddings = getattr(response, "embeddings", [])
            span.set_attribute("gen_ai.response.embeddings", len(embeddings))

            return response

        except Exception as e:
            span.record_exception(e)
            span.set_attribute("error", True)
            raise


async def _wrap_async_embed(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for cohere.AsyncClient.embed."""
    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", "embed-english-v3.0")

    with tracer.start_span(
        name=f"cohere.embed/{model}",
        kind=SpanKind.LLM_CALL,
        attributes={
            "gen_ai.system": "cohere",
            "gen_ai.request.model": model,
            "gen_ai.operation": "embed",
        },
    ) as span:
        texts = kwargs.get("texts", args[0] if args else [])
        span.set_attribute("gen_ai.input.count", len(texts) if isinstance(texts, list) else 1)

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

            span.set_attribute("gen_ai.latency_ms", latency_ms)
            embeddings = getattr(response, "embeddings", [])
            span.set_attribute("gen_ai.response.embeddings", len(embeddings))

            return response

        except Exception as e:
            span.record_exception(e)
            span.set_attribute("error", True)
            raise
