"""
OpenAI integration — wraps openai.OpenAI to auto-trace all chat completions.

Supports both non-streaming and streaming (``stream=True``) calls.
For streams, chunks are yielded immediately while the wrapper silently
buffers text.  Once the iterator is exhausted the full trace is dispatched.

Usage:
    import openai
    from syntropy import Syntropy
    from syntropy.integrations.openai import wrap_openai

    syn = Syntropy(api_key="syn_...", agent_id="my-agent")
    client = wrap_openai(openai.OpenAI(), syntropy=syn)

    # Non-streaming (auto-traced):
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": "Hello!"}],
    )

    # Streaming (auto-traced on completion):
    for chunk in client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": "Hello!"}],
        stream=True,
    ):
        print(chunk.choices[0].delta.content or "", end="")
"""

from __future__ import annotations

import time
from typing import Any, Generator, Iterator

from syntropy.client import Syntropy
from syntropy.trace_context import is_inside_trace


# ----------------------------------------------------------------
# Helpers
# ----------------------------------------------------------------

def _extract_prompt(kwargs: dict[str, Any]) -> str:
    messages = kwargs.get("messages", [])
    if messages and isinstance(messages[-1], dict):
        return messages[-1].get("content", str(messages))
    return str(messages)


def _wrap_stream(
    stream: Any,
    kwargs: dict[str, Any],
    start: float,
    syntropy: Syntropy,
    agent_id: str,
) -> Any:
    """Yield chunks from an OpenAI stream, buffer text, trace on end."""
    text_parts: list[str] = []
    model = kwargs.get("model", "unknown")
    tokens_in = 0
    tokens_out = 0
    finish_reason: str | None = None
    stream_error: BaseException | None = None

    try:
        for chunk in stream:
            # Buffer text
            delta = getattr(chunk.choices[0], "delta", None) if chunk.choices else None
            if delta and getattr(delta, "content", None):
                text_parts.append(delta.content)
            if chunk.choices and getattr(chunk.choices[0], "finish_reason", None):
                finish_reason = chunk.choices[0].finish_reason
            if hasattr(chunk, "model") and chunk.model:
                model = chunk.model
            # OpenAI may include usage in the final chunk
            if hasattr(chunk, "usage") and chunk.usage:
                tokens_in = getattr(chunk.usage, "prompt_tokens", 0)
                tokens_out = getattr(chunk.usage, "completion_tokens", 0)
            yield chunk
    except BaseException as exc:
        stream_error = exc
        raise
    finally:
        elapsed_ms = int((time.perf_counter() - start) * 1000)
        syntropy.trace(
            prompt=_extract_prompt(kwargs),
            response="".join(text_parts),
            model=model,
            agent_id=agent_id,
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            latency_ms=elapsed_ms,
            status="error" if stream_error else "success",
            metadata={
                "provider": "openai",
                "stream": True,
                "temperature": kwargs.get("temperature"),
                "finish_reason": finish_reason,
                **({
                    "error": str(stream_error),
                    "partialResponse": True,
                } if stream_error else {}),
            },
        )


# ----------------------------------------------------------------
# Traced proxy classes
# ----------------------------------------------------------------


class _TracedCompletions:
    """Proxy around openai.ChatCompletions that records traces."""

    def __init__(
        self,
        original: Any,
        syntropy: Syntropy,
        agent_id: str,
    ):
        self._original = original
        self._syntropy = syntropy
        self._agent_id = agent_id

    def create(self, **kwargs: Any) -> Any:
        """Intercept chat.completions.create() and trace the result."""
        # Skip if already inside an orchestration-level trace (e.g. LangChain)
        if is_inside_trace():
            return self._original.create(**kwargs)

        start = time.perf_counter()
        is_stream = kwargs.get("stream", False)

        response = self._original.create(**kwargs)

        # ---- Streaming path ----
        if is_stream:
            return _wrap_stream(
                response, kwargs, start, self._syntropy, self._agent_id,
            )

        # ---- Non-streaming path ----
        elapsed_ms = int((time.perf_counter() - start) * 1000)

        prompt = _extract_prompt(kwargs)

        choice = response.choices[0] if response.choices else None
        output = (
            choice.message.content or ""
            if choice and hasattr(choice, "message")
            else ""
        )

        usage = response.usage
        tokens_in = usage.prompt_tokens if usage else 0
        tokens_out = usage.completion_tokens if usage else 0

        self._syntropy.trace(
            prompt=prompt,
            response=output,
            model=response.model or kwargs.get("model", "unknown"),
            agent_id=self._agent_id,
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            latency_ms=elapsed_ms,
            metadata={
                "provider": "openai",
                "stream": False,
                "temperature": kwargs.get("temperature"),
            },
        )

        return response

    def __getattr__(self, name: str) -> Any:
        return getattr(self._original, name)


class _TracedChat:
    """Proxy around openai.Chat that injects traced completions."""

    def __init__(self, original: Any, syntropy: Syntropy, agent_id: str):
        self._original = original
        self._syntropy = syntropy
        self._agent_id = agent_id
        self.completions = _TracedCompletions(
            original.completions, syntropy, agent_id
        )

    def __getattr__(self, name: str) -> Any:
        return getattr(self._original, name)


class TracedOpenAI:
    """Wrapper around openai.OpenAI that auto-traces chat completions."""

    def __init__(
        self,
        client: Any,
        syntropy: Syntropy,
        agent_id: str = "",
    ):
        self._client = client
        self._syntropy = syntropy
        self._agent_id = agent_id or syntropy._config.agent_id
        self.chat = _TracedChat(client.chat, syntropy, self._agent_id)

    def __getattr__(self, name: str) -> Any:
        return getattr(self._client, name)


def wrap_openai(
    client: Any,
    *,
    syntropy: Syntropy,
    agent_id: str = "",
) -> TracedOpenAI:
    """
    Wrap an OpenAI client to auto-trace all chat.completions.create() calls.
    Supports both streaming and non-streaming.

    Args:
        client: An openai.OpenAI() instance.
        syntropy: A Syntropy client instance.
        agent_id: Override agent_id (defaults to syntropy.config.agent_id).

    Returns:
        A drop-in replacement for the OpenAI client with tracing enabled.
    """
    return TracedOpenAI(client, syntropy, agent_id)
