"""OpenAI integration with auto-instrumentation via monkey patching."""

import functools
import logging
from typing import Any

from fenra._context import get_context
from fenra._core import enqueue_transaction

logger = logging.getLogger(__name__)

_patched_chat_sync = False
_patched_chat_async = False
_patched_responses_sync = False
_patched_responses_async = False
_patched_images_sync = False
_patched_images_async = False
_patched_beta_parse_sync = False
_patched_beta_parse_async = False

# Store original methods for unpatching
_original_chat_create: Any = None
_original_chat_create_async: Any = None
_original_responses_create: Any = None
_original_responses_create_async: Any = None
_original_images_generate: Any = None
_original_images_generate_async: Any = None
_original_beta_parse: Any = None
_original_beta_parse_async: Any = None

def _is_openai_endpoint(client: Any) -> bool:
    """
    Check if client is actually hitting OpenAI, not other OpenAI-compatible endpoints.
    """

    base_url = str(getattr(client, "base_url", ""))
    # OpenAI domains
    if "api.openai.com" in base_url:
        return True
    # Azure OpenAI
    if "openai.azure.com" in base_url:
        return True
    if "generativelanguage.googleapis.com" in base_url:
        return False
    # Default to False for custom models that use OpenAI-compatible endpoints
    return False


def _track_chat_completion_response(
    response: Any,
    model: str,
    context: dict[str, Any],
) -> None:
    """Extract usage from Chat Completions response and queue transaction with full context."""
    if not hasattr(response, "usage") or response.usage is None:
        return

    usage = response.usage

    metrics: dict[str, Any] = {
        "input_tokens": getattr(usage, "prompt_tokens", 0),
        "output_tokens": getattr(usage, "completion_tokens", 0),
        "total_tokens": getattr(usage, "total_tokens", 0),
    }

    # Include reasoning tokens if present (o1/o3 models)
    if hasattr(usage, "completion_tokens_details") and usage.completion_tokens_details:
        details = usage.completion_tokens_details
        if hasattr(details, "reasoning_tokens") and details.reasoning_tokens:
            metrics["reasoning_tokens"] = details.reasoning_tokens

    # Include cached tokens if present
    if hasattr(usage, "prompt_tokens_details") and usage.prompt_tokens_details:
        details = usage.prompt_tokens_details
        if hasattr(details, "cached_tokens") and details.cached_tokens:
            metrics["cached_tokens"] = details.cached_tokens

    transaction = {
        "provider": "openai",
        "model": model,
        "usage": [
            {
                "type": "tokens",
                "metrics": metrics,
            }
        ],
        "context": context,  # Full merged context from set_context() calls
    }

    try:
        transaction["provider_usage_raw"] = {
            "prompt_tokens": getattr(usage, "prompt_tokens", None),
            "completion_tokens": getattr(usage, "completion_tokens", None),
            "total_tokens": getattr(usage, "total_tokens", None),
        }
    except Exception:
        pass

    enqueue_transaction(transaction)


def _track_responses_api_response(
    response: Any,
    model: str,
    context: dict[str, Any],
    tools: list[dict[str, Any]] | None = None,
) -> None:
    """Extract usage from Responses API response and queue transaction with full context."""
    if not hasattr(response, "usage") or response.usage is None:
        return

    usage = response.usage

    # Build token metrics (Responses API uses different field names)
    metrics: dict[str, Any] = {
        "input_tokens": getattr(usage, "input_tokens", 0),
        "output_tokens": getattr(usage, "output_tokens", 0),
        "total_tokens": getattr(usage, "total_tokens", 0),
    }

    # Include reasoning tokens if present (o1/o3 models)
    if hasattr(usage, "output_tokens_details") and usage.output_tokens_details:
        details = usage.output_tokens_details
        if hasattr(details, "reasoning_tokens") and details.reasoning_tokens:
            metrics["reasoning_tokens"] = details.reasoning_tokens

    # Include cached tokens if present
    if hasattr(usage, "input_tokens_details") and usage.input_tokens_details:
        details = usage.input_tokens_details
        if hasattr(details, "cached_tokens") and details.cached_tokens:
            metrics["cached_tokens"] = details.cached_tokens

    usage_entries: list[dict[str, Any]] = [
        {
            "type": "tokens",
            "metrics": metrics,
        }
    ]

    # Add requests usage when web_search tool is used
    if tools and any(t.get("type") in ("web_search", "web_search_preview") for t in tools):
        usage_entries.append(
            {
                "type": "requests",
                "metrics": {"count": 1, "request_type": "web_search"},
            }
        )

    transaction: dict[str, Any] = {
        "provider": "openai",
        "model": model,
        "usage": usage_entries,
        "context": context,
    }

    # Include raw usage for debugging if available
    try:
        transaction["provider_usage_raw"] = {
            "input_tokens": getattr(usage, "input_tokens", None),
            "output_tokens": getattr(usage, "output_tokens", None),
            "total_tokens": getattr(usage, "total_tokens", None),
        }
    except Exception:
        pass

    enqueue_transaction(transaction)


