"""
Together AI SDK instrumentation.

Patches Together AI SDK methods to capture:
- Chat completions (sync and async)
- Function/tool calls
- 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")

# Track if already instrumented
_instrumented = False


def instrument_together(module: Any) -> None:
    """
    Apply instrumentation to Together AI module.

    Supports both v1.x (ChatCompletions) and v2.x (Completions) class structures
    via runtime detection.
    """
    global _instrumented
    if _instrumented:
        return

    try:
        # Detect class structure at runtime
        chat_mod = getattr(module, "resources", None)
        chat_mod = getattr(chat_mod, "chat", None) if chat_mod else None
        chat_mod = getattr(chat_mod, "completions", None) if chat_mod else None

        if chat_mod is None:
            logger.warning("Together AI SDK: cannot find chat.completions module")
            return

        # Detect class name at runtime — varies across Together SDK versions:
        #   v1.x: ChatCompletions / AsyncChatCompletions
        #   v2.0-2.1: Completions / AsyncCompletions (matches OpenAI/Groq)
        #   v2.2+: CompletionsResource / AsyncCompletionsResource
        if hasattr(chat_mod, "CompletionsResource"):
            sync_cls = "CompletionsResource"
            async_cls = "AsyncCompletionsResource" if hasattr(chat_mod, "AsyncCompletionsResource") else None
        elif hasattr(chat_mod, "Completions"):
            sync_cls = "Completions"
            async_cls = "AsyncCompletions" if hasattr(chat_mod, "AsyncCompletions") else None
        elif hasattr(chat_mod, "ChatCompletions"):
            sync_cls = "ChatCompletions"
            async_cls = "AsyncChatCompletions" if hasattr(chat_mod, "AsyncChatCompletions") else None
        else:
            logger.warning("Together AI SDK: unrecognized chat completions class structure")
            return

        # Patch sync chat completions
        wrapt.wrap_function_wrapper(
            module,
            f"resources.chat.completions.{sync_cls}.create",
            _wrap_chat_create,
        )

        # Patch async chat completions
        if async_cls:
            wrapt.wrap_function_wrapper(
                module,
                f"resources.chat.completions.{async_cls}.create",
                _wrap_async_chat_create,
            )

        _instrumented = True
        logger.debug("Instrumented Together AI SDK (class: %s)", sync_cls)

    except Exception as e:
        logger.warning(f"Failed to instrument Together 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."""
    return {
        "gen_ai.system": "together",
        "gen_ai.request.model": kwargs.get("model", "unknown"),
        "gen_ai.request.stream": kwargs.get("stream", False),
        "gen_ai.request.temperature": kwargs.get("temperature"),
        "gen_ai.request.max_tokens": kwargs.get("max_tokens"),
        "gen_ai.request.has_tools": kwargs.get("tools") is not None,
    }


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]):  # Limit to 10 messages
        span.set_attribute(f"gen_ai.prompt.{i}.role", msg.get("role", ""))
        content = msg.get("content", "")
        if isinstance(content, str):
            # Truncate long content
            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."""
    # Model info
    span.set_attribute("gen_ai.response.model", getattr(response, "model", "unknown"))
    span.set_attribute("gen_ai.response.id", getattr(response, "id", ""))
    span.set_attribute("gen_ai.latency_ms", latency_ms)

    # 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_choice = choices[0]
        message = getattr(first_choice, "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))
                for j, tc in enumerate(tool_calls[:5]):  # Limit to 5
                    span.set_attribute(
                        f"gen_ai.completion.tool_call.{j}.name",
                        getattr(tc.function, "name", ""),
                    )

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


def _wrap_chat_create(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for together.chat.completions.create (sync)."""
    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", "unknown")
    trace_content = _should_trace_content()
    stream = kwargs.get("stream", False)

    # NON-STREAMING: use context manager (span auto-ended on block exit)
    if not stream:
        with tracer.start_span(
            name=f"together.chat.completions.create/{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

    # STREAMING: manual span lifecycle — span.end() deferred to generator finally
    span = tracer.create_span(
        name=f"together.chat.completions.create/{model}",
        kind=SpanKind.LLM_CALL,
        attributes=_extract_model_info(kwargs),
    )
    messages = kwargs.get("messages", [])
    _capture_messages(span, messages, trace_content)

    try:
        start_time = time.perf_counter()
        result = wrapped(*args, **kwargs)
        return _wrap_stream_sync(result, span, start_time, trace_content, tracer)
    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]))
        span.end()
        tracer.export_span(span)
        raise


