"""Google Cloud Code Assist OAuth streaming provider."""

import json
import random
import time
from collections.abc import AsyncIterator
from typing import Any

import aiohttp

from otto.auth.base import OAuthProvider, OAuthProviderError, ProviderChunk
from otto.auth.registry import register_provider
from otto.log import get_logger

# === Constants ===

_DEFAULT_ENDPOINT = "https://cloudcode-pa.googleapis.com"
_STREAM_PATH = "/v1internal:streamGenerateContent?alt=sse"
_NODE_USER_AGENT = "google-api-nodejs-client/9.15.1"
_X_GOOG_API_CLIENT = "gl-node/22.17.0"
_CLIENT_METADATA = json.dumps(
    {
        "ideType": "IDE_UNSPECIFIED",
        "platform": "PLATFORM_UNSPECIFIED",
        "pluginType": "GEMINI",
    }
)

log = get_logger(__name__)


# === Helpers ===


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


def _json_dumps_compact(value: Any) -> str:
    return json.dumps(value, separators=(",", ":"))


@register_provider
class GoogleCloudCodeProvider(OAuthProvider):
    """OAuth provider for `gemini-cli/*` models via Cloud Code Assist."""

    provider_name = "google_cloud_code"
    model_prefixes = ["gemini-cli"]

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

        if credentials.get("type") != "oauth":
            raise OAuthProviderError(
                "Google Cloud Code provider requires OAuth credentials. Run: otto auth login google"
            )

        access = credentials.get("access")
        if not isinstance(access, str) or not access.strip():
            raise OAuthProviderError(
                "Missing Google OAuth access token. Run: otto auth login google"
            )

        project_id = credentials.get("projectId")
        if not isinstance(project_id, str) or not project_id.strip():
            raise OAuthProviderError(
                "Missing Google project ID in credentials. "
                "Set GOOGLE_CLOUD_PROJECT and re-authenticate."
            )

        self._access_token = access
        self._project_id = project_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_effort: str | None = None,
        **kwargs: Any,
    ) -> AsyncIterator[ProviderChunk]:
        _ = reasoning_effort, kwargs
        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,
        )

        headers = {
            "Authorization": f"Bearer {self._access_token}",
            "Content-Type": "application/json",
            "Accept": "text/event-stream",
            "User-Agent": _NODE_USER_AGENT,
            "X-Goog-Api-Client": _X_GOOG_API_CLIENT,
            "Client-Metadata": _CLIENT_METADATA,
        }

        url = f"{_DEFAULT_ENDPOINT}{_STREAM_PATH}"

        try:
            async with aiohttp.ClientSession() as session:
                response = await self._request_with_retry(
                    session,
                    "POST",
                    url,
                    headers=headers,
                    json_body=request_body,
                )
                async with response:
                    if response.status != 200:
                        error_text = await response.text()
                        raise OAuthProviderError(
                            f"Cloud Code Assist request failed ({response.status}): {error_text}"
                        )

                    async for event in self._stream_response(response):
                        yield event
        except OAuthProviderError as exc:
            raise RuntimeError(str(exc))
        except Exception as exc:  # pragma: no cover - defensive path
            log.exception("google cloud code streaming error")
            raise RuntimeError(str(exc))

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

    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,
    ) -> dict[str, Any]:
        inner_request: dict[str, Any] = {
            "contents": self._convert_messages(messages),
        }

        if system_prompt:
            inner_request["systemInstruction"] = {
                "parts": [{"text": _sanitize_surrogates(system_prompt)}],
            }

        generation_config: dict[str, Any] = {}
        if temperature is not None:
            generation_config["temperature"] = temperature
        if max_tokens is not None:
            generation_config["maxOutputTokens"] = max_tokens
        if generation_config:
            inner_request["generationConfig"] = generation_config

        if session_id:
            inner_request["sessionId"] = session_id

        if tools:
            converted_tools = self._convert_tools(tools)
            if converted_tools:
                inner_request["tools"] = converted_tools
                if tool_choice is not None:
                    inner_request["toolConfig"] = {
                        "functionCallingConfig": {
                            "mode": self._map_tool_choice(tool_choice),
                        }
                    }

        request_id = f"otto-{int(time.time())}-{random.randint(100000, 999999)}"

        return {
            "project": self._project_id,
            "model": model_id,
            "request": inner_request,
            "userAgent": _NODE_USER_AGENT,
            "requestId": request_id,
        }

    def _convert_messages(self, messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
        converted: list[dict[str, Any]] = []
        pending_tool_parts: list[dict[str, Any]] = []

        def _flush_pending_tool_parts() -> None:
            if pending_tool_parts:
                converted.append({"role": "user", "parts": pending_tool_parts.copy()})
                pending_tool_parts.clear()

        for message in messages:
            role = str(message.get("role", "user"))
            content = message.get("content", "")

            if role == "system":
                continue

            if role == "tool":
                tool_name = str(message.get("name") or "tool")
                pending_tool_parts.append(
                    {
                        "functionResponse": {
                            "name": tool_name,
                            "response": {
                                "result": _sanitize_surrogates(str(content)),
                            },
                        }
                    }
                )
                continue

            _flush_pending_tool_parts()
            parts: list[dict[str, Any]] = []

            if isinstance(content, str):
                if content:
                    parts.append({"text": _sanitize_surrogates(content)})
            elif isinstance(content, list):
                for item in content:
                    if not isinstance(item, dict):
                        continue
                    if item.get("type") == "text":
                        text = str(item.get("text", ""))
                        if text:
                            parts.append({"text": _sanitize_surrogates(text)})
                    elif item.get("type") == "image_url":
                        image_obj = item.get("image_url")
                        if not isinstance(image_obj, dict):
                            continue
                        url = str(image_obj.get("url", ""))
                        if not url.startswith("data:"):
                            continue
                        metadata, _, encoded = url.partition(",")
                        mime_type = metadata.split(";")[0].replace("data:", "") or "image/png"
                        if encoded:
                            parts.append(
                                {
                                    "inlineData": {
                                        "mimeType": mime_type,
                                        "data": encoded,
                                    }
                                }
                            )

            if role == "assistant":
                role = "model"

                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 {}
                        name = str(function_data.get("name", ""))
                        args_raw = function_data.get("arguments", "{}")
                        try:
                            args = json.loads(str(args_raw))
                        except json.JSONDecodeError:
                            args = {}

                        function_call_part: dict[str, Any] = {
                            "functionCall": {
                                "name": name,
                                "args": args,
                            }
                        }

                        extra_content = tool_call.get("extra_content")
                        if isinstance(extra_content, dict):
                            google_extra = extra_content.get("google")
                            if isinstance(google_extra, dict):
                                thought_signature = str(
                                    google_extra.get("thought_signature", "")
                                ).strip()
                                if thought_signature:
                                    function_call_part["thoughtSignature"] = thought_signature

                        parts.append(function_call_part)

            if parts:
                converted.append({"role": role, "parts": parts})

        _flush_pending_tool_parts()
        return converted

    @staticmethod
    def _convert_tools(tools: list[dict[str, Any]]) -> list[dict[str, Any]]:
        function_declarations: 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 {}

            declaration: dict[str, Any] = {
                "name": str(function_data.get("name", "")),
                "description": str(function_data.get("description", "")),
            }

            parameters = function_data.get("parameters")
            if isinstance(parameters, dict):
                declaration["parameters"] = parameters

            function_declarations.append(declaration)

        if not function_declarations:
            return []

        return [{"functionDeclarations": function_declarations}]

    @staticmethod
    def _map_tool_choice(tool_choice: str | dict[str, Any]) -> str:
        if isinstance(tool_choice, str):
            return {
                "auto": "AUTO",
                "none": "NONE",
                "required": "ANY",
            }.get(tool_choice, "AUTO")

        if isinstance(tool_choice, dict):
            return "ANY"

        return "AUTO"

    async def _stream_response(
        self,
        response: aiohttp.ClientResponse,
    ) -> AsyncIterator[ProviderChunk]:
        buffer = ""
        tool_call_index = 0
        usage: dict[str, int] | None = None
        has_content = False
        stop_reason = "stop"

        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:].lstrip() for line in raw_event.split("\n") 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:
                    payload = json.loads(data_str)
                except json.JSONDecodeError:
                    continue

                response_payload = payload.get("response", payload)
                if not isinstance(response_payload, dict):
                    continue

                candidates = response_payload.get("candidates")
                if isinstance(candidates, list):
                    for candidate in candidates:
                        if not isinstance(candidate, dict):
                            continue

                        content = candidate.get("content")
                        if not isinstance(content, dict):
                            continue

                        parts = content.get("parts")
                        if not isinstance(parts, list):
                            continue

                        for part in parts:
                            if not isinstance(part, dict):
                                continue

                            text = part.get("text")
                            if isinstance(text, str) and text:
                                has_content = True
                                if part.get("thought"):
                                    continue
                                else:
                                    yield self._text_chunk(text)

                            function_call = part.get("functionCall")
                            if isinstance(function_call, dict):
                                has_content = True
                                name = str(function_call.get("name", ""))
                                args = function_call.get("args", {})
                                call_id = f"call_{tool_call_index}"
                                arguments = _json_dumps_compact(args)
                                thought_signature_raw = part.get("thoughtSignature")
                                thought_signature = (
                                    str(thought_signature_raw).strip()
                                    if isinstance(thought_signature_raw, str)
                                    else None
                                )

                                yield self._tool_call_chunk(
                                    tool_call_index,
                                    call_id,
                                    name,
                                    arguments,
                                    thought_signature=thought_signature,
                                )
                                tool_call_index += 1

                        finish_reason = candidate.get("finishReason")
                        if isinstance(finish_reason, str) and finish_reason:
                            stop_reason = self._map_stop_reason(finish_reason)

                usage_metadata = response_payload.get("usageMetadata")
                if isinstance(usage_metadata, dict):
                    usage = {
                        "prompt_tokens": int(usage_metadata.get("promptTokenCount", 0)),
                        "completion_tokens": int(usage_metadata.get("candidatesTokenCount", 0))
                        + int(usage_metadata.get("thoughtsTokenCount", 0)),
                        "total_tokens": int(usage_metadata.get("totalTokenCount", 0)),
                    }

        if not has_content:
            raise OAuthProviderError("Cloud Code Assist API returned an empty response")

        yield self._final_chunk(stop_reason, usage)

    @staticmethod
    def _map_stop_reason(finish_reason: str) -> str:
        return {
            "STOP": "stop",
            "MAX_TOKENS": "length",
            "SAFETY": "content_filter",
            "RECITATION": "content_filter",
            "OTHER": "stop",
        }.get(finish_reason, "stop")