def _track_images_response(
    model: str,
    n: int,
    size: str,
    quality: str,
    context: dict[str, Any],
    response: Any = None,
) -> None:
    """Track image generation usage and queue transaction with full context.

    Prefers quality/size/n from the API response when available (reflects actual
    usage); otherwise uses request kwargs.
    """
    # Prefer response when available (API returns what was actually used)
    if response is not None:
        quality = getattr(response, "quality", None) or quality
        size = getattr(response, "size", None) or size
        if hasattr(response, "data") and response.data is not None:
            n = len(response.data)

    # Parse size to get pixels (e.g., "1024x1024" -> 1024)
    size_px: int | None = None
    if size and "x" in str(size):
        try:
            size_px = int(str(size).split("x")[0])
        except (ValueError, TypeError):
            pass

    metrics: dict[str, Any] = {
        "generated": n,
        "size_px": size_px,
        "quality": quality or None,
    }
    # Omit None values for API payload
    metrics = {k: v for k, v in metrics.items() if v is not None}

    transaction = {
        "provider": "openai",
        "model": model,
        "usage": [
            {
                "type": "images",
                "metrics": metrics,
            }
        ],
        "context": context,
    }
    enqueue_transaction(transaction)


def _track_openai_stream(
    stream: Any,
    model: str,
    context: dict[str, Any],
    is_responses_api: bool = False,
    tools: list[dict[str, Any]] | None = None,
) -> Any:
    """
    Wrap a streaming response to track usage from the final chunk.

    Returns a wrapped stream that tracks usage when the stream is exhausted.
    """
    last_chunk = None

    def wrapped_stream() -> Any:
        nonlocal last_chunk
        try:
            for chunk in stream:
                last_chunk = chunk
                yield chunk
        finally:
            # After stream is exhausted, track usage from final chunk
            if last_chunk and hasattr(last_chunk, "usage") and last_chunk.usage:
                try:
                    if is_responses_api:
                        _track_responses_api_response(
                            last_chunk, model, context, tools=tools
                        )
                    else:
                        _track_chat_completion_response(last_chunk, model, context)
                except Exception as e:
                    logger.error(
                        f"Error tracking OpenAI streaming usage: {e}", exc_info=True
                    )

    return wrapped_stream()


async def _track_openai_stream_async(
    stream: Any,
    model: str,
    context: dict[str, Any],
    is_responses_api: bool = False,
    tools: list[dict[str, Any]] | None = None,
) -> Any:
    """
    Wrap an async streaming response to track usage from the final chunk.

    Returns a wrapped async stream that tracks usage when the stream is exhausted.
    """
    last_chunk = None

    try:
        async for chunk in stream:
            last_chunk = chunk
            yield chunk
    finally:
        # After stream is exhausted, track usage from final chunk
        if last_chunk and hasattr(last_chunk, "usage") and last_chunk.usage:
            try:
                if is_responses_api:
                    _track_responses_api_response(
                        last_chunk, model, context, tools=tools
                    )
                else:
                    _track_chat_completion_response(last_chunk, model, context)
            except Exception as e:
                logger.error(
                    f"Error tracking OpenAI streaming usage: {e}", exc_info=True
                )


# ============================================================================
# Chat Completions API Patching
# ============================================================================

