from __future__ import annotations

# standard
import json
import time
from collections.abc import AsyncIterator

# third party
import httpx

# custom
from ..base import BaseEngine
from ..types import (
    Message,
    Performance,
    Response,
    StopReason,
    Tool,
    ToolCall,
    Usage,
    make_non_stream_performance,
    resolve_tokens,
)


class OpenAIEngine(BaseEngine):
    def __init__(
        self,
        model: str,
        api_key: str,
        provider: str = "openai",
        base_url: str = "https://api.openai.com/v1",
        max_tokens: int = 8192,
        reasoning_effort: str | None = None,
    ) -> None:
        self.provider = provider
        self.model = model
        self.api_key = api_key
        self.base_url = base_url.rstrip("/")
        self.max_tokens = max_tokens
        # NOTE "low" | "medium" | "high" — only sent for reasoning-capable models
        self.reasoning_effort = reasoning_effort

    ### PAYLOAD BUILDERS

    def _build_messages(self, messages: list[Message]) -> list[dict]:
        # NOTE reasoning is only echoed on the last assistant turn so that cross-engine
        # conversation history doesn't send provider-specific data to the wrong API.
        last_assistant_idx = max(
            (i for i, m in enumerate(messages) if m.role.value == "assistant"),
            default=-1,
        )

        result = []
        for idx, m in enumerate(messages):
            if m.role.value == "tool":
                result.append(
                    {
                        "role": "tool",
                        "tool_call_id": m.tool_call_id,
                        "content": m.content,
                    }
                )
            elif m.tool_calls:
                msg: dict = {
                    "role": "assistant",
                    "content": m.content,
                    "tool_calls": [
                        {
                            "id": tc.id,
                            "type": "function",
                            "function": {
                                "name": tc.name,
                                "arguments": json.dumps(tc.arguments),
                            },
                        }
                        for tc in m.tool_calls
                    ],
                }
                if idx == last_assistant_idx and m.reasoning_content is not None:
                    msg["reasoning_content"] = m.reasoning_content
                result.append(msg)
            else:
                entry: dict = {"role": m.role.value, "content": m.content}
                if (
                    m.role.value == "assistant"
                    and idx == last_assistant_idx
                    and m.reasoning_content is not None
                ):
                    entry["reasoning_content"] = m.reasoning_content
                result.append(entry)
        return result

    def _build_tools(self, tools: list[Tool]) -> list[dict]:
        return [
            {
                "type": "function",
                "function": {
                    "name": t.name,
                    "description": t.description,
                    "parameters": t.parameters,
                },
            }
            for t in tools
        ]

    def _build_payload(self, messages: list[Message], tools: list[Tool] | None) -> dict:
        payload: dict = {
            "model": self.model,
            "messages": self._build_messages(messages),
            "max_completion_tokens": self.max_tokens,
        }
        if self.reasoning_effort is not None:
            payload["reasoning_effort"] = self.reasoning_effort
        if tools:
            payload["tools"] = self._build_tools(tools)
        return payload

    ### RESPONSE PARSERS

    def _parse_response(self, raw: dict) -> Response:
        choice = raw["choices"][0]
        msg = choice["message"]
        content: str | None = msg.get("content")
        # NOTE deepseek reasoner and other reasoning models expose chain-of-thought here.
        # NOTE Not sent back in subsequent requests — informational only.
        reasoning_content: str | None = msg.get("reasoning_content")
        tool_calls: list[ToolCall] | None = None

        if msg.get("tool_calls"):
            tool_calls = [
                ToolCall(
                    id=tc["id"],
                    name=tc["function"]["name"],
                    arguments=json.loads(tc["function"]["arguments"]),
                )
                for tc in msg["tool_calls"]
            ]

        usage = raw.get("usage", {})
        in_tok = usage.get("prompt_tokens", 0)
        out_tok = usage.get("completion_tokens", 0)
        tot_tok = usage.get("total_tokens", 0)
        cache_read_tokens = usage.get("prompt_tokens_details", {}).get("cached_tokens", 0)
        input_tokens, output_tokens, total_tokens = resolve_tokens(
            in_tok, out_tok, tot_tok
        )
        finish_reason = choice.get("finish_reason", "stop")
        stop_map = {
            "stop": StopReason.END_TURN,
            "tool_calls": StopReason.TOOL_USE,
            "length": StopReason.MAX_TOKENS,
        }

        return Response(
            provider=self.provider,
            model=self.model,
            content=content,
            tool_calls=tool_calls,
            stop_reason=stop_map.get(finish_reason, StopReason.END_TURN),
            reasoning_content=reasoning_content,
            usage=Usage(
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                total_tokens=total_tokens,
                cache_read_tokens=cache_read_tokens,
            ),
        )

    ### INTERFACES

    def _headers(self) -> dict:
        return {"Authorization": f"Bearer {self.api_key}"}

    async def chat(
        self,
        messages: list[Message],
        tools: list[Tool] | None = None,
    ) -> Response:
        t_start = time.monotonic()
        async with httpx.AsyncClient() as client:
            resp = await client.post(
                f"{self.base_url}/chat/completions",
                headers=self._headers(),
                json=self._build_payload(messages, tools),
                timeout=120,
            )
            if not resp.is_success:
                raise RuntimeError(
                    f"OpenAI-compatible API error {resp.status_code}: {resp.text}"
                )
            response = self._parse_response(resp.json())
        total_duration = time.monotonic() - t_start
        response.performance = make_non_stream_performance(response, total_duration)
        return response

    async def stream(
        self,
        messages: list[Message],
        tools: list[Tool] | None = None,
    ) -> AsyncIterator[Response]:
        payload = self._build_payload(messages, tools)
        payload["stream"] = True
        payload["stream_options"] = {"include_usage": True}

        stop_reason: str = StopReason.END_TURN.value
        stop_map = {
            "stop": StopReason.END_TURN,
            "tool_calls": StopReason.TOOL_USE,
            "length": StopReason.MAX_TOKENS,
        }
        input_tokens = 0
        output_tokens = 0
        total_tokens = 0
        cache_read_tokens = 0

        # tool call accumulators keyed by index inside choices[0].delta.tool_calls.
        tool_buf: dict[int, dict] = {}

        t_start = time.monotonic()
        t_first_chunk: float | None = None
        t_reasoning_start: float | None = None
        t_reasoning_end: float | None = None
        t_content_start: float | None = None

        async with httpx.AsyncClient() as client:
            async with client.stream(
                "POST",
                f"{self.base_url}/chat/completions",
                headers=self._headers(),
                json=payload,
                timeout=120,
            ) as resp:
                resp.raise_for_status()
                if self.reasoning_effort is not None:
                    # Model will reason silently — tokens are not streamed.
                    # Yield a sentinel immediately so the consumer isn't left in
                    # the dark; record the request time as reasoning start.
                    now = time.monotonic()
                    t_first_chunk = now
                    t_reasoning_start = t_start
                    yield Response(
                        provider=self.provider,
                        model=self.model,
                        streaming=True,
                        synthetic=True,
                        reasoning_content="Reasoning in progress - tokens not available for this model",
                    )
                async for line in resp.aiter_lines():
                    if not line.startswith("data: "):  # pragma: no cover
                        continue
                    data = line[6:]
                    if data == "[DONE]":
                        break
                    now = time.monotonic()
                    if t_first_chunk is None:
                        t_first_chunk = now

                    chunk = json.loads(data)
                    # With include_usage, a final chunk arrives with usage and no choices.
                    if usage := chunk.get("usage"):
                        in_tok = usage.get("prompt_tokens", 0)
                        out_tok = usage.get("completion_tokens", 0)
                        tot_tok = usage.get("total_tokens", 0)
                        cache_read_tokens = usage.get("prompt_tokens_details", {}).get("cached_tokens", 0)
                        input_tokens, output_tokens, total_tokens = resolve_tokens(
                            in_tok, out_tok, tot_tok
                        )
                    if not chunk.get("choices"):
                        continue
                    choice = chunk["choices"][0]
                    delta = choice.get("delta", {})

                    if fr := choice.get("finish_reason"):
                        if tool_buf:
                            tcs = [
                                ToolCall(
                                    id=buf.get("id") or f"tc_{i}",
                                    name=buf["name"],
                                    arguments=(
                                        json.loads(buf["arguments"])
                                        if buf["arguments"]
                                        else {}
                                    ),
                                )
                                for i, buf in sorted(tool_buf.items())
                            ]
                            stop_reason = fr
                            yield Response(
                                provider=self.provider,
                                model=self.model,
                                streaming=True,
                                tool_calls=tcs,
                            )
                            tool_buf.clear()

                    if rc := delta.get("reasoning_content"):
                        if t_reasoning_start is None:
                            t_reasoning_start = now
                        yield Response(
                            provider=self.provider,
                            model=self.model,
                            streaming=True,
                            reasoning_content=rc,
                        )

                    if text := delta.get("content"):
                        if t_content_start is None:
                            t_content_start = now
                            if (
                                t_reasoning_start is not None
                                and t_reasoning_end is None
                            ):
                                t_reasoning_end = now
                        yield Response(
                            provider=self.provider,
                            model=self.model,
                            streaming=True,
                            content=text,
                        )

                    for tc_delta in delta.get("tool_calls", []):
                        i = tc_delta.get("index", 0)
                        if i not in tool_buf:
                            tool_buf[i] = {"id": "", "name": "", "arguments": ""}
                        if tc_id := tc_delta.get("id"):
                            tool_buf[i]["id"] = tc_id
                        fn = tc_delta.get("function", {})
                        if name := fn.get("name"):
                            tool_buf[i]["name"] = name
                        if args := fn.get("arguments"):
                            tool_buf[i]["arguments"] += args

        t_end = time.monotonic()
        total_duration = t_end - t_start
        latency = (t_first_chunk - t_start) if t_first_chunk else total_duration
        reasoning_duration = (
            ((t_reasoning_end or t_end) - t_reasoning_start)
            if t_reasoning_start
            else 0.0
        )
        content_duration = (t_end - t_content_start) if t_content_start else 0.0
        throughput = int(output_tokens / total_duration) if total_duration > 0 else 0

        yield Response(
            provider=self.provider,
            model=self.model,
            streaming=True,
            stop_reason=stop_map.get(stop_reason, StopReason.END_TURN),
            usage=Usage(
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                total_tokens=total_tokens,
                cache_read_tokens=cache_read_tokens,
            ),
            performance=Performance(
                latency=round(latency, 2),
                reasoning_duration=round(reasoning_duration, 2),
                content_duration=round(content_duration, 2),
                total_duration=round(total_duration, 2),
                throughput=throughput,
            ),
        )
