"""
LiteLLM custom provider — OpenAI Codex via chatgpt.com/backend-api.

Auth tokens are read from CODEX_AUTH_FILE (default /root/.codex/auth.json)
and auto-refreshed using the OpenAI OAuth endpoint before expiry.
"""
import base64
import json
import os
import time
from typing import AsyncIterator, Callable, Iterator, Optional, Union

import httpx
import litellm

from litellm.llms.custom_llm import CustomLLM, CustomLLMError
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.types.utils import GenericStreamingChunk, Usage
from litellm.utils import ModelResponse

AUTH_FILE = os.environ.get("CODEX_AUTH_FILE", "/root/.codex/auth.json")

# Self-register so the Router can validate this provider before custom_llm_setup() runs
if "codex" not in litellm.provider_list:
    litellm.provider_list.append("codex")
if "codex" not in litellm._custom_providers:
    litellm._custom_providers.append("codex")
TOKEN_URL = "https://auth.openai.com/oauth/token"
API_URL = "https://chatgpt.com/backend-api/codex/responses"


# ---------------------------------------------------------------------------
# Auth helpers
# ---------------------------------------------------------------------------

def _load_auth() -> dict:
    with open(AUTH_FILE) as f:
        return json.load(f)


def _save_auth(data: dict) -> None:
    with open(AUTH_FILE, "w") as f:
        json.dump(data, f, indent=2)


def _jwt_exp(token: str) -> int:
    try:
        payload = token.split(".")[1]
        payload += "=" * (4 - len(payload) % 4)
        return json.loads(base64.urlsafe_b64decode(payload)).get("exp", 0)
    except Exception:
        return 0


def _client_id(auth: dict) -> str:
    """Resolve the OAuth client from config or the ID token audience."""
    override = os.environ.get("CODEX_CLIENT_ID")
    if override:
        return override

    try:
        payload = auth["tokens"]["id_token"].split(".")[1]
        payload += "=" * (-len(payload) % 4)
        audience = json.loads(base64.urlsafe_b64decode(payload))["aud"]
        if isinstance(audience, str) and audience:
            return audience
        if (
            isinstance(audience, list)
            and len(audience) == 1
            and isinstance(audience[0], str)
            and audience[0]
        ):
            return audience[0]
    except (KeyError, IndexError, TypeError, ValueError, json.JSONDecodeError):
        pass

    raise ValueError(
        "Codex OAuth client ID is unavailable: set CODEX_CLIENT_ID or "
        "provide an id_token with a single aud value"
    )