def patch_openai() -> None:
    """
    Patch OpenAI Chat Completions client to auto-track usage.

    This patches the synchronous `Completions.create` method.
    """
    global _patched_chat_sync, _original_chat_create

    if _patched_chat_sync:
        logger.debug("OpenAI Chat Completions sync already patched, skipping")
        return

    try:
        from openai.resources.chat.completions import Completions
    except ImportError:
        logger.warning(
            "OpenAI SDK not installed. Install with: pip install 'fenra[openai]'"
        )
        return

    original_create = Completions.create
    _original_chat_create = original_create

    @functools.wraps(original_create)
    def patched_create(self: Completions, *args: Any, **kwargs: Any) -> Any:
        if not _is_openai_endpoint(self._client):
            return original_create(self, *args, **kwargs)

        is_streaming = kwargs.get("stream", False)

        response = original_create(self, *args, **kwargs)

        context = get_context()

        model = kwargs.get("model") or getattr(response, "model", "unknown")

        try:
            if is_streaming:
                return _track_openai_stream(response, model, context, is_responses_api=False)
            else:
                _track_chat_completion_response(response, model, context)
        except Exception as e:
            logger.error(f"Error tracking OpenAI usage: {e}", exc_info=True)

        return response

    Completions.create = patched_create  # type: ignore[assignment]
    _patched_chat_sync = True
    logger.info("OpenAI Chat Completions SDK patched for auto-instrumentation")


def patch_openai_async() -> None:
    """
    Patch OpenAI Chat Completions async client to auto-track usage.

    This patches the asynchronous `AsyncCompletions.create` method.
    """
    global _patched_chat_async, _original_chat_create_async

    if _patched_chat_async:
        logger.debug("OpenAI Chat Completions async already patched, skipping")
        return

    try:
        from openai.resources.chat.completions import AsyncCompletions
    except ImportError:
        logger.warning(
            "OpenAI SDK not installed. Install with: pip install 'fenra[openai]'"
        )
        return

    original_create = AsyncCompletions.create
    _original_chat_create_async = original_create

    @functools.wraps(original_create)
    async def patched_create(
        self: AsyncCompletions, *args: Any, **kwargs: Any
    ) -> Any:
        if not _is_openai_endpoint(self._client):
            return await original_create(self, *args, **kwargs)

        is_streaming = kwargs.get("stream", False)

        response = await original_create(self, *args, **kwargs)

        context = get_context()

        model = kwargs.get("model") or getattr(response, "model", "unknown")

        try:
            if is_streaming:
                return _track_openai_stream_async(response, model, context, is_responses_api=False)
            else:
                _track_chat_completion_response(response, model, context)
        except Exception as e:
            logger.error(f"Error tracking OpenAI usage: {e}", exc_info=True)

        return response

    AsyncCompletions.create = patched_create  # type: ignore[assignment]
    _patched_chat_async = True
    logger.info("OpenAI Chat Completions async SDK patched for auto-instrumentation")


# ============================================================================
# Responses API Patching
# ============================================================================


def patch_openai_responses() -> None:
    """
    Patch OpenAI Responses API client to auto-track usage.

    This patches the synchronous `Responses.create` method.
    """
    global _patched_responses_sync, _original_responses_create

    if _patched_responses_sync:
        logger.debug("OpenAI Responses API sync already patched, skipping")
        return

    try:
        from openai.resources.responses import Responses
    except ImportError:
        logger.warning(
            "OpenAI SDK not installed. Install with: pip install 'fenra[openai]'"
        )
        return

    original_create = Responses.create
    _original_responses_create = original_create

    @functools.wraps(original_create)
    def patched_create(self: Responses, *args: Any, **kwargs: Any) -> Any:
        if not _is_openai_endpoint(self._client):
            return original_create(self, *args, **kwargs)

        is_streaming = kwargs.get("stream", False)

        response = original_create(self, *args, **kwargs)

        context = get_context()

        model = kwargs.get("model") or getattr(response, "model", "unknown")

        tools = kwargs.get("tools")
        try:
            if is_streaming:
                return _track_openai_stream(
                    response, model, context, is_responses_api=True, tools=tools
                )
            else:
                _track_responses_api_response(response, model, context, tools=tools)
        except Exception as e:
            logger.error(f"Error tracking OpenAI Responses API usage: {e}", exc_info=True)

        return response

    Responses.create = patched_create  # type: ignore[assignment]
    _patched_responses_sync = True
    logger.info("OpenAI Responses API SDK patched for auto-instrumentation")


