"""
Azure OpenAI client wrapper for Raindrop observability.

Wraps an AzureOpenAI client to capture chat.completions.create() calls.

Usage:
    from raindrop_azure_openai import create_raindrop_azure_openai
    from openai import AzureOpenAI

    raindrop = create_raindrop_azure_openai(api_key="rk_...")
    client = AzureOpenAI(azure_endpoint="...", api_key="...", api_version="...")
    wrapped = raindrop["wrap"](client)

    response = wrapped.chat.completions.create(model="gpt-4o", messages=[...])
    raindrop["flush"]()
"""

from __future__ import annotations

import json
import uuid
from typing import Any, Dict, Optional

import raindrop.analytics as raindrop


def create_raindrop_azure_openai(
    api_key: str,
    user_id: Optional[str] = None,
    convo_id: Optional[str] = None,
) -> dict:
    """Create a Raindrop-instrumented Azure OpenAI wrapper."""
    raindrop.init(api_key=api_key)

    def wrap(client: Any) -> Any:
        return wrap_azure_client(client, user_id=user_id, convo_id=convo_id)

    return {
        "wrap": wrap,
        "flush": raindrop.flush,
        "shutdown": raindrop.shutdown,
    }


def wrap_azure_client(
    client: Any,
    user_id: Optional[str] = None,
    convo_id: Optional[str] = None,
) -> Any:
    """Wrap an AzureOpenAI client to capture chat completions."""
    try:
        completions = client.chat.completions
        original_create = completions.create

        def wrapped_create(**kwargs: Any) -> Any:
            return _instrument_create(original_create, kwargs, user_id, convo_id)

        completions.create = wrapped_create
    except Exception:
        pass

    return client


def _instrument_create(
    original_fn: Any,
    kwargs: Dict[str, Any],
    user_id: Optional[str],
    convo_id: Optional[str],
) -> Any:
    try:
        model_id = kwargs.get("model")
        input_text = _extract_input(kwargs.get("messages"))
        event_id = str(uuid.uuid4())
    except Exception:
        return original_fn(**kwargs)

    try:
        response = original_fn(**kwargs)
    except Exception:
        try:
            raindrop.track_ai(
                user_id=user_id or "unknown",
                event="ai_generation",
                event_id=event_id,
                input=input_text,
                model=model_id,
                convo_id=convo_id,
            )
        except Exception:
            pass
        raise

    try:
        output_text = _extract_output(response)
        response_model = getattr(response, "model", model_id)
        usage = getattr(response, "usage", None)
        prompt_tokens = getattr(usage, "prompt_tokens", None) if usage else None
        completion_tokens = getattr(usage, "completion_tokens", None) if usage else None

        properties: Dict[str, Any] = {}
        if prompt_tokens is not None:
            properties["ai.usage.prompt_tokens"] = prompt_tokens
        if completion_tokens is not None:
            properties["ai.usage.completion_tokens"] = completion_tokens

        raindrop.track_ai(
            user_id=user_id or "unknown",
            event="ai_generation",
            event_id=event_id,
            input=input_text,
            output=output_text,
            model=response_model,
            properties=properties or None,
            convo_id=convo_id,
        )
    except Exception:
        pass

    return response


def _extract_input(messages: Any) -> Optional[str]:
    if not isinstance(messages, list):
        return None
    try:
        user_msgs = [m for m in messages if isinstance(m, dict) and m.get("role") == "user"]
        last_msg = user_msgs[-1] if user_msgs else (messages[-1] if messages else None)
        if not last_msg or not isinstance(last_msg, dict):
            return None
        content = last_msg.get("content", "")
        if isinstance(content, str):
            return content or None
        try:
            return json.dumps(content, default=str)
        except Exception:
            return str(content)
    except Exception:
        return None


def _extract_output(response: Any) -> Optional[str]:
    try:
        choices = getattr(response, "choices", None)
        if choices and len(choices) > 0:
            message = getattr(choices[0], "message", None)
            if message:
                return getattr(message, "content", None)
    except Exception:
        pass
    return None
