"""OpenAI Codex OAuth streaming provider."""

import base64
import json
from collections.abc import AsyncIterator
from typing import Any

import aiohttp

from otto.log import get_logger

from ..base import OAuthProvider, OAuthProviderError, ProviderChunk
from ..registry import register_provider

DEFAULT_CODEX_ENDPOINT = "https://chatgpt.com/backend-api/codex/responses"
JWT_CLAIM_PATH = "https://api.openai.com/auth"
_REASONING_EFFORTS = {"none", "low", "medium", "high", "xhigh"}
_CONNECT_TIMEOUT_SECONDS = 20
_READ_TIMEOUT_SECONDS = 45
_MAX_RETRIES = 3
_MAX_RETRY_DELAY_SECONDS = 30.0

log = get_logger(__name__)


def _decode_jwt_payload(token: str) -> dict[str, Any] | None:
    try:
        parts = token.split(".")
        if len(parts) != 3:
            return None
        payload = parts[1]
        payload += "=" * (-len(payload) % 4)
        return json.loads(base64.urlsafe_b64decode(payload))
    except Exception:
        return None


def _extract_account_id(token: str) -> str | None:
    payload = _decode_jwt_payload(token)
    if not payload:
        return None
    auth_claim = payload.get(JWT_CLAIM_PATH)
    if not isinstance(auth_claim, dict):
        return None
    account_id = auth_claim.get("chatgpt_account_id")
    return account_id if isinstance(account_id, str) and account_id else None


def _sanitize_text(text: str) -> str:
    return text.encode("utf-8", errors="surrogatepass").decode("utf-8", errors="replace")