def patch_openai_responses_async() -> None:
    """
    Patch OpenAI Responses API async client to auto-track usage.

    This patches the asynchronous `AsyncResponses.create` method.
    """
    global _patched_responses_async, _original_responses_create_async

    if _patched_responses_async:
        logger.debug("OpenAI Responses API async already patched, skipping")
        return

    try:
        from openai.resources.responses import AsyncResponses
    except ImportError:
        logger.warning(
            "OpenAI SDK not installed. Install with: pip install 'fenra[openai]'"
        )
        return

    original_create = AsyncResponses.create
    _original_responses_create_async = original_create

    @functools.wraps(original_create)
    async def patched_create(
        self: AsyncResponses, *args: Any, **kwargs: Any
    ) -> Any:
        if not _is_openai_endpoint(self._client):
            return await original_create(self, *args, **kwargs)

        is_streaming = kwargs.get("stream", False)

        response = await original_create(self, *args, **kwargs)

        context = get_context()

        model = kwargs.get("model") or getattr(response, "model", "unknown")

        tools = kwargs.get("tools")
        try:
            if is_streaming:
                return _track_openai_stream_async(
                    response, model, context, is_responses_api=True, tools=tools
                )
            else:
                _track_responses_api_response(response, model, context, tools=tools)
        except Exception as e:
            logger.error(f"Error tracking OpenAI Responses API usage: {e}", exc_info=True)

        return response

    AsyncResponses.create = patched_create  # type: ignore[assignment]
    _patched_responses_async = True
    logger.info("OpenAI Responses API async SDK patched for auto-instrumentation")


# ============================================================================
# Images API Patching
# ============================================================================


def patch_openai_images() -> None:
    """
    Patch OpenAI Images client to auto-track usage.

    This patches the synchronous `Images.generate` method.
    """
    global _patched_images_sync, _original_images_generate

    if _patched_images_sync:
        logger.debug("OpenAI Images sync already patched, skipping")
        return

    try:
        from openai.resources.images import Images
    except ImportError:
        logger.warning(
            "OpenAI SDK not installed. Install with: pip install 'fenra[openai]'"
        )
        return

    original_generate = Images.generate
    _original_images_generate = original_generate

    @functools.wraps(original_generate)
    def patched_generate(self: Any, *args: Any, **kwargs: Any) -> Any:
        if not _is_openai_endpoint(self._client):
            return original_generate(self, *args, **kwargs)

        response = original_generate(self, *args, **kwargs)

        context = get_context()
        model = kwargs.get("model", "unknown")
        n = kwargs.get("n", 1)
        size = kwargs.get("size", "")
        quality = kwargs.get("quality", "")

        try:
            _track_images_response(
                model=model,
                n=n,
                size=size,
                quality=quality,
                context=context,
                response=response,
            )
        except Exception as e:
            logger.error(f"Error tracking OpenAI Images usage: {e}", exc_info=True)

        return response

    Images.generate = patched_generate  # type: ignore[assignment]
    _patched_images_sync = True
    logger.info("OpenAI Images SDK patched for auto-instrumentation")


def patch_openai_images_async() -> None:
    """
    Patch OpenAI Images async client to auto-track usage.

    This patches the asynchronous `AsyncImages.generate` method.
    """
    global _patched_images_async, _original_images_generate_async

    if _patched_images_async:
        logger.debug("OpenAI Images async already patched, skipping")
        return

    try:
        from openai.resources.images import AsyncImages
    except ImportError:
        logger.warning(
            "OpenAI SDK not installed. Install with: pip install 'fenra[openai]'"
        )
        return

    original_generate = AsyncImages.generate
    _original_images_generate_async = original_generate

    @functools.wraps(original_generate)
    async def patched_generate(self: Any, *args: Any, **kwargs: Any) -> Any:
        if not _is_openai_endpoint(self._client):
            return await original_generate(self, *args, **kwargs)

        response = await original_generate(self, *args, **kwargs)

        context = get_context()
        model = kwargs.get("model", "unknown")
        n = kwargs.get("n", 1)
        size = kwargs.get("size", "")
        quality = kwargs.get("quality", "")

        try:
            _track_images_response(
                model=model,
                n=n,
                size=size,
                quality=quality,
                context=context,
                response=response,
            )
        except Exception as e:
            logger.error(f"Error tracking OpenAI Images usage: {e}", exc_info=True)

        return response

    AsyncImages.generate = patched_generate  # type: ignore[assignment]
    _patched_images_async = True
    logger.info("OpenAI Images async SDK patched for auto-instrumentation")


