"""
Hugging Face Inference Client instrumentation.

Patches huggingface_hub InferenceClient and AsyncInferenceClient methods to capture:
- Chat completions (sync and async)
- Text generation (sync and async)
- Streaming responses

Based on ARCHITECTURE_V2.md Section 6.3.
Reference: https://huggingface.co/docs/huggingface_hub/en/package_reference/inference_client
"""

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_huggingface(module: Any) -> None:
    """
    Apply instrumentation to huggingface_hub module.

    This patches:
    - huggingface_hub.InferenceClient.chat_completion
    - huggingface_hub.AsyncInferenceClient.chat_completion
    - huggingface_hub.InferenceClient.text_generation
    - huggingface_hub.AsyncInferenceClient.text_generation
    """
    global _instrumented
    if _instrumented:
        return

    try:
        # Patch sync chat completion
        wrapt.wrap_function_wrapper(
            module,
            "InferenceClient.chat_completion",
            _wrap_chat_completion,
        )

        # Patch async chat completion
        wrapt.wrap_function_wrapper(
            module,
            "AsyncInferenceClient.chat_completion",
            _wrap_async_chat_completion,
        )

        # Patch sync text generation
        wrapt.wrap_function_wrapper(
            module,
            "InferenceClient.text_generation",
            _wrap_text_generation,
        )

        # Patch async text generation
        wrapt.wrap_function_wrapper(
            module,
            "AsyncInferenceClient.text_generation",
            _wrap_async_text_generation,
        )

        _instrumented = True
        logger.debug("Instrumented Hugging Face InferenceClient")

    except Exception as e:
        logger.warning(f"Failed to instrument Hugging Face InferenceClient: {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 _get_model_from_kwargs_or_instance(kwargs: Dict[str, Any], instance: Any) -> str:
    """Extract model name from kwargs or InferenceClient.model attribute."""
    model = kwargs.get("model")
    if model:
        return model
    # InferenceClient stores default model in self.model
    return getattr(instance, "model", None) or "unknown"


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

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

    stream = kwargs.get("stream", False)
    temperature = kwargs.get("temperature")
    max_tokens = kwargs.get("max_tokens")
    tools = kwargs.get("tools")

    info: Dict[str, Any] = {
        "gen_ai.system": "huggingface",
        "gen_ai.request.model": _get_model_from_kwargs_or_instance(kwargs, instance),
        "gen_ai.request.has_tools": isinstance(tools, (list, tuple)) and len(tools) > 0,
    }

    if is_json_primitive(stream):
        info["gen_ai.request.stream"] = stream
    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 _extract_text_gen_model_info(kwargs: Dict[str, Any], instance: Any) -> Dict[str, Any]:
    """Extract model information from text_generation request kwargs.

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

    stream = kwargs.get("stream", False)
    temperature = kwargs.get("temperature")
    max_tokens = kwargs.get("max_new_tokens", 100)

    info: Dict[str, Any] = {
        "gen_ai.system": "huggingface",
        "gen_ai.request.model": _get_model_from_kwargs_or_instance(kwargs, instance),
    }

    if is_json_primitive(stream):
        info["gen_ai.request.stream"] = stream
    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]):
        if isinstance(msg, dict):
            span.set_attribute(f"gen_ai.prompt.{i}.role", msg.get("role", ""))
            content = msg.get("content", "")
        else:
            # ChatCompletionInputMessage dataclass
            span.set_attribute(f"gen_ai.prompt.{i}.role", getattr(msg, "role", ""))
            content = 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_chat_response(
    span: Any,
    response: Any,
    latency_ms: float,
    trace_content: bool,
) -> None:
    """Capture chat_completion response data into span.

    Response is a ChatCompletionOutput dataclass with:
    - choices: list[ChatCompletionOutputComplete]
    - usage: ChatCompletionOutputUsage (prompt_tokens, completion_tokens, total_tokens)
    - model: str
    - id: str
    """
    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 = 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]):
                    func = getattr(tc, "function", None)
                    if func:
                        span.set_attribute(
                            f"gen_ai.completion.tool_call.{j}.name",
                            getattr(func, "name", ""),
                        )

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


def _capture_text_gen_response(
    span: Any,
    response: Any,
    latency_ms: float,
    trace_content: bool,
) -> None:
    """Capture text_generation response data into span.

    Response can be:
    - str (default, no details)
    - TextGenerationOutput (with details=True)
    """
    span.set_attribute("gen_ai.latency_ms", latency_ms)

    if isinstance(response, str):
        # Simple string response (no details)
        if trace_content:
            span.set_attribute(
                "gen_ai.completion.content",
                response[:10000] if len(response) > 10000 else response,
            )
    else:
        # TextGenerationOutput with details
        generated_text = getattr(response, "generated_text", "")
        if trace_content and generated_text:
            span.set_attribute(
                "gen_ai.completion.content",
                generated_text[:10000] if len(generated_text) > 10000 else generated_text,
            )

        details = getattr(response, "details", None)
        if details:
            span.set_attribute(
                "gen_ai.usage.completion_tokens",
                getattr(details, "generated_tokens", 0),
            )
            span.set_attribute(
                "gen_ai.completion.finish_reason",
                str(getattr(details, "finish_reason", "")),
            )


# =============================================================================
# Chat Completion Wrappers
# =============================================================================


def _wrap_chat_completion(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for InferenceClient.chat_completion (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 = _get_model_from_kwargs_or_instance(kwargs, instance)
    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"huggingface.chat_completion/{model}",
            kind=SpanKind.LLM_CALL,
            attributes=_extract_chat_model_info(kwargs, instance),
        ) as span:
            # Messages can be in args[0] or kwargs["messages"]
            messages = kwargs.get("messages") or (args[0] if args else [])
            if isinstance(messages, list):
                _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_chat_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"huggingface.chat_completion/{model}",
        kind=SpanKind.LLM_CALL,
        attributes=_extract_chat_model_info(kwargs, instance),
    )
    messages = kwargs.get("messages") or (args[0] if args else [])
    if isinstance(messages, list):
        _capture_messages(span, messages, trace_content)

    try:
        start_time = time.perf_counter()
        result = wrapped(*args, **kwargs)
        return _wrap_chat_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_completion(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for AsyncInferenceClient.chat_completion (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 = _get_model_from_kwargs_or_instance(kwargs, instance)
    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"huggingface.chat_completion/{model}",
            kind=SpanKind.LLM_CALL,
            attributes=_extract_chat_model_info(kwargs, instance),
        ) as span:
            messages = kwargs.get("messages") or (args[0] if args else [])
            if isinstance(messages, list):
                _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_chat_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"huggingface.chat_completion/{model}",
        kind=SpanKind.LLM_CALL,
        attributes=_extract_chat_model_info(kwargs, instance),
    )
    messages = kwargs.get("messages") or (args[0] if args else [])
    if isinstance(messages, list):
        _capture_messages(span, messages, trace_content)

    try:
        start_time = time.perf_counter()
        result = await wrapped(*args, **kwargs)
        return _wrap_chat_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


# =============================================================================
# Text Generation Wrappers
# =============================================================================


def _wrap_text_generation(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for InferenceClient.text_generation (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 = _get_model_from_kwargs_or_instance(kwargs, instance)
    trace_content = _should_trace_content()

    with tracer.start_span(
        name=f"huggingface.text_generation/{model}",
        kind=SpanKind.LLM_CALL,
        attributes=_extract_text_gen_model_info(kwargs, instance),
    ) as span:
        # Capture prompt (first positional arg or kwargs["prompt"])
        prompt = kwargs.get("prompt") or (args[0] if args else "")
        if trace_content and isinstance(prompt, str) and prompt:
            span.set_attribute(
                "gen_ai.prompt.0.content",
                prompt[:10000] if len(prompt) > 10000 else prompt,
            )

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

            _capture_text_gen_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_text_generation(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for AsyncInferenceClient.text_generation (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 = _get_model_from_kwargs_or_instance(kwargs, instance)
    trace_content = _should_trace_content()

    with tracer.start_span(
        name=f"huggingface.text_generation/{model}",
        kind=SpanKind.LLM_CALL,
        attributes=_extract_text_gen_model_info(kwargs, instance),
    ) as span:
        prompt = kwargs.get("prompt") or (args[0] if args else "")
        if trace_content and isinstance(prompt, str) and prompt:
            span.set_attribute(
                "gen_ai.prompt.0.content",
                prompt[:10000] if len(prompt) > 10000 else prompt,
            )

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

            _capture_text_gen_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 Wrappers
# =============================================================================


def _wrap_chat_stream_sync(
    stream: Iterator[Any],
    span: Any,
    start_time: float,
    trace_content: bool,
    tracer: Any,
) -> Iterator[Any]:
    """Wrap a sync chat streaming response (Iterable[ChatCompletionStreamOutput])."""
    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:
                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)

            # Usage may appear in 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_chat_stream_async(
    stream: AsyncIterator[Any],
    span: Any,
    start_time: float,
    trace_content: bool,
    tracer: Any,
) -> AsyncIterator[Any]:
    """Wrap an async chat streaming response (AsyncIterable[ChatCompletionStreamOutput])."""
    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)

            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)