@register_provider
class OpenAICodexProvider(OAuthProvider):
    provider_name = "openai_codex"
    model_prefixes = ["codex"]

    def __init__(self, credentials: dict[str, Any]):
        super().__init__(credentials)

        cred_type = credentials.get("type")
        if cred_type != "oauth":
            raise OAuthProviderError(
                "OpenAI Codex provider requires OAuth credentials. Run: otto auth login openai"
            )

        access = credentials.get("access")
        if not isinstance(access, str) or not access:
            raise OAuthProviderError("Missing access token in OpenAI OAuth credentials")

        account_id = credentials.get("accountId")
        if not isinstance(account_id, str) or not account_id:
            account_id = _extract_account_id(access)
        if not account_id:
            raise OAuthProviderError("Failed to extract accountId from OpenAI access token")

        self._access_token = access
        self._account_id = account_id

    async def stream_completion(
        self,
        model: str,
        messages: list[dict[str, Any]],
        *,
        tools: list[dict[str, Any]] | None = None,
        tool_choice: str | dict[str, Any] | None = None,
        temperature: float | None = None,
        max_tokens: int | None = None,
        system_prompt: str | None = None,
        session_id: str | None = None,
        reasoning: dict[str, Any] | None = None,
        reasoning_effort: str | None = None,
    ) -> AsyncIterator[ProviderChunk]:
        model_id = self._extract_model_id(model)
        request_body = self._build_request(
            model_id=model_id,
            messages=messages,
            tools=tools,
            tool_choice=tool_choice,
            temperature=temperature,
            max_tokens=max_tokens,
            system_prompt=system_prompt,
            session_id=session_id,
            reasoning=reasoning,
            reasoning_effort=reasoning_effort,
        )
        headers = self._build_headers(session_id)

        try:
            timeout = aiohttp.ClientTimeout(
                total=None,
                connect=_CONNECT_TIMEOUT_SECONDS,
                sock_connect=_CONNECT_TIMEOUT_SECONDS,
                sock_read=_READ_TIMEOUT_SECONDS,
            )
            async with aiohttp.ClientSession(timeout=timeout) as session:
                response = await self._request_with_retry(
                    session,
                    "POST",
                    DEFAULT_CODEX_ENDPOINT,
                    headers=headers,
                    json_body=request_body,
                    max_retries=_MAX_RETRIES,
                    max_retry_delay=_MAX_RETRY_DELAY_SECONDS,
                )
                async with response:
                    if response.status != 200:
                        body = await response.text()
                        raise OAuthProviderError(
                            f"OpenAI Codex request failed ({response.status}): {body}"
                        )

                    async for event in self._stream_response(response):
                        yield event
        except TimeoutError as exc:
            raise RuntimeError("OpenAI Codex stream timed out") from exc
        except Exception as exc:
            raise RuntimeError(str(exc))

    @staticmethod
    def _extract_model_id(model: str) -> str:
        model_lower = model.lower()
        if model_lower.startswith("codex/"):
            return model.split("/", 1)[1]
        return model

    def _build_headers(self, session_id: str | None) -> dict[str, str]:
        headers = {
            "Authorization": f"Bearer {self._access_token}",
            "chatgpt-account-id": self._account_id,
            "OpenAI-Beta": "responses=experimental",
            "Accept": "text/event-stream",
            "Content-Type": "application/json",
            "originator": "otto",
        }
        if session_id:
            headers["session_id"] = session_id
        return headers

    def _build_request(
        self,
        *,
        model_id: str,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]] | None,
        tool_choice: str | dict[str, Any] | None,
        temperature: float | None,
        max_tokens: int | None,
        system_prompt: str | None,
        session_id: str | None,
        reasoning: dict[str, Any] | None,
        reasoning_effort: str | None,
    ) -> dict[str, Any]:
        effective_reasoning = dict(reasoning or {})
        if reasoning_effort is not None:
            effective_reasoning["effort"] = reasoning_effort

        request: dict[str, Any] = {
            "model": model_id,
            "store": False,
            "stream": True,
            "input": self._convert_messages(messages),
            "tool_choice": tool_choice if tool_choice is not None else "auto",
            "parallel_tool_calls": True,
            "reasoning": {
                "effort": self._normalize_reasoning_effort(
                    str(effective_reasoning.get("effort", "medium"))
                ),
                "summary": str(effective_reasoning.get("summary", "auto")),
            },
        }

        if system_prompt:
            request["instructions"] = _sanitize_text(system_prompt)
        if temperature is not None:
            request["temperature"] = temperature
        if max_tokens is not None:
            request["max_output_tokens"] = max_tokens
        if session_id:
            request["prompt_cache_key"] = session_id
        if tools:
            request["tools"] = self._convert_tools(tools)

        return request

    @staticmethod
    def _normalize_reasoning_effort(effort: str) -> str:
        normalized = effort.lower().strip()
        if normalized in _REASONING_EFFORTS:
            return normalized
        if normalized in {"off", "disabled"}:
            return "none"
        if normalized in {"minimal", "on", "enabled"}:
            return "low" if normalized == "minimal" else "medium"
        return "medium"

    def _convert_messages(self, messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
        codex_messages: list[dict[str, Any]] = []
        for message in messages:
            role = str(message.get("role", "user"))
            content = message.get("content", "")

            if role == "system":
                continue

            if role == "tool":
                codex_messages.append(
                    {
                        "type": "function_call_output",
                        "call_id": str(message.get("tool_call_id", "")),
                        "output": _sanitize_text(str(content)),
                    }
                )
                continue

            if role == "assistant":
                tool_calls = message.get("tool_calls")
                if isinstance(tool_calls, list):
                    for tool_call in tool_calls:
                        if not isinstance(tool_call, dict):
                            continue
                        function = tool_call.get("function")
                        function_data = function if isinstance(function, dict) else {}
                        codex_messages.append(
                            {
                                "type": "function_call",
                                "call_id": str(tool_call.get("id", "")),
                                "name": str(function_data.get("name", "")),
                                "arguments": str(function_data.get("arguments", "{}")),
                            }
                        )
                if isinstance(content, str) and content.strip():
                    codex_messages.append(
                        {
                            "type": "message",
                            "role": "assistant",
                            "content": _sanitize_text(content),
                        }
                    )
                continue

            if isinstance(content, str) and content.strip():
                codex_messages.append(
                    {
                        "type": "message",
                        "role": "user",
                        "content": _sanitize_text(content),
                    }
                )
            elif isinstance(content, list):
                parts: list[dict[str, Any]] = []
                for item in content:
                    if not isinstance(item, dict):
                        continue
                    if item.get("type") == "text":
                        text = str(item.get("text", ""))
                        if text:
                            parts.append({"type": "input_text", "text": _sanitize_text(text)})
                    elif item.get("type") == "image_url":
                        image_obj = item.get("image_url")
                        if isinstance(image_obj, dict):
                            url = str(image_obj.get("url", ""))
                            if url:
                                parts.append({"type": "input_image", "image_url": url})
                if parts:
                    codex_messages.append({"type": "message", "role": "user", "content": parts})

        return codex_messages

    @staticmethod
    def _convert_tools(tools: list[dict[str, Any]]) -> list[dict[str, Any]]:
        converted: list[dict[str, Any]] = []
        for tool in tools:
            if tool.get("type") != "function":
                continue
            function = tool.get("function")
            function_data = function if isinstance(function, dict) else {}
            converted.append(
                {
                    "type": "function",
                    "name": str(function_data.get("name", "")),
                    "description": str(function_data.get("description", "")),
                    "parameters": function_data.get(
                        "parameters", {"type": "object", "properties": {}}
                    ),
                }
            )
        return converted

    async def _stream_response(
        self, response: aiohttp.ClientResponse
    ) -> AsyncIterator[ProviderChunk]:
        buffer = ""
        tool_calls: dict[str, dict[str, str]] = {}
        completed_calls: set[str] = set()
        tool_call_indices: dict[str, int] = {}
        tool_call_ids_by_output_index: dict[int, str] = {}
        next_tool_call_index = 0
        usage: dict[str, int] | None = None
        has_content = False
        finish_reason = "stop"

        def resolve_tool_call_index(call_id: str) -> int:
            nonlocal next_tool_call_index
            existing_index = tool_call_indices.get(call_id)
            if existing_index is not None:
                return existing_index
            next_index = next_tool_call_index
            tool_call_indices[call_id] = next_index
            next_tool_call_index += 1
            return next_index

        def _event_output_index(event: dict[str, Any]) -> int | None:
            raw_output_index = event.get("output_index")
            if isinstance(raw_output_index, int):
                return raw_output_index
            try:
                return int(raw_output_index)
            except (TypeError, ValueError):
                return None

        def _resolve_call_id(event: dict[str, Any], item_data: dict[str, Any] | None = None) -> str:
            event_call_id = event.get("call_id")
            if isinstance(event_call_id, str) and event_call_id:
                return event_call_id

            if item_data is not None:
                item_call_id = item_data.get("call_id")
                if isinstance(item_call_id, str) and item_call_id:
                    return item_call_id

            output_index = _event_output_index(event)
            if output_index is not None:
                mapped_call_id = tool_call_ids_by_output_index.get(output_index)
                if mapped_call_id:
                    return mapped_call_id

            return ""

        def _remember_call_id(event: dict[str, Any], call_id: str) -> None:
            if not call_id:
                return
            output_index = _event_output_index(event)
            if output_index is None:
                return
            tool_call_ids_by_output_index[output_index] = call_id

        async for chunk in response.content:
            buffer += (
                chunk.decode("utf-8", errors="replace").replace("\r\n", "\n").replace("\r", "\n")
            )

            while "\n\n" in buffer:
                raw_event, buffer = buffer.split("\n\n", 1)
                data_lines = [
                    line[5:].strip() for line in raw_event.splitlines() if line.startswith("data:")
                ]
                if not data_lines:
                    continue

                data_str = "\n".join(data_lines).strip()
                if not data_str or data_str == "[DONE]":
                    continue

                try:
                    event = json.loads(data_str)
                except json.JSONDecodeError:
                    continue

                event_type = str(event.get("type", ""))

                if event_type == "error":
                    message = str(event.get("message", "") or event.get("code", "Unknown error"))
                    raise OAuthProviderError(message)

                if event_type == "response.failed":
                    details = event.get("response")
                    details_data = details if isinstance(details, dict) else {}
                    error_obj = details_data.get("error")
                    error_data = error_obj if isinstance(error_obj, dict) else {}
                    raise OAuthProviderError(
                        str(error_data.get("message", "Codex response failed"))
                    )

                if event_type == "response.output_text.delta":
                    delta = event.get("delta")
                    if isinstance(delta, str) and delta:
                        has_content = True
                        yield self._text_chunk(delta)
                    continue

                if event_type == "response.reasoning_summary_text.delta":
                    delta = event.get("delta")
                    if isinstance(delta, str) and delta:
                        has_content = True
                    continue

                if event_type == "response.output_item.added":
                    item = event.get("item")
                    item_data = item if isinstance(item, dict) else {}
                    if item_data.get("type") == "function_call":
                        call_id = _resolve_call_id(event, item_data)
                        if not call_id:
                            continue
                        _remember_call_id(event, call_id)
                        name = str(item_data.get("name", ""))
                        arguments = str(item_data.get("arguments", ""))
                        has_content = True
                        existing = tool_calls.get(call_id, {"name": "", "arguments": ""})
                        if name:
                            existing["name"] = name
                        if arguments:
                            existing["arguments"] += arguments
                        tool_calls[call_id] = existing
                        tool_call_index = resolve_tool_call_index(call_id)
                        yield self._tool_call_chunk(
                            tool_call_index, call_id, existing["name"] or None, None
                        )
                    continue

                if event_type == "response.function_call_arguments.delta":
                    call_id = _resolve_call_id(event)
                    delta = event.get("delta")
                    delta_text = delta if isinstance(delta, str) else ""
                    if call_id and delta_text:
                        _remember_call_id(event, call_id)
                        has_content = True
                        tool_data = tool_calls.setdefault(call_id, {"name": "", "arguments": ""})
                        tool_data["arguments"] += delta_text
                        tool_call_index = resolve_tool_call_index(call_id)
                        yield self._tool_call_chunk(tool_call_index, call_id, None, delta_text)
                    continue

                if event_type == "response.output_item.done":
                    item = event.get("item")
                    item_data = item if isinstance(item, dict) else {}
                    if item_data.get("type") == "function_call":
                        call_id = _resolve_call_id(event, item_data)
                        if not call_id:
                            continue
                        _remember_call_id(event, call_id)
                        name = str(item_data.get("name", ""))
                        arguments = str(item_data.get("arguments", ""))
                        has_content = True
                        tool_data = tool_calls.setdefault(call_id, {"name": "", "arguments": ""})
                        previous_name = tool_data.get("name", "")
                        previous_arguments = tool_data.get("arguments", "")
                        if name:
                            tool_data["name"] = name
                        emitted_arguments: str | None = None
                        if arguments:
                            if arguments.startswith(previous_arguments):
                                emitted_arguments = arguments[len(previous_arguments) :]
                            elif arguments != previous_arguments:
                                emitted_arguments = arguments
                            tool_data["arguments"] = arguments
                        tool_call_index = resolve_tool_call_index(call_id)
                        if (name and not previous_name) or emitted_arguments:
                            yield self._tool_call_chunk(
                                tool_call_index,
                                call_id,
                                name if name and not previous_name else None,
                                emitted_arguments if emitted_arguments else None,
                            )
                        completed_calls.add(call_id)
                    continue

                if event_type == "response.function_call_arguments.done":
                    call_id = _resolve_call_id(event)
                    if not call_id or call_id in completed_calls or call_id not in tool_calls:
                        continue
                    _remember_call_id(event, call_id)
                    completed_calls.add(call_id)
                    continue

                if event_type in {"response.done", "response.completed"}:
                    response_data_raw = event.get("response")
                    response_data = response_data_raw if isinstance(response_data_raw, dict) else {}
                    usage_raw = response_data.get("usage")
                    usage_data = usage_raw if isinstance(usage_raw, dict) else {}
                    if usage_data:
                        usage = {
                            "prompt_tokens": int(usage_data.get("input_tokens", 0)),
                            "completion_tokens": int(usage_data.get("output_tokens", 0)),
                            "total_tokens": int(usage_data.get("total_tokens", 0)),
                        }

        if not has_content:
            raise OAuthProviderError("Codex API returned an empty response")

        yield self._final_chunk(finish_reason, usage)
