from __future__ import annotations

from typing import Any


def wrap_openai(client: Any) -> Any:
    """Wrap an OpenAI client to automatically trace chat.completions.create() calls."""
    try:
        original_create = client.chat.completions.create
    except AttributeError:
        print("[ZappyBee] Could not wrap client - no chat.completions.create method found.")
        return client

    async def traced_create(*args: Any, **kwargs: Any) -> Any:
        from ..client import TokenCat

        model = kwargs.get("model", "unknown")
        messages = kwargs.get("messages", [])
        stream = kwargs.get("stream", False)
        trace_input = {"model": model, "message_count": len(messages)}

        trace = await TokenCat.start_trace(
            f"openai/{model}",
            input=trace_input,
        )

        step = await trace.start_step(
            "chat.completions.create",
            "llm_call",
            input=trace_input,
            model=model,
        )

        try:
            response = await original_create(*args, **kwargs)

            if stream:
                return _wrap_stream(response, step, trace, model)

            prompt_tokens = getattr(getattr(response, "usage", None), "prompt_tokens", None)
            completion_tokens = getattr(getattr(response, "usage", None), "completion_tokens", None)

            output = None
            if hasattr(response, "choices") and response.choices:
                choice = response.choices[0]
                msg = getattr(choice, "message", None)
                if msg:
                    output = {"role": getattr(msg, "role", "assistant"), "content": getattr(msg, "content", "")}

            await step.end(
                output=output,
                status="success",
                model=model,
                prompt_tokens=prompt_tokens,
                completion_tokens=completion_tokens,
            )
            await trace.end(output=output, status="success")
            return response

        except Exception as e:
            await step.end(status="error", error_message=str(e), model=model)
            await trace.end(status="error", error_message=str(e))
            raise

    client.chat.completions.create = traced_create
    return client


async def _wrap_stream(stream: Any, step: Any, trace: Any, model: str) -> Any:
    """Wrap an async stream to capture usage on completion."""
    prompt_tokens = 0
    completion_tokens = 0
    collected_content = ""

    try:
        async for chunk in stream:
            if hasattr(chunk, "usage") and chunk.usage:
                prompt_tokens = getattr(chunk.usage, "prompt_tokens", 0) or prompt_tokens
                completion_tokens = getattr(chunk.usage, "completion_tokens", 0) or completion_tokens

            if hasattr(chunk, "choices") and chunk.choices:
                delta = getattr(chunk.choices[0], "delta", None)
                if delta and getattr(delta, "content", None):
                    collected_content += delta.content

            yield chunk

        output = {"role": "assistant", "content": collected_content} if collected_content else None
        await step.end(
            output=output,
            status="success",
            model=model,
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
        )
        await trace.end(output=output, status="success")
    except Exception as e:
        await step.end(status="error", error_message=str(e), model=model)
        await trace.end(status="error", error_message=str(e))
        raise
