"""Obol interceptor for OpenAI SDK.

Patches chat.completions.create() to capture tool/function calls.

Usage:
    from obol.interceptors import openai_sdk
    openai_sdk.install()

    # Then use OpenAI normally - all tool calls are logged
"""
import functools
from obol.interceptors._base import send_event, _content_hash

_installed = False
_original_create = None
_original_acreate = None


def _extract_tool_calls(response):
    """Extract tool calls from OpenAI response."""
    tool_calls = []

    # Handle different response formats
    if hasattr(response, "choices") and response.choices:
        choice = response.choices[0]
        message = getattr(choice, "message", None)

        if message:
            # Function calling (legacy)
            if hasattr(message, "function_call") and message.function_call:
                tool_calls.append({
                    "type": "function",
                    "name": message.function_call.name,
                    "arguments_hash": _content_hash(message.function_call.arguments),
                })

            # Tool calls (new format)
            if hasattr(message, "tool_calls") and message.tool_calls:
                for tc in message.tool_calls:
                    if hasattr(tc, "function"):
                        tool_calls.append({
                            "type": "tool",
                            "name": tc.function.name,
                            "arguments_hash": _content_hash(tc.function.arguments),
                        })

    return tool_calls


def _count_messages(messages):
    """Count messages by role."""
    if not messages:
        return {}
    counts = {}
    for msg in messages:
        role = msg.get("role", "unknown") if isinstance(msg, dict) else getattr(msg, "role", "unknown")
        counts[role] = counts.get(role, 0) + 1
    return counts


def _traced_create(original):
    """Create a traced version of chat.completions.create."""
    @functools.wraps(original)
    def wrapper(self, *args, **kwargs):
        model = kwargs.get("model", "unknown")
        messages = kwargs.get("messages", [])
        tools = kwargs.get("tools", [])
        functions = kwargs.get("functions", [])  # Legacy format

        # Log request start
        send_event(
            platform="openai",
            action="ChatCompletion.create",
            target=model,
            result="pending",
            metadata={
                "model": model,
                "message_counts": _count_messages(messages),
                "tool_count": len(tools) + len(functions),
                "has_tools": bool(tools or functions),
            }
        )

        try:
            response = original(self, *args, **kwargs)

            # Extract tool calls from response
            tool_calls = _extract_tool_calls(response)

            # Log completion
            send_event(
                platform="openai",
                action="ChatCompletion.create",
                target=model,
                result="success",
                metadata={
                    "model": model,
                    "tool_calls": tool_calls,
                    "tool_call_count": len(tool_calls),
                    "finish_reason": response.choices[0].finish_reason if response.choices else None,
                    "usage": {
                        "prompt_tokens": getattr(response.usage, "prompt_tokens", 0) if response.usage else 0,
                        "completion_tokens": getattr(response.usage, "completion_tokens", 0) if response.usage else 0,
                    } if hasattr(response, "usage") else None,
                }
            )

            # Log individual tool calls as separate events
            for tc in tool_calls:
                send_event(
                    platform="openai",
                    action=f"ToolCall.{tc['name']}",
                    target=tc["arguments_hash"],
                    result="invoked",
                    metadata={
                        "model": model,
                        "call_type": tc["type"],
                    }
                )

            return response

        except Exception as e:
            # Log error
            send_event(
                platform="openai",
                action="ChatCompletion.create",
                target=model,
                result="error",
                metadata={
                    "model": model,
                    "error_type": type(e).__name__,
                    "error_message": str(e)[:500],
                }
            )
            raise

    return wrapper


def _traced_acreate(original):
    """Create a traced version of chat.completions.create (async)."""
    @functools.wraps(original)
    async def wrapper(self, *args, **kwargs):
        model = kwargs.get("model", "unknown")
        messages = kwargs.get("messages", [])
        tools = kwargs.get("tools", [])
        functions = kwargs.get("functions", [])

        # Log request start
        send_event(
            platform="openai",
            action="ChatCompletion.create",
            target=model,
            result="pending",
            metadata={
                "model": model,
                "message_counts": _count_messages(messages),
                "tool_count": len(tools) + len(functions),
                "has_tools": bool(tools or functions),
                "async": True,
            }
        )

        try:
            response = await original(self, *args, **kwargs)

            # Extract tool calls from response
            tool_calls = _extract_tool_calls(response)

            # Log completion
            send_event(
                platform="openai",
                action="ChatCompletion.create",
                target=model,
                result="success",
                metadata={
                    "model": model,
                    "tool_calls": tool_calls,
                    "tool_call_count": len(tool_calls),
                    "finish_reason": response.choices[0].finish_reason if response.choices else None,
                    "async": True,
                }
            )

            # Log individual tool calls
            for tc in tool_calls:
                send_event(
                    platform="openai",
                    action=f"ToolCall.{tc['name']}",
                    target=tc["arguments_hash"],
                    result="invoked",
                    metadata={
                        "model": model,
                        "call_type": tc["type"],
                        "async": True,
                    }
                )

            return response

        except Exception as e:
            send_event(
                platform="openai",
                action="ChatCompletion.create",
                target=model,
                result="error",
                metadata={
                    "model": model,
                    "error_type": type(e).__name__,
                    "error_message": str(e)[:500],
                    "async": True,
                }
            )
            raise

    return wrapper


def install():
    """Install the OpenAI SDK interceptor.

    Patches chat.completions.create() to capture all LLM tool calls.
    Safe to call multiple times.
    """
    global _installed, _original_create, _original_acreate

    if _installed:
        return

    try:
        from openai import OpenAI
        from openai.resources.chat import Completions
    except ImportError:
        return  # OpenAI not installed

    # Patch sync create
    _original_create = Completions.create
    Completions.create = _traced_create(_original_create)

    # Patch async create if available
    try:
        from openai.resources.chat import AsyncCompletions
        _original_acreate = AsyncCompletions.create
        AsyncCompletions.create = _traced_acreate(_original_acreate)
    except (ImportError, AttributeError):
        pass

    _installed = True


def uninstall():
    """Remove the OpenAI SDK interceptor."""
    global _installed, _original_create, _original_acreate

    if not _installed:
        return

    try:
        from openai.resources.chat import Completions
        if _original_create:
            Completions.create = _original_create
    except ImportError:
        pass

    try:
        from openai.resources.chat import AsyncCompletions
        if _original_acreate:
            AsyncCompletions.create = _original_acreate
    except (ImportError, AttributeError):
        pass

    _installed = False


def is_installed():
    """Check if the interceptor is currently installed."""
    return _installed