def _refresh(auth: dict) -> dict:
    resp = httpx.post(
        TOKEN_URL,
        json={
            "client_id": _client_id(auth),
            "grant_type": "refresh_token",
            "refresh_token": auth["tokens"]["refresh_token"],
        },
        timeout=30,
    )
    resp.raise_for_status()
    data = resp.json()
    auth["tokens"]["access_token"] = data["access_token"]
    if "refresh_token" in data:
        auth["tokens"]["refresh_token"] = data["refresh_token"]
    auth["last_refresh"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
    _save_auth(auth)
    return auth


def _get_auth(force: bool = False) -> dict:
    auth = _load_auth()
    if force or time.time() > _jwt_exp(auth["tokens"]["access_token"]) - 60:
        auth = _refresh(auth)
    return auth


# ---------------------------------------------------------------------------
# Request / response helpers
# ---------------------------------------------------------------------------

def _split_messages(messages: list) -> tuple[str, list]:
    """Convert OpenAI chat messages into Responses API input items."""
    instructions_parts = []
    input_msgs = []
    for msg in messages:
        role = msg.get("role", "user")
        content = msg.get("content", "")
        if role in ("system", "developer"):
            instructions_parts.append(content)
        elif role == "assistant" and msg.get("tool_calls"):
            if content:
                input_msgs.append({"role": role, "content": content})
            for tool_call in msg["tool_calls"]:
                if not isinstance(tool_call, dict):
                    tool_call = tool_call.model_dump()
                function = tool_call.get("function", {})
                input_msgs.append({
                    "type": "function_call",
                    "call_id": tool_call["id"],
                    "name": function["name"],
                    "arguments": function.get("arguments", ""),
                })
        elif role == "tool":
            input_msgs.append({
                "type": "function_call_output",
                "call_id": msg["tool_call_id"],
                "output": content,
            })
        else:
            input_msgs.append({"role": role, "content": content})
    return "\n\n".join(instructions_parts), input_msgs


def _resolve_model(model: str) -> str:
    """Translate internal cx-* names back to the real Codex API model names (gpt-*)."""
    if model.startswith("cx-"):
        return "gpt-" + model[3:]
    return model


def _build_body(model: str, messages: list, optional_params: dict, stream: bool) -> dict:
    instructions, input_msgs = _split_messages(messages)
    body: dict = {
        "model": _resolve_model(model),
        "instructions": instructions or "You are a helpful assistant.",
        "input": input_msgs,
        "stream": True,  # API only supports streaming; non-stream calls aggregate internally
        "store": False,
        "reasoning": {"effort": optional_params.get("reasoning_effort", "medium")},
    }

    tools = []
    for tool in optional_params.get("tools") or []:
        if not isinstance(tool, dict):
            tool = tool.model_dump()
        if tool.get("type") == "function" and "function" in tool:
            function = tool["function"]
            tools.append({
                "type": "function",
                "name": function["name"],
                "description": function.get("description", ""),
                "parameters": function.get("parameters", {"type": "object"}),
                **({"strict": function["strict"]} if "strict" in function else {}),
            })
        else:
            tools.append(tool)
    if tools:
        body["tools"] = tools

    tool_choice = optional_params.get("tool_choice")
    if isinstance(tool_choice, dict) and "function" in tool_choice:
        function = tool_choice["function"]
        body["tool_choice"] = {"type": "function", "name": function["name"]}
    elif tool_choice is not None:
        body["tool_choice"] = tool_choice
    if "parallel_tool_calls" in optional_params:
        body["parallel_tool_calls"] = optional_params["parallel_tool_calls"]
    return body


def _auth_headers(auth: dict) -> dict:
    return {
        "Authorization": f"Bearer {auth['tokens']['access_token']}",
        "ChatGPT-Account-Id": auth["tokens"]["account_id"],
        "Content-Type": "application/json",
        "User-Agent": "opencode/1.0",
        "originator": "opencode",
    }


def _accumulate_tool_chunk(tool_calls: dict, chunk: GenericStreamingChunk) -> None:
    tool_use = chunk.get("tool_use")
    if not tool_use:
        return
    index = tool_use.get("index", 0)
    current = tool_calls.setdefault(index, {
        "id": "",
        "type": "function",
        "function": {"name": "", "arguments": ""},
    })
    if tool_use.get("id"):
        current["id"] = tool_use["id"]
    function = tool_use.get("function") or {}
    if function.get("name"):
        current["function"]["name"] = function["name"]
    current["function"]["arguments"] += function.get("arguments") or ""


def _event_chunk(event: dict, state: dict) -> Optional[GenericStreamingChunk]:
    """Translate a Responses API SSE event into LiteLLM's generic chunk."""
    etype = event.get("type", "")
    if etype == "response.output_text.delta":
        return GenericStreamingChunk(
            text=event.get("delta", ""),
            is_finished=False,
            finish_reason="",
            usage=None,
        )
    if etype == "response.output_item.added":
        item = event.get("item", {})
        if item.get("type") == "function_call":
            index = event.get("output_index", 0)
            state["tool_calls"] = True
            state.setdefault("tools", {})[index] = item
            return GenericStreamingChunk(
                text="",
                tool_use={
                    "index": index,
                    "id": item.get("call_id"),
                    "type": "function",
                    "function": {
                        "name": item.get("name", ""),
                        "arguments": item.get("arguments", ""),
                    },
                },
                is_finished=False,
                finish_reason="",
                usage=None,
            )
    if etype == "response.function_call_arguments.delta":
        index = event.get("output_index", 0)
        return GenericStreamingChunk(
            text="",
            tool_use={
                "index": index,
                "id": None,
                "type": "function",
                "function": {"name": None, "arguments": event.get("delta", "")},
            },
            is_finished=False,
            finish_reason="",
            usage=None,
        )
    if etype in ("response.completed", "response.done"):
        usage_data = event.get("response", {}).get("usage", {})
        return GenericStreamingChunk(
            text="",
            is_finished=True,
            finish_reason="tool_calls" if state.get("tool_calls") else "stop",
            usage={
                "prompt_tokens": usage_data.get("input_tokens", 0),
                "completion_tokens": usage_data.get("output_tokens", 0),
                "total_tokens": usage_data.get("input_tokens", 0)
                + usage_data.get("output_tokens", 0),
            },
        )
    return None


def _iter_sse(response: httpx.Response) -> Iterator[GenericStreamingChunk]:
    state: dict = {}
    for line in response.iter_lines():
        if not line or line.startswith(":") or not line.startswith("data: "):
            continue
        raw = line[6:]
        if raw == "[DONE]":
            break
        try:
            event = json.loads(raw)
        except json.JSONDecodeError:
            continue
        chunk = _event_chunk(event, state)
        if chunk is not None:
            yield chunk
        if chunk is not None and chunk["is_finished"]:
            break


async def _aiter_sse(response: httpx.Response) -> AsyncIterator[GenericStreamingChunk]:
    state: dict = {}
    async for line in response.aiter_lines():
        if not line or line.startswith(":") or not line.startswith("data: "):
            continue
        raw = line[6:]
        if raw == "[DONE]":
            break
        try:
            event = json.loads(raw)
        except json.JSONDecodeError:
            continue
        chunk = _event_chunk(event, state)
        if chunk is not None:
            yield chunk
        if chunk is not None and chunk["is_finished"]:
            break


# ---------------------------------------------------------------------------
# LiteLLM CustomLLM subclass
# ---------------------------------------------------------------------------

class CodexProvider(CustomLLM):

    def completion(
        self,
        model: str,
        messages: list,
        api_base: str,
        custom_prompt_dict: dict,
        model_response: ModelResponse,
        print_verbose: Callable,
        encoding,
        api_key,
        logging_obj,
        optional_params: dict,
        acompletion=None,
        litellm_params=None,
        logger_fn=None,
        headers={},
        timeout: Optional[Union[float, httpx.Timeout]] = None,
        client: Optional[HTTPHandler] = None,
    ) -> ModelResponse:
        model_name = model.split("/", 1)[-1] if "/" in model else model
        body = _build_body(model_name, messages, optional_params, stream=True)
        text_parts = []
        tool_calls: dict = {}
        prompt_tokens = completion_tokens = 0
        for attempt in range(2):
            try:
                auth = _get_auth(force=(attempt == 1))
            except Exception as e:
                raise CustomLLMError(status_code=401, message=f"Codex auth error: {e}")
            try:
                with httpx.Client(timeout=120) as c:
                    with c.stream("POST", API_URL, json=body, headers=_auth_headers(auth)) as resp:
                        resp.raise_for_status()
                        for chunk in _iter_sse(resp):
                            text_parts.append(chunk["text"])
                            _accumulate_tool_chunk(tool_calls, chunk)
                            if chunk.get("usage"):
                                prompt_tokens = chunk["usage"].get("prompt_tokens", 0)
                                completion_tokens = chunk["usage"].get("completion_tokens", 0)
                break
            except httpx.HTTPStatusError as e:
                if attempt == 0 and e.response.status_code == 401:
                    text_parts.clear()
                    tool_calls.clear()
                    continue
                raise CustomLLMError(status_code=e.response.status_code, message=str(e))

        model_response.choices[0].message.content = "".join(text_parts)  # type: ignore
        if tool_calls:
            model_response.choices[0].message.tool_calls = [  # type: ignore
                tool_calls[index] for index in sorted(tool_calls)
            ]
        model_response.choices[0].finish_reason = (  # type: ignore
            "tool_calls" if tool_calls else "stop"
        )
        model_response.model = _resolve_model(model_name)
        model_response.usage = Usage(
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
            total_tokens=prompt_tokens + completion_tokens,
        )
        return model_response

    def streaming(
        self,
        model: str,
        messages: list,
        api_base: str,
        custom_prompt_dict: dict,
        model_response: ModelResponse,
        print_verbose: Callable,
        encoding,
        api_key,
        logging_obj,
        optional_params: dict,
        acompletion=None,
        litellm_params=None,
        logger_fn=None,
        headers={},
        timeout: Optional[Union[float, httpx.Timeout]] = None,
        client: Optional[HTTPHandler] = None,
    ) -> Iterator[GenericStreamingChunk]:
        model_name = model.split("/", 1)[-1] if "/" in model else model
        body = _build_body(model_name, messages, optional_params, stream=True)
        for attempt in range(2):
            try:
                auth = _get_auth(force=(attempt == 1))
            except Exception as e:
                raise CustomLLMError(status_code=401, message=f"Codex auth error: {e}")
            try:
                with httpx.Client(timeout=120) as c:
                    with c.stream("POST", API_URL, json=body, headers=_auth_headers(auth)) as resp:
                        resp.raise_for_status()
                        yield from _iter_sse(resp)
                return
            except httpx.HTTPStatusError as e:
                if attempt == 0 and e.response.status_code == 401:
                    continue
                raise CustomLLMError(status_code=e.response.status_code, message=str(e))

    async def acompletion(
        self,
        model: str,
        messages: list,
        api_base: str,
        custom_prompt_dict: dict,
        model_response: ModelResponse,
        print_verbose: Callable,
        encoding,
        api_key,
        logging_obj,
        optional_params: dict,
        acompletion=None,
        litellm_params=None,
        logger_fn=None,
        headers={},
        timeout: Optional[Union[float, httpx.Timeout]] = None,
        client: Optional[AsyncHTTPHandler] = None,
    ) -> ModelResponse:
        model_name = model.split("/", 1)[-1] if "/" in model else model
        body = _build_body(model_name, messages, optional_params, stream=True)
        text_parts = []
        tool_calls: dict = {}
        prompt_tokens = completion_tokens = 0
        for attempt in range(2):
            try:
                auth = _get_auth(force=(attempt == 1))
            except Exception as e:
                raise CustomLLMError(status_code=401, message=f"Codex auth error: {e}")
            try:
                async with httpx.AsyncClient(timeout=120) as c:
                    async with c.stream("POST", API_URL, json=body, headers=_auth_headers(auth)) as resp:
                        resp.raise_for_status()
                        async for chunk in _aiter_sse(resp):
                            text_parts.append(chunk["text"])
                            _accumulate_tool_chunk(tool_calls, chunk)
                            if chunk.get("usage"):
                                prompt_tokens = chunk["usage"].get("prompt_tokens", 0)
                                completion_tokens = chunk["usage"].get("completion_tokens", 0)
                break
            except httpx.HTTPStatusError as e:
                if attempt == 0 and e.response.status_code == 401:
                    text_parts.clear()
                    tool_calls.clear()
                    continue
                raise CustomLLMError(status_code=e.response.status_code, message=str(e))

        model_response.choices[0].message.content = "".join(text_parts)  # type: ignore
        if tool_calls:
            model_response.choices[0].message.tool_calls = [  # type: ignore
                tool_calls[index] for index in sorted(tool_calls)
            ]
        model_response.choices[0].finish_reason = (  # type: ignore
            "tool_calls" if tool_calls else "stop"
        )
        model_response.model = _resolve_model(model_name)
        model_response.usage = Usage(
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
            total_tokens=prompt_tokens + completion_tokens,
        )
        return model_response

    async def astreaming(
        self,
        model: str,
        messages: list,
        api_base: str,
        custom_prompt_dict: dict,
        model_response: ModelResponse,
        print_verbose: Callable,
        encoding,
        api_key,
        logging_obj,
        optional_params: dict,
        acompletion=None,
        litellm_params=None,
        logger_fn=None,
        headers={},
        timeout: Optional[Union[float, httpx.Timeout]] = None,
        client: Optional[AsyncHTTPHandler] = None,
    ) -> AsyncIterator[GenericStreamingChunk]:
        model_name = model.split("/", 1)[-1] if "/" in model else model
        body = _build_body(model_name, messages, optional_params, stream=True)
        for attempt in range(2):
            try:
                auth = _get_auth(force=(attempt == 1))
            except Exception as e:
                raise CustomLLMError(status_code=401, message=f"Codex auth error: {e}")
            try:
                async with httpx.AsyncClient(timeout=120) as c:
                    async with c.stream("POST", API_URL, json=body, headers=_auth_headers(auth)) as resp:
                        resp.raise_for_status()
                        async for chunk in _aiter_sse(resp):
                            yield chunk
                return
            except httpx.HTTPStatusError as e:
                if attempt == 0 and e.response.status_code == 401:
                    continue
                raise CustomLLMError(status_code=e.response.status_code, message=str(e))


codex_provider = CodexProvider()