"""
Instructor patches for from_provider(), patch(), and Instructor.create().

Wraps the Instructor factory functions to auto-register Risicare tracking,
and wraps Instructor.create() to create extraction spans around structured
output attempts.

Does NOT suppress provider instrumentation — Instructor delegates to the
underlying LLM client (OpenAI, Anthropic, etc.), so provider patches handle
LLM spans naturally as children.
"""

from __future__ import annotations

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

import wrapt

from risicare.integrations._base import (
    create_framework_attributes,
    get_framework_version,
    get_tracer,
    is_tracing_enabled,
    record_error,
    safe_set_attribute,
    should_trace_content,
    truncate_content,
)

logger = logging.getLogger(__name__)

T = TypeVar("T")


def _extract_response_model_name(kwargs: Dict[str, Any]) -> str:
    """Extract the response model class name from create() kwargs."""
    response_model = kwargs.get("response_model")
    if response_model is None:
        return "unknown"
    return getattr(response_model, "__name__", str(response_model))


def _extract_mode(instance: Any) -> str:
    """Extract the extraction mode from an Instructor client."""
    mode = getattr(instance, "mode", None)
    if mode is None:
        return "unknown"
    # Mode is typically an enum with a .value
    return str(getattr(mode, "value", mode))


def _wrap_create(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for Instructor.create() (sync)."""
    if not is_tracing_enabled():
        return wrapped(*args, **kwargs)

    tracer = get_tracer()
    if tracer is None:
        return wrapped(*args, **kwargs)

    from risicare_core import SpanKind

    model_name = _extract_response_model_name(kwargs)
    version = get_framework_version("instructor")

    attrs = create_framework_attributes("instructor", version)
    attrs["framework.span_kind"] = "extraction"
    attrs["framework.instructor.response_model"] = model_name
    attrs["framework.instructor.mode"] = _extract_mode(instance)

    max_retries = kwargs.get("max_retries", getattr(instance, "max_retries", None))
    if max_retries is not None:
        safe_set_attribute(None, "", None)  # no-op, just for import validation
        attrs["framework.instructor.max_retries"] = max_retries

    with tracer.start_span(
        name=f"instructor.create/{model_name}",
        kind=SpanKind.INTERNAL,
        attributes=attrs,
    ) as span:
        try:
            start_time = time.perf_counter()
            result = wrapped(*args, **kwargs)
            latency_ms = (time.perf_counter() - start_time) * 1000
            span.set_attribute("gen_ai.latency_ms", latency_ms)

            if should_trace_content() and result is not None:
                result_str = str(result)
                safe_set_attribute(
                    span,
                    "gen_ai.response.content",
                    truncate_content(result_str, 5000),
                )

            return result
        except Exception as e:
            record_error(span, e)
            raise


async def _wrap_create_async(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """Wrapper for AsyncInstructor.create() (async)."""
    if not is_tracing_enabled():
        return await wrapped(*args, **kwargs)

    tracer = get_tracer()
    if tracer is None:
        return await wrapped(*args, **kwargs)

    from risicare_core import SpanKind

    model_name = _extract_response_model_name(kwargs)
    version = get_framework_version("instructor")

    attrs = create_framework_attributes("instructor", version)
    attrs["framework.span_kind"] = "extraction"
    attrs["framework.instructor.response_model"] = model_name
    attrs["framework.instructor.mode"] = _extract_mode(instance)

    max_retries = kwargs.get("max_retries", getattr(instance, "max_retries", None))
    if max_retries is not None:
        attrs["framework.instructor.max_retries"] = max_retries

    with tracer.start_span(
        name=f"instructor.create/{model_name}",
        kind=SpanKind.INTERNAL,
        attributes=attrs,
    ) as span:
        try:
            start_time = time.perf_counter()
            result = await wrapped(*args, **kwargs)
            latency_ms = (time.perf_counter() - start_time) * 1000
            span.set_attribute("gen_ai.latency_ms", latency_ms)

            if should_trace_content() and result is not None:
                result_str = str(result)
                safe_set_attribute(
                    span,
                    "gen_ai.response.content",
                    truncate_content(result_str, 5000),
                )

            return result
        except Exception as e:
            record_error(span, e)
            raise


def _wrap_from_provider(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """
    Wrapper for instructor.from_provider().

    Calls the original factory, then wraps the returned client's create()
    method if not already wrapped.
    """
    client = wrapped(*args, **kwargs)

    # Mark the client so we don't double-wrap
    if getattr(client, "_risicare_instrumented", False):
        return client

    try:
        # Wrap create() on the returned client instance
        original_create = client.create

        def instrumented_create(*a: Any, **kw: Any) -> Any:
            return _wrap_create(original_create, client, a, kw)

        client.create = instrumented_create

        # Wrap create_partial if it exists
        original_create_partial = getattr(client, "create_partial", None)
        if original_create_partial is not None:
            def instrumented_create_partial(*a: Any, **kw: Any) -> Any:
                return _wrap_create(original_create_partial, client, a, kw)
            client.create_partial = instrumented_create_partial

        client._risicare_instrumented = True
    except Exception as e:
        logger.debug(f"Failed to wrap Instructor client methods: {e}")

    return client


def _wrap_patch_legacy(
    wrapped: Callable[..., T],
    instance: Any,
    args: tuple,
    kwargs: Dict[str, Any],
) -> T:
    """
    Wrapper for instructor.patch() (legacy API).

    Same approach as from_provider — calls original, wraps returned client.
    """
    client = wrapped(*args, **kwargs)

    if getattr(client, "_risicare_instrumented", False):
        return client

    try:
        original_create = getattr(client, "create", None)
        if original_create is not None:
            def instrumented_create(*a: Any, **kw: Any) -> Any:
                return _wrap_create(original_create, client, a, kw)
            client.create = instrumented_create
            client._risicare_instrumented = True
    except Exception as e:
        logger.debug(f"Failed to wrap legacy patched client: {e}")

    return client


def patch_instructor(module: Any) -> None:
    """Apply wrapt patches to Instructor module."""
    # Patch from_provider (modern API)
    try:
        wrapt.wrap_function_wrapper(module, "from_provider", _wrap_from_provider)
    except Exception as e:
        logger.debug(f"Failed to patch instructor.from_provider: {e}")

    # Patch legacy patch() function
    try:
        if hasattr(module, "patch"):
            wrapt.wrap_function_wrapper(module, "patch", _wrap_patch_legacy)
    except Exception as e:
        logger.debug(f"Failed to patch instructor.patch: {e}")

    # Patch Instructor.create directly for users who construct clients manually
    try:
        wrapt.wrap_function_wrapper(
            "instructor.client", "Instructor.create", _wrap_create
        )
    except Exception as e:
        logger.debug(f"Failed to patch Instructor.create: {e}")

    # Patch AsyncInstructor.create
    try:
        wrapt.wrap_function_wrapper(
            "instructor.client", "AsyncInstructor.create", _wrap_create_async
        )
    except Exception as e:
        logger.debug(f"Failed to patch AsyncInstructor.create: {e}")