async def _wrap_async_chat_create(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for together.chat.completions.create (async)."""
    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", "unknown")
    trace_content = _should_trace_content()
    stream = kwargs.get("stream", False)

    # NON-STREAMING: use context manager (span auto-ended on block exit)
    if not stream:
        with tracer.start_span(
            name=f"together.chat.completions.create/{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

    # STREAMING: manual span lifecycle — span.end() deferred to generator finally
    span = tracer.create_span(
        name=f"together.chat.completions.create/{model}",
        kind=SpanKind.LLM_CALL,
        attributes=_extract_model_info(kwargs),
    )
    messages = kwargs.get("messages", [])
    _capture_messages(span, messages, trace_content)

    try:
        start_time = time.perf_counter()
        result = await wrapped(*args, **kwargs)
        return _wrap_stream_async(result, span, start_time, trace_content, tracer)
    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]))
        span.end()
        tracer.export_span(span)
        raise


def _wrap_stream_sync(
    stream: Iterator[Any],
    span: Any,
    start_time: float,
    trace_content: bool,
    tracer: Any,
) -> Iterator[Any]:
    """Wrap a sync streaming response."""
    chunk_count = 0
    content_buffer = []
    content_length = 0
    usage_data = {}

    try:
        for chunk in stream:
            chunk_count += 1

            # Accumulate content (bounded)
            if trace_content and content_length < 10000:
                delta = getattr(getattr(chunk, "choices", [{}])[0] if chunk.choices else None, "delta", None)
                if delta:
                    content = getattr(delta, "content", "")
                    if content:
                        content_buffer.append(content)
                        content_length += len(content)

            # Capture usage from final chunk
            usage = getattr(chunk, "usage", None)
            if usage:
                usage_data = {
                    "prompt_tokens": getattr(usage, "prompt_tokens", 0),
                    "completion_tokens": getattr(usage, "completion_tokens", 0),
                    "total_tokens": getattr(usage, "total_tokens", 0),
                }

            yield chunk

    except Exception as e:
        span.record_exception(e)
        span.set_attribute("error", True)
        span.set_attribute("error.type", type(e).__name__)
        raise
    finally:
        latency_ms = (time.perf_counter() - start_time) * 1000
        span.set_attribute("gen_ai.latency_ms", latency_ms)
        span.set_attribute("gen_ai.response.chunks", chunk_count)
        span.set_attribute("gen_ai.response.stream", True)

        if usage_data:
            span.set_attribute("gen_ai.usage.prompt_tokens", usage_data.get("prompt_tokens", 0))
            span.set_attribute("gen_ai.usage.completion_tokens", usage_data.get("completion_tokens", 0))
            span.set_attribute("gen_ai.usage.total_tokens", usage_data.get("total_tokens", 0))

        if trace_content and content_buffer:
            full_content = "".join(content_buffer)
            span.set_attribute(
                "gen_ai.completion.content",
                full_content[:10000] if len(full_content) > 10000 else full_content,
            )

        from risicare_core import SpanStatus

        if span.status == SpanStatus.UNSET:
            span.mark_ok()
        span.end()
        tracer.export_span(span)


async def _wrap_stream_async(
    stream: AsyncIterator[Any],
    span: Any,
    start_time: float,
    trace_content: bool,
    tracer: Any,
) -> AsyncIterator[Any]:
    """Wrap an async streaming response."""
    chunk_count = 0
    content_buffer = []
    content_length = 0
    usage_data = {}

    try:
        async for chunk in stream:
            chunk_count += 1

            # Accumulate content (bounded)
            if trace_content and content_length < 10000:
                choices = getattr(chunk, "choices", [])
                if choices:
                    delta = getattr(choices[0], "delta", None)
                    if delta:
                        content = getattr(delta, "content", "")
                        if content:
                            content_buffer.append(content)
                            content_length += len(content)

            # Capture usage from final chunk
            usage = getattr(chunk, "usage", None)
            if usage:
                usage_data = {
                    "prompt_tokens": getattr(usage, "prompt_tokens", 0),
                    "completion_tokens": getattr(usage, "completion_tokens", 0),
                    "total_tokens": getattr(usage, "total_tokens", 0),
                }

            yield chunk

    except Exception as e:
        span.record_exception(e)
        span.set_attribute("error", True)
        span.set_attribute("error.type", type(e).__name__)
        raise
    finally:
        latency_ms = (time.perf_counter() - start_time) * 1000
        span.set_attribute("gen_ai.latency_ms", latency_ms)
        span.set_attribute("gen_ai.response.chunks", chunk_count)
        span.set_attribute("gen_ai.response.stream", True)

        if usage_data:
            span.set_attribute("gen_ai.usage.prompt_tokens", usage_data.get("prompt_tokens", 0))
            span.set_attribute("gen_ai.usage.completion_tokens", usage_data.get("completion_tokens", 0))
            span.set_attribute("gen_ai.usage.total_tokens", usage_data.get("total_tokens", 0))

        if trace_content and content_buffer:
            full_content = "".join(content_buffer)
            span.set_attribute(
                "gen_ai.completion.content",
                full_content[:10000] if len(full_content) > 10000 else full_content,
            )

        from risicare_core import SpanStatus

        if span.status == SpanStatus.UNSET:
            span.mark_ok()
        span.end()
        tracer.export_span(span)