# ============================================================================
# Beta Chat Completions Parse Patching
# ============================================================================


def patch_openai_beta_parse() -> None:
    """
    Patch OpenAI Beta Chat Completions parse to auto-track usage.

    This patches the synchronous `beta.chat.completions.parse` method.
    Response format is identical to Chat Completions, so we reuse
    _track_chat_completion_response.
    """
    global _patched_beta_parse_sync, _original_beta_parse

    if _patched_beta_parse_sync:
        logger.debug("OpenAI Beta parse sync already patched, skipping")
        return

    try:
        from openai.resources.beta.chat.completions import Completions as BetaCompletions
    except ImportError:
        logger.debug(
            "OpenAI Beta chat completions not available (optional). "
            "Skip patching beta parse."
        )
        return

    if not hasattr(BetaCompletions, "parse"):
        logger.debug("OpenAI Beta Completions has no parse method, skipping")
        return

    original_parse = BetaCompletions.parse
    _original_beta_parse = original_parse

    @functools.wraps(original_parse)
    def patched_parse(self: Any, *args: Any, **kwargs: Any) -> Any:
        if not _is_openai_endpoint(self._client):
            return original_parse(self, *args, **kwargs)

        response = original_parse(self, *args, **kwargs)

        context = get_context()
        model = kwargs.get("model") or getattr(response, "model", "unknown")

        try:
            _track_chat_completion_response(response, model, context)
        except Exception as e:
            logger.error(
                f"Error tracking OpenAI Beta parse usage: {e}", exc_info=True
            )

        return response

    BetaCompletions.parse = patched_parse  # type: ignore[assignment]
    _patched_beta_parse_sync = True
    logger.info("OpenAI Beta chat.completions.parse SDK patched for auto-instrumentation")


def patch_openai_beta_parse_async() -> None:
    """
    Patch OpenAI Beta Chat Completions async parse to auto-track usage.

    This patches the asynchronous `beta.chat.completions.parse` method.
    """
    global _patched_beta_parse_async, _original_beta_parse_async

    if _patched_beta_parse_async:
        logger.debug("OpenAI Beta parse async already patched, skipping")
        return

    try:
        from openai.resources.beta.chat.completions import (
            AsyncCompletions as AsyncBetaCompletions,
        )
    except ImportError:
        logger.debug(
            "OpenAI Beta async chat completions not available (optional). "
            "Skip patching beta parse async."
        )
        return

    if not hasattr(AsyncBetaCompletions, "parse"):
        logger.debug("OpenAI AsyncBeta Completions has no parse method, skipping")
        return

    original_parse = AsyncBetaCompletions.parse
    _original_beta_parse_async = original_parse

    @functools.wraps(original_parse)
    async def patched_parse(self: Any, *args: Any, **kwargs: Any) -> Any:
        if not _is_openai_endpoint(self._client):
            return await original_parse(self, *args, **kwargs)

        response = await original_parse(self, *args, **kwargs)

        context = get_context()
        model = kwargs.get("model") or getattr(response, "model", "unknown")

        try:
            _track_chat_completion_response(response, model, context)
        except Exception as e:
            logger.error(
                f"Error tracking OpenAI Beta parse usage: {e}", exc_info=True
            )

        return response

    AsyncBetaCompletions.parse = patched_parse  # type: ignore[assignment]
    _patched_beta_parse_async = True
    logger.info(
        "OpenAI Beta chat.completions.parse async SDK patched for auto-instrumentation"
    )


# ============================================================================
# Unpatch Functions
# ============================================================================


def unpatch_openai() -> None:
    """Restore original OpenAI Chat Completions create method."""
    global _patched_chat_sync, _original_chat_create

    if not _patched_chat_sync or _original_chat_create is None:
        return

    try:
        from openai.resources.chat.completions import Completions

        Completions.create = _original_chat_create  # type: ignore[assignment]
        _patched_chat_sync = False
        _original_chat_create = None
        logger.info("OpenAI Chat Completions SDK unpatched")
    except ImportError:
        pass


def unpatch_openai_async() -> None:
    """Restore original OpenAI Chat Completions async create method."""
    global _patched_chat_async, _original_chat_create_async

    if not _patched_chat_async or _original_chat_create_async is None:
        return

    try:
        from openai.resources.chat.completions import AsyncCompletions

        AsyncCompletions.create = _original_chat_create_async  # type: ignore[assignment]
        _patched_chat_async = False
        _original_chat_create_async = None
        logger.info("OpenAI Chat Completions async SDK unpatched")
    except ImportError:
        pass


def unpatch_openai_responses() -> None:
    """Restore original OpenAI Responses API create method."""
    global _patched_responses_sync, _original_responses_create

    if not _patched_responses_sync or _original_responses_create is None:
        return

    try:
        from openai.resources.responses import Responses

        Responses.create = _original_responses_create  # type: ignore[assignment]
        _patched_responses_sync = False
        _original_responses_create = None
        logger.info("OpenAI Responses API SDK unpatched")
    except ImportError:
        pass


def unpatch_openai_responses_async() -> None:
    """Restore original OpenAI Responses API async create method."""
    global _patched_responses_async, _original_responses_create_async

    if not _patched_responses_async or _original_responses_create_async is None:
        return

    try:
        from openai.resources.responses import AsyncResponses

        AsyncResponses.create = _original_responses_create_async  # type: ignore[assignment]
        _patched_responses_async = False
        _original_responses_create_async = None
        logger.info("OpenAI Responses API async SDK unpatched")
    except ImportError:
        pass


def unpatch_openai_images() -> None:
    """Restore original OpenAI Images generate method."""
    global _patched_images_sync, _original_images_generate

    if not _patched_images_sync or _original_images_generate is None:
        return

    try:
        from openai.resources.images import Images

        Images.generate = _original_images_generate  # type: ignore[assignment]
        _patched_images_sync = False
        _original_images_generate = None
        logger.info("OpenAI Images SDK unpatched")
    except ImportError:
        pass


def unpatch_openai_images_async() -> None:
    """Restore original OpenAI Images async generate method."""
    global _patched_images_async, _original_images_generate_async

    if not _patched_images_async or _original_images_generate_async is None:
        return

    try:
        from openai.resources.images import AsyncImages

        AsyncImages.generate = _original_images_generate_async  # type: ignore[assignment]
        _patched_images_async = False
        _original_images_generate_async = None
        logger.info("OpenAI Images async SDK unpatched")
    except ImportError:
        pass


def unpatch_openai_beta_parse() -> None:
    """Restore original OpenAI Beta Chat Completions parse method."""
    global _patched_beta_parse_sync, _original_beta_parse

    if not _patched_beta_parse_sync or _original_beta_parse is None:
        return

    try:
        from openai.resources.beta.chat.completions import Completions as BetaCompletions

        BetaCompletions.parse = _original_beta_parse  # type: ignore[assignment]
        _patched_beta_parse_sync = False
        _original_beta_parse = None
        logger.info("OpenAI Beta chat.completions.parse SDK unpatched")
    except ImportError:
        pass


def unpatch_openai_beta_parse_async() -> None:
    """Restore original OpenAI Beta Chat Completions async parse method."""
    global _patched_beta_parse_async, _original_beta_parse_async

    if not _patched_beta_parse_async or _original_beta_parse_async is None:
        return

    try:
        from openai.resources.beta.chat.completions import (
            AsyncCompletions as AsyncBetaCompletions,
        )

        AsyncBetaCompletions.parse = _original_beta_parse_async  # type: ignore[assignment]
        _patched_beta_parse_async = False
        _original_beta_parse_async = None
        logger.info("OpenAI Beta chat.completions.parse async SDK unpatched")
    except ImportError:
        pass


def unpatch_openai_all() -> None:
    """Restore all original OpenAI SDK methods."""
    unpatch_openai()
    unpatch_openai_async()
    unpatch_openai_responses()
    unpatch_openai_responses_async()
    unpatch_openai_images()
    unpatch_openai_images_async()
    unpatch_openai_beta_parse()
    unpatch_openai_beta_parse_async()
