import json
import logging
import os
import time
import uuid
from pathlib import Path
from typing import Any, Dict, List, Optional

import requests

from marsys.models.adapters.base import APIProviderAdapter, AsyncBaseAPIAdapter
from marsys.models.adapters.google import GoogleAdapter
from marsys.models.response_models import (
    ErrorResponse,
    HarmonizedResponse,
    ResponseMetadata,
    ToolCall,
    UsageInfo,
)

logger = logging.getLogger(__name__)


# --- Google OAuth Adapter (Antigravity Gateway) ---

GOOGLE_OAUTH_MODELS = {
    "claude-sonnet-4-6",
    "claude-opus-4-5-thinking",
    "claude-opus-4-6-thinking",
    "gemini-3.1-pro-high",
    "gemini-3.1-pro-medium",
    "gemini-3.1-pro-low",
    "gemini-3-flash",
    "gemini-2.5-flash",
    "gemini-2.5-pro",
    "gemini-2.5-flash-thinking",
    "gpt-oss-120b-medium",
}

GOOGLE_OAUTH_MODEL_ALIASES = {
    "sonnet": "claude-sonnet-4-6",
    "opus": "claude-opus-4-6-thinking",
    "gemini": "gemini-3.1-pro-high",
    "gemini-high": "gemini-3.1-pro-high",
    "gemini-medium": "gemini-3.1-pro-medium",
    "gemini-low": "gemini-3.1-pro-low",
    "gemini-flash": "gemini-3-flash",
    "flash": "gemini-2.5-flash",
    "gpt-oss": "gpt-oss-120b-medium",
}


class DiscoveryResult:
    """Result from loadCodeAssist discovery call."""
    __slots__ = ("project_id", "current_tier", "paid_tier", "allowed_tiers", "needs_onboarding")

    def __init__(
        self,
        project_id: Optional[str],
        current_tier: Optional[str] = None,
        paid_tier: Optional[str] = None,
        allowed_tiers: Optional[List[Dict[str, Any]]] = None,
        needs_onboarding: bool = False,
    ):
        self.project_id = project_id
        self.current_tier = current_tier
        self.paid_tier = paid_tier
        self.allowed_tiers = allowed_tiers or []
        self.needs_onboarding = needs_onboarding

    @property
    def effective_tier(self) -> Optional[str]:
        return self.paid_tier or self.current_tier


class GoogleOAuthAdapter(GoogleAdapter):
    """
    Adapter for Google Antigravity API using Google OAuth credentials.

    Extends GoogleAdapter to reuse Gemini format conversion.

    Uses the loadCodeAssist endpoint to resolve the user's real GCP project ID
    instead of relying on hardcoded defaults.
    Wraps requests in Cloud Code envelope and uses SSE streaming.

    Endpoint: https://cloudcode-pa.googleapis.com/v1internal:streamGenerateContent?alt=sse
    Auth: Google OAuth 2.0 with PKCE
    """

    streaming = True

    _DEFAULT_VERSION = "1.18.3"
    _INFERENCE_ENDPOINTS = [
        "https://cloudcode-pa.googleapis.com",
        "https://daily-cloudcode-pa.sandbox.googleapis.com",
    ]
    _DISCOVERY_ENDPOINTS = [
        "https://cloudcode-pa.googleapis.com",
        "https://daily-cloudcode-pa.sandbox.googleapis.com",
        "https://autopush-cloudcode-pa.sandbox.googleapis.com",
    ]
    _REVOKED_PROJECT_ID = "rising-fact-p41fc"
    _PROJECT_CACHE_TTL = 86400  # 24 hours

    def __init__(self, model_name: str, **kwargs):
        # Resolve model alias
        resolved = GOOGLE_OAUTH_MODEL_ALIASES.get(model_name, model_name)

        if resolved not in GOOGLE_OAUTH_MODELS:
            raise ValueError(
                f"Model '{model_name}' (resolved: '{resolved}') is not supported by google-oauth. "
                f"Supported: {', '.join(sorted(GOOGLE_OAUTH_MODELS))}"
            )

        # Initialize GoogleAdapter parent with empty api_key (we use OAuth)
        super().__init__(
            model_name=resolved,
            api_key="",
            base_url=self._INFERENCE_ENDPOINTS[0],
            max_tokens=kwargs.pop("max_tokens", 8192),
            temperature=kwargs.pop("temperature", 0.7),
            thinking_budget=kwargs.pop("thinking_budget", 2000),
        )

        # Store credentials path for refresh
        self._credentials_path = kwargs.get("credentials_path")

        # Load OAuth credentials
        self.credentials = self._load_google_credentials(kwargs.get("credentials_path"))
        self.access_token = self.credentials["access_token"]

        # Auto-refresh token if needed
        if kwargs.get("auto_refresh", True):
            self._refresh_token_if_needed()

        # Resolve version
        self._version = os.getenv("MARSYS_ANTIGRAVITY_VERSION") or self._fetch_antigravity_version()

        # Detect platform for User-Agent
        import platform as _platform
        sys_name = _platform.system().lower()
        machine = _platform.machine().lower()
        self._platform = sys_name
        self._platform_upper = sys_name.capitalize()
        if machine in ("x86_64", "amd64"):
            self._arch = "amd64"
        elif machine in ("arm64", "aarch64"):
            self._arch = "arm64"
        else:
            self._arch = machine

        # Tier info from discovery (populated during project resolution)
        self._tier_info: Optional[DiscoveryResult] = None

        # Resolve project ID (replaces hardcoded fallback)
        self.project_id = self._get_or_resolve_project_id()

    def _load_google_credentials(self, credentials_path: Optional[str] = None) -> Dict[str, Any]:
        """Load Google OAuth credentials from credential store or default paths."""
        from marsys.agents.exceptions import AgentConfigurationError

        # Try explicit path first
        search_paths = []
        if credentials_path:
            search_paths.append(Path(credentials_path).expanduser())

        # Then try env var
        env_path = os.getenv("MARSYS_GOOGLE_OAUTH_PATH")
        if env_path:
            search_paths.append(Path(env_path).expanduser())

        # Then try default paths
        search_paths.extend([
            Path.home() / ".marsys" / "oauth" / "google" / "credentials.json",
            Path.home() / ".config" / "opencode" / "antigravity-accounts.json",
        ])

        for cred_path in search_paths:
            if not cred_path.exists():
                continue

            try:
                with open(cred_path) as f:
                    data = json.load(f)

                # Handle antigravity-accounts.json format (array of accounts)
                if isinstance(data, list) and data:
                    # Use first active account
                    account = data[0]
                    access_token = account.get("authToken") or account.get("access_token")
                    refresh_token = account.get("refreshToken") or account.get("refresh_token")
                    project_id = account.get("projectId") or account.get("project_id")
                    expires_at = account.get("expiresAt") or account.get("expires_at")
                elif isinstance(data, dict):
                    # MARSYS credential format
                    access_token = data.get("access_token")
                    refresh_token = data.get("refresh_token")
                    project_id = data.get("project_id")
                    expires_at = data.get("expires_at")
                else:
                    continue

                if not access_token:
                    continue

                # Check expiration
                if expires_at:
                    expires_at_sec = expires_at / 1000 if expires_at > 10**12 else expires_at
                    if time.time() > expires_at_sec:
                        # Token expired, try to refresh
                        logger.warning(f"Google OAuth token expired, will attempt refresh")

                # Extract project_id_resolved_at timestamp if available
                if isinstance(data, dict):
                    project_id_resolved_at = data.get("project_id_resolved_at", 0)
                else:
                    project_id_resolved_at = 0

                self._credentials_path = str(cred_path)
                return {
                    "access_token": access_token,
                    "refresh_token": refresh_token,
                    "project_id": project_id,
                    "project_id_resolved_at": project_id_resolved_at,
                    "expires_at": expires_at,
                }

            except (json.JSONDecodeError, OSError) as e:
                logger.debug(f"Failed to read credentials from {cred_path}: {e}")
                continue

        raise AgentConfigurationError(
            "Google OAuth credentials not found",
            config_field="credentials_path",
            suggestion=(
                "Run 'marsys oauth add <profile> --provider google-oauth' to authenticate, "
                "or set MARSYS_GOOGLE_OAUTH_PATH to your credentials file"
            )
        )

    def _refresh_token_if_needed(self) -> bool:
        """Refresh OAuth token if expiring within buffer."""
        try:
            from marsys.models.credentials import OAuthTokenRefresher
            refreshed = OAuthTokenRefresher.refresh_if_needed(
                self._credentials_path,
                "google-oauth"
            )
            if refreshed:
                # Reload access token from credentials file
                self.credentials = self._load_google_credentials(self._credentials_path)
                self.access_token = self.credentials["access_token"]
            return refreshed
        except Exception as e:
            logger.warning(f"Failed to auto-refresh Google OAuth token: {e}")
            return False

    def _fetch_antigravity_version(self) -> str:
        """Fetch latest Antigravity version from auto-updater, fallback to default."""
        try:
            resp = requests.get(
                "https://antigravity-auto-updater-974169037036.us-central1.run.app",
                timeout=5,
            )
            if resp.status_code == 200:
                version = resp.text.strip()
                if version and version[0].isdigit():
                    logger.debug(f"Fetched Antigravity version: {version}")
                    return version
        except Exception as e:
            logger.debug(f"Failed to fetch Antigravity version: {e}")
        return self._DEFAULT_VERSION

    def _get_or_resolve_project_id(self) -> str:
        """Get project ID from credentials cache or resolve via loadCodeAssist."""
        cached_id = self.credentials.get("project_id")
        resolved_at = self.credentials.get("project_id_resolved_at", 0)

        # Use cached project ID if it's not the revoked one and was resolved recently
        if (
            cached_id
            and cached_id != self._REVOKED_PROJECT_ID
            and (time.time() - resolved_at) < self._PROJECT_CACHE_TTL
        ):
            tier_label = getattr(self._tier_info, "effective_tier", "unknown") if self._tier_info else "cached"
            logger.info(f"GoogleOAuth using project_id={cached_id}, tier={tier_label}")
            return cached_id

        # Need to resolve
        resolved = self._resolve_project_id()
        if resolved:
            return resolved

        if cached_id and cached_id != self._REVOKED_PROJECT_ID:
            logger.warning("loadCodeAssist failed, using cached project_id")
            return cached_id

        logger.error(
            "Could not resolve project ID via loadCodeAssist and no valid cached ID. "
            "Re-authenticate with: marsys oauth add <profile> --provider google-oauth"
        )
        raise ValueError("No valid Google Cloud project ID available")

    def _resolve_project_id(self) -> Optional[str]:
        """
        Resolve the user's GCP project ID via loadCodeAssist.

        Tries production endpoint first, then daily sandbox and autopush.
        Only calls onboardUser if loadCodeAssist indicates the user has no
        currentTier AND allowedTiers are available (using the default tier,
        NOT a hardcoded "free-tier").

        Uses ideType=ANTIGRAVITY to match the Antigravity IDE's discovery calls.
        """
        metadata = {
            "ideType": "ANTIGRAVITY",
            "platform": "PLATFORM_UNSPECIFIED",
            "pluginType": "GEMINI",
        }
        headers = {
            "Authorization": f"Bearer {self.access_token}",
            "Content-Type": "application/json",
            "User-Agent": f"antigravity/{self._version} {self._platform}/{self._arch}",
            "X-Goog-Api-Client": "google-cloud-sdk vscode_cloudshelleditor/0.1",
            "Client-Metadata": json.dumps(metadata),
        }
        body = {"metadata": metadata}

        # Try loadCodeAssist on each discovery endpoint
        discovery: Optional[DiscoveryResult] = None
        for endpoint in self._DISCOVERY_ENDPOINTS:
            discovery = self._try_load_code_assist(endpoint, headers, body)
            if discovery and discovery.project_id:
                self._tier_info = discovery
                logger.info(
                    f"Discovery: project={discovery.project_id}, "
                    f"tier={discovery.effective_tier or 'unknown'}"
                )
                self._persist_project_id(discovery.project_id)
                return discovery.project_id

        # Only attempt onboardUser when loadCodeAssist succeeded but indicated
        # that onboarding is needed (no currentTier, no project) AND there are
        # allowed tiers to choose from.
        if discovery and discovery.needs_onboarding and discovery.allowed_tiers:
            default_tier = self._find_default_tier(discovery.allowed_tiers)
            if default_tier:
                logger.info(f"No current tier found, onboarding with default tier: {default_tier}")
                onboard_body = {"tierId": default_tier, "metadata": metadata}
                for endpoint in self._DISCOVERY_ENDPOINTS:
                    project_id = self._try_onboard_user(endpoint, headers, onboard_body)
                    if project_id:
                        self._persist_project_id(project_id)
                        return project_id

        return None

    def _try_load_code_assist(
        self, endpoint: str, headers: Dict[str, str], body: Dict[str, Any]
    ) -> Optional[DiscoveryResult]:
        """Try loadCodeAssist on a single endpoint. Returns DiscoveryResult or None."""
        url = f"{endpoint}/v1internal:loadCodeAssist"
        try:
            resp = requests.post(url, json=body, headers=headers, timeout=10)
            logger.info(f"loadCodeAssist {endpoint} -> {resp.status_code}: {resp.text[:1000]}")
            if resp.status_code != 200:
                return None
            data = resp.json()
            return self._extract_discovery_info(data)
        except Exception as e:
            logger.debug(f"loadCodeAssist failed on {endpoint}: {e}")
            return None

    def _try_onboard_user(
        self, endpoint: str, headers: Dict[str, str], body: Dict[str, Any]
    ) -> Optional[str]:
        """Try onboardUser to auto-provision a project."""
        url = f"{endpoint}/v1internal:onboardUser"
        try:
            resp = requests.post(url, json=body, headers=headers, timeout=15)
            if resp.status_code != 200:
                logger.debug(f"onboardUser {endpoint} returned {resp.status_code}: {resp.text[:500]}")
                return None
            data = resp.json()

            # Handle Long Running Operation response
            if not data.get("done", True) and data.get("name"):
                data = self._poll_operation(endpoint, headers, data["name"])
                if data is None:
                    return None

            return self._extract_project_id_from_data(data)
        except Exception as e:
            logger.debug(f"onboardUser failed on {endpoint}: {e}")
            return None

    def _poll_operation(
        self, endpoint: str, headers: Dict[str, str], operation_name: str,
        max_polls: int = 12, interval: float = 5.0,
    ) -> Optional[Dict[str, Any]]:
        """Poll a Long Running Operation until done or timeout."""
        url = f"{endpoint}/v1internal/{operation_name}"
        for i in range(max_polls):
            time.sleep(interval)
            try:
                resp = requests.get(url, headers=headers, timeout=10)
                if resp.status_code != 200:
                    logger.debug(f"LRO poll {i+1}/{max_polls} returned {resp.status_code}")
                    continue
                data = resp.json()
                if data.get("done", False):
                    return data.get("response", data)
            except Exception as e:
                logger.debug(f"LRO poll {i+1}/{max_polls} failed: {e}")
        logger.warning(f"LRO {operation_name} did not complete within {max_polls * interval}s")
        return None

    @staticmethod
    def _find_default_tier(allowed_tiers: List[Dict[str, Any]]) -> Optional[str]:
        """Find the default tier from allowedTiers list.

        Returns the tier with isDefault=true, or the first tier if none marked default.
        """
        if not allowed_tiers:
            return None
        for tier in allowed_tiers:
            if tier.get("isDefault", False):
                return tier.get("id") or tier.get("tierId")
        # No explicit default — use first tier
        first = allowed_tiers[0]
        return first.get("id") or first.get("tierId")

    @staticmethod
    def _extract_project_id_from_data(data: Dict[str, Any]) -> Optional[str]:
        """Extract project ID string from a response dict."""
        project = data.get("cloudaicompanionProject")
        if project is None:
            return None
        if isinstance(project, dict):
            return project.get("id")
        if isinstance(project, str) and project:
            return project
        return None

    @staticmethod
    def _extract_discovery_info(data: Dict[str, Any]) -> Optional[DiscoveryResult]:
        """Extract full discovery info from loadCodeAssist response.

        Parses cloudaicompanionProject, currentTier, paidTier, and allowedTiers
        to determine effective tier and whether onboarding is needed.
        """
        project_id = GoogleOAuthAdapter._extract_project_id_from_data(data)

        current_tier_data = data.get("currentTier")
        current_tier = None
        if isinstance(current_tier_data, dict):
            current_tier = current_tier_data.get("id") or current_tier_data.get("tierId")
        elif isinstance(current_tier_data, str):
            current_tier = current_tier_data

        paid_tier_data = data.get("paidTier")
        paid_tier = None
        if isinstance(paid_tier_data, dict):
            paid_tier = paid_tier_data.get("id") or paid_tier_data.get("tierId")
        elif isinstance(paid_tier_data, str):
            paid_tier = paid_tier_data

        allowed_tiers = data.get("allowedTiers", [])
        if not isinstance(allowed_tiers, list):
            allowed_tiers = []

        needs_onboarding = current_tier is None and project_id is None

        return DiscoveryResult(
            project_id=project_id,
            current_tier=current_tier,
            paid_tier=paid_tier,
            allowed_tiers=allowed_tiers,
            needs_onboarding=needs_onboarding,
        )

    def _persist_project_id(self, project_id: str) -> None:
        """Persist resolved project_id to credentials file."""
        if not self._credentials_path:
            return
        try:
            from marsys.models.credentials import update_google_oauth_project_id
            update_google_oauth_project_id(self._credentials_path, project_id)
            self.credentials["project_id"] = project_id
            self.credentials["project_id_resolved_at"] = int(time.time())
            logger.info(f"Resolved and cached project_id: {project_id}")
        except Exception as e:
            logger.warning(f"Failed to persist project_id: {e}")

    def get_headers(self) -> Dict[str, str]:
        """Return OAuth + Antigravity headers for inference requests.

        Includes X-Goog-Api-Client and Client-Metadata headers matching
        the Antigravity IDE's real inference behavior (per API spec).
        These headers are used by Google's backend for tier-based routing.
        """
        metadata = {
            "ideType": "ANTIGRAVITY",
            "platform": "PLATFORM_UNSPECIFIED",
            "pluginType": "GEMINI",
        }
        headers = {
            "Authorization": f"Bearer {self.access_token}",
            "Content-Type": "application/json",
            "User-Agent": f"antigravity/{self._version} {self._platform}/{self._arch}",
            "Accept": "text/event-stream",
            "X-Goog-Api-Client": "google-cloud-sdk vscode_cloudshelleditor/0.1",
            "Client-Metadata": json.dumps(metadata),
        }
        if self.model_name.endswith("-thinking"):
            headers["anthropic-beta"] = "interleaved-thinking-2025-05-14"
        return headers

    def _clean_tool_schemas(self, tools: List[Dict]) -> List[Dict]:
        """Strip unsupported JSON Schema fields for Antigravity API."""
        import copy

        unsupported_keys = {
            "const", "$ref", "$defs", "$schema", "default",
            "examples", "additionalProperties", "title"
        }

        def _clean(node):
            if isinstance(node, dict):
                cleaned = {}
                for k, v in node.items():
                    if k in unsupported_keys:
                        if k == "const":
                            cleaned["enum"] = [v]
                            continue
                        if k == "title" and not isinstance(v, str):
                            # Actual property named "title" (dict/list value),
                            # not a JSON Schema annotation (string value) — keep it.
                            cleaned[k] = _clean(v)
                            continue
                        continue
                    cleaned[k] = _clean(v)
                return cleaned
            elif isinstance(node, list):
                return [_clean(item) for item in node]
            return node

        return _clean(copy.deepcopy(tools))

    def format_request_payload(self, messages: List[Dict], **kwargs) -> Dict[str, Any]:
        """
        Build Cloud Code envelope wrapping the Gemini payload from super().

        Extracts system messages into a top-level systemInstruction field
        (required by the Antigravity gateway) and wraps in the Cloud Code envelope.
        """
        # Get standard Gemini payload from GoogleAdapter
        gemini_payload = super().format_request_payload(messages, **kwargs)

        # Clean tool schemas if present
        if "tools" in gemini_payload:
            gemini_payload["tools"] = self._clean_tool_schemas(gemini_payload["tools"])

        # Extract system instructions from contents into systemInstruction field.
        # GoogleAdapter converts system messages to user role - we need to pull
        # the first message out if it was originally a system message.
        self._extract_system_instruction(messages, gemini_payload)

        # Wrap in Cloud Code envelope (matching NoeFabris/opencode format).
        return {
            "project": self.project_id,
            "model": self.model_name,
            "request": gemini_payload,
            "requestType": "agent",
            "userAgent": "antigravity",
            "requestId": f"agent-{uuid.uuid4()}",
        }

    @staticmethod
    def _extract_system_instruction(
        original_messages: List[Dict], gemini_payload: Dict[str, Any]
    ) -> None:
        """
        Extract system messages from contents into a top-level systemInstruction field.

        The Antigravity gateway expects system instructions in a dedicated
        systemInstruction field rather than inline in contents.
        """
        if not original_messages:
            return

        # Collect all system messages from the original messages
        system_texts = []
        for msg in original_messages:
            if msg.get("role") == "system":
                content = msg.get("content", "")
                if isinstance(content, list):
                    # Multi-part content
                    parts_text = []
                    for part in content:
                        if isinstance(part, dict) and part.get("type") == "text":
                            parts_text.append(part.get("text", ""))
                        elif isinstance(part, str):
                            parts_text.append(part)
                    content = "\n".join(parts_text)
                if content:
                    system_texts.append(content)

        if not system_texts:
            return

        # Build the systemInstruction field
        combined_system = "\n\n".join(system_texts)
        gemini_payload["systemInstruction"] = {
            "parts": [{"text": combined_system}],
        }

        # Remove the corresponding converted messages from contents.
        # GoogleAdapter converts system -> user role, so the first N user messages
        # that originated from system messages need to be removed.
        contents = gemini_payload.get("contents", [])
        system_count = len(system_texts)
        removed = 0
        new_contents = []
        for msg_content in contents:
            if removed < system_count and msg_content.get("role") == "user":
                # Check if this user message text matches a system message
                parts = msg_content.get("parts", [])
                if parts and len(parts) == 1 and "text" in parts[0]:
                    text = parts[0]["text"]
                    if text in system_texts:
                        removed += 1
                        continue
            new_contents.append(msg_content)
        gemini_payload["contents"] = new_contents

    def get_endpoint_url(self, base: Optional[str] = None) -> str:
        """Return streaming endpoint URL."""
        host = base or self.base_url
        return f"{host}/v1internal:streamGenerateContent?alt=sse"

    def _sync_stream_response(
        self, endpoint: str, headers: Dict[str, str], payload: Dict[str, Any]
    ) -> Dict[str, Any]:
        """Handle synchronous SSE streaming from Antigravity API."""
        try:
            import httpx
        except ImportError:
            raise ImportError(
                "httpx is required for Google OAuth adapter. "
                "Install with: pip install httpx"
            )

        result = {
            "candidates": [],
            "usageMetadata": {},
        }

        accumulated_text = ""
        accumulated_tool_calls = []
        reasoning_details = []

        with httpx.Client(timeout=180.0) as client:
            with client.stream(
                "POST", endpoint, json=payload, headers=headers
            ) as response:

                if response.status_code != 200:
                    response.read()
                    body_text = response.text if hasattr(response, "text") else ""
                    logger.warning(
                        f"Inference {response.status_code}: {body_text[:2000]}"
                    )
                    return response

                for line in response.iter_lines():
                    if line.startswith("data: "):
                        json_str = line[6:]
                    elif line.startswith("data:"):
                        json_str = line[5:]
                    else:
                        continue

                    try:
                        data = json.loads(json_str)

                        # Unwrap Cloud Code envelope
                        inner = data.get("response", data)

                        # Accumulate from candidates
                        candidates = inner.get("candidates", [])
                        if candidates:
                            candidate = candidates[0]
                            parts = candidate.get("content", {}).get("parts", [])
                            for part in parts:
                                if isinstance(part, dict):
                                    if "text" in part:
                                        accumulated_text += part["text"]
                                        if "thought_signature" in part:
                                            reasoning_details.append({
                                                "type": "text",
                                                "thought_signature": part["thought_signature"]
                                            })
                                    elif "functionCall" in part:
                                        fc = part["functionCall"]
                                        tool_id = f"call_{uuid.uuid4().hex[:8]}_{fc.get('name', 'unknown')}"
                                        accumulated_tool_calls.append({
                                            "id": tool_id,
                                            "functionCall": fc,
                                        })
                                        if "thoughtSignature" in part:
                                            reasoning_details.append({
                                                "type": "function_call",
                                                "tool_call_id": tool_id,
                                                "function_name": fc.get("name", ""),
                                                "thought_signature": part["thoughtSignature"]
                                            })

                        # Accumulate usage
                        usage = inner.get("usageMetadata")
                        if usage:
                            result["usageMetadata"] = usage

                    except json.JSONDecodeError:
                        continue

        # Build final response in Gemini format for parent's harmonize_response
        parts = []
        if accumulated_text:
            parts.append({"text": accumulated_text})
        for tc in accumulated_tool_calls:
            parts.append({"functionCall": tc["functionCall"]})

        if parts:
            result["candidates"] = [{
                "content": {"role": "model", "parts": parts},
                "finishReason": "STOP",
            }]

        # Attach reasoning details for later use
        result["_reasoning_details"] = reasoning_details

        return result

    async def _async_stream_response(
        self, endpoint: str, headers: Dict[str, str], payload: Dict[str, Any]
    ) -> Dict[str, Any]:
        """Handle async SSE streaming from Antigravity API."""
        try:
            import httpx
        except ImportError:
            raise ImportError(
                "httpx is required for Google OAuth adapter. "
                "Install with: pip install httpx"
            )

        result = {
            "candidates": [],
            "usageMetadata": {},
        }

        accumulated_text = ""
        accumulated_tool_calls = []
        reasoning_details = []

        async with httpx.AsyncClient(timeout=180.0) as client:
            async with client.stream(
                "POST", endpoint, json=payload, headers=headers
            ) as response:

                if response.status_code != 200:
                    await response.aread()
                    return response

                async for line in response.aiter_lines():
                    if line.startswith("data: "):
                        json_str = line[6:]
                    elif line.startswith("data:"):
                        json_str = line[5:]
                    else:
                        continue

                    try:
                        data = json.loads(json_str)

                        # Unwrap Cloud Code envelope
                        inner = data.get("response", data)

                        # Accumulate from candidates
                        candidates = inner.get("candidates", [])
                        if candidates:
                            candidate = candidates[0]
                            parts = candidate.get("content", {}).get("parts", [])
                            for part in parts:
                                if isinstance(part, dict):
                                    if "text" in part:
                                        accumulated_text += part["text"]
                                        if "thought_signature" in part:
                                            reasoning_details.append({
                                                "type": "text",
                                                "thought_signature": part["thought_signature"]
                                            })
                                    elif "functionCall" in part:
                                        fc = part["functionCall"]
                                        tool_id = f"call_{uuid.uuid4().hex[:8]}_{fc.get('name', 'unknown')}"
                                        accumulated_tool_calls.append({
                                            "id": tool_id,
                                            "functionCall": fc,
                                        })
                                        if "thoughtSignature" in part:
                                            reasoning_details.append({
                                                "type": "function_call",
                                                "tool_call_id": tool_id,
                                                "function_name": fc.get("name", ""),
                                                "thought_signature": part["thoughtSignature"]
                                            })

                        # Accumulate usage
                        usage = inner.get("usageMetadata")
                        if usage:
                            result["usageMetadata"] = usage

                    except json.JSONDecodeError:
                        continue

        # Build final response in Gemini format
        parts = []
        if accumulated_text:
            parts.append({"text": accumulated_text})
        for tc in accumulated_tool_calls:
            parts.append({"functionCall": tc["functionCall"]})

        if parts:
            result["candidates"] = [{
                "content": {"role": "model", "parts": parts},
                "finishReason": "STOP",
            }]

        result["_reasoning_details"] = reasoning_details

        return result

    def harmonize_response(
        self, raw_response: Dict[str, Any], request_start_time: float
    ) -> HarmonizedResponse:
        """Unwrap Cloud Code envelope then delegate to GoogleAdapter."""
        # Extract reasoning details before parent strips them
        extra_reasoning = raw_response.pop("_reasoning_details", [])

        # Delegate to GoogleAdapter's harmonize_response
        response = super().harmonize_response(raw_response, request_start_time)

        # Override provider in metadata
        response.metadata.provider = "google-oauth"

        # Merge any additional reasoning details
        if extra_reasoning:
            existing = response.reasoning_details or []
            response.reasoning_details = existing + extra_reasoning if existing else extra_reasoning

        return response

    def _compute_retry_delay(self, attempt: int, response=None) -> float:
        """Compute retry delay with exponential backoff and jitter."""
        import random

        base_delay = 5.0

        # Check for retry-after headers
        if response is not None and hasattr(response, "headers"):
            retry_after = response.headers.get("retry-after")
            if retry_after:
                try:
                    return float(retry_after)
                except (ValueError, TypeError):
                    pass
            retry_after_ms = response.headers.get("retry-after-ms")
            if retry_after_ms:
                try:
                    return float(retry_after_ms) / 1000.0
                except (ValueError, TypeError):
                    pass

        # Check for RESOURCE_EXHAUSTED (per-minute quota) → longer backoff
        if response is not None and hasattr(response, "status_code"):
            if response.status_code == 429:
                # Try to detect RESOURCE_EXHAUSTED from response body
                try:
                    if hasattr(response, "text"):
                        body_text = response.text
                    elif hasattr(response, "_content"):
                        body_text = response._content.decode("utf-8", errors="replace")
                    else:
                        body_text = ""
                    if "RESOURCE_EXHAUSTED" in body_text:
                        base_delay = 60.0
                except Exception:
                    base_delay = 30.0

        delay = base_delay * (2 ** attempt)
        delay = min(delay, 120.0)  # Cap at 2 minutes
        jitter = delay * 0.3 * (random.random() * 2 - 1)  # ±30% jitter
        return max(1.0, delay + jitter)

    def run_streaming(self, messages: List[Dict], **kwargs) -> HarmonizedResponse:
        """Execute API request with synchronous SSE streaming and retry.

        Tries inference endpoints in order (production first, then sandbox fallback).
        Falls back to the next endpoint on connection errors.
        """
        from marsys.agents.exceptions import ModelAPIError

        max_retries = 3
        last_error = None
        headers = self.get_headers()
        payload = self.format_request_payload(messages, **kwargs)

        for attempt in range(max_retries + 1):
            request_start_time = time.time()

            # Pick endpoint: on retry after connection error, try next endpoint
            endpoint_idx = min(attempt, len(self._INFERENCE_ENDPOINTS) - 1)
            endpoint = self.get_endpoint_url(self._INFERENCE_ENDPOINTS[endpoint_idx])

            logger.info(
                f"Inference -> {endpoint}, model={payload.get('model')}, "
                f"project={payload.get('project')}"
            )

            try:
                raw_response = self._sync_stream_response(endpoint, headers, payload)

                if not isinstance(raw_response, dict):
                    status_code = getattr(raw_response, "status_code", 0)

                    # On 401, try token refresh once
                    if status_code == 401 and attempt == 0:
                        logger.info("Got 401, attempting token refresh")
                        self._refresh_token_if_needed()
                        headers = self.get_headers()
                        continue

                    # On 429 or 5xx, retry with backoff
                    if status_code in (429, 500, 502, 503, 504) and attempt < max_retries:
                        delay = self._compute_retry_delay(attempt, raw_response)
                        body_text = ""
                        if hasattr(raw_response, "text"):
                            body_text = raw_response.text[:500]
                        elif hasattr(raw_response, "_content"):
                            body_text = raw_response._content.decode("utf-8", errors="replace")[:500]
                        logger.warning(
                            f"Got {status_code} from Antigravity API ({endpoint}): {body_text}, "
                            f"retrying in {delay:.1f}s (attempt {attempt + 1}/{max_retries})"
                        )
                        time.sleep(delay)
                        continue

                    raise ModelAPIError.from_provider_response(
                        provider="google-oauth",
                        response=raw_response,
                    )

                return self.harmonize_response(raw_response, request_start_time)

            except ModelAPIError as e:
                last_error = e
                if e.is_retryable and attempt < max_retries:
                    delay = self._compute_retry_delay(attempt)
                    logger.warning(
                        f"Retryable error from Antigravity API: {e.developer_message}, "
                        f"retrying in {delay:.1f}s (attempt {attempt + 1}/{max_retries})"
                    )
                    time.sleep(delay)
                    continue
                raise
            except Exception as e:
                raise ModelAPIError.from_provider_response(
                    provider="google-oauth",
                    exception=e
                )

        raise last_error  # Should not reach here, but safety net

    def handle_api_error(self, error: Exception, response=None) -> ErrorResponse:
        """Handle Antigravity-specific API errors."""
        from marsys.agents.exceptions import ModelAPIError

        api_error = ModelAPIError.from_provider_response(
            provider="google-oauth", response=response, exception=error
        )

        if api_error.is_critical():
            raise api_error

        return ErrorResponse(
            error=api_error.developer_message,
            error_code=api_error.api_error_code,
            error_type=api_error.api_error_type,
            provider=api_error.provider,
            model=self.model_name,
            classification={
                "category": api_error.classification,
                "is_retryable": api_error.is_retryable,
                "retry_after": api_error.retry_after,
                "suggested_action": api_error.suggested_action,
            },
        )


class AsyncGoogleOAuthAdapter(AsyncBaseAPIAdapter, GoogleOAuthAdapter):
    """Async version of Google OAuth adapter."""

    streaming = True

    def __init__(self, model_name: str, **kwargs):
        GoogleOAuthAdapter.__init__(self, model_name, **kwargs)
        self._session = None

    async def arun_streaming(self, messages: List[Dict], **kwargs) -> HarmonizedResponse:
        """Execute async API request with SSE streaming and retry.

        Tries inference endpoints in order (matching sync run_streaming).
        Falls back to the next endpoint on connection errors.
        """
        import asyncio
        from marsys.agents.exceptions import ModelAPIError

        max_retries = 3
        last_error = None
        headers = self.get_headers()
        payload = self.format_request_payload(messages, **kwargs)

        for attempt in range(max_retries + 1):
            request_start_time = time.time()

            # Pick endpoint: on retry after connection error, try next endpoint
            endpoint_idx = min(attempt, len(self._INFERENCE_ENDPOINTS) - 1)
            endpoint = self.get_endpoint_url(self._INFERENCE_ENDPOINTS[endpoint_idx])

            try:
                raw_response = await self._async_stream_response(endpoint, headers, payload)

                if not isinstance(raw_response, dict):
                    status_code = getattr(raw_response, "status_code", 0)

                    # On 401, try token refresh once
                    if status_code == 401 and attempt == 0:
                        logger.info("Got 401, attempting token refresh")
                        self._refresh_token_if_needed()
                        headers = self.get_headers()
                        continue

                    if status_code in (429, 500, 502, 503, 504) and attempt < max_retries:
                        delay = self._compute_retry_delay(attempt, raw_response)
                        body_text = ""
                        if hasattr(raw_response, "text"):
                            body_text = raw_response.text[:500]
                        elif hasattr(raw_response, "_content"):
                            body_text = raw_response._content.decode("utf-8", errors="replace")[:500]
                        logger.warning(
                            f"Got {status_code} from Antigravity API ({endpoint}): {body_text}, "
                            f"retrying in {delay:.1f}s (attempt {attempt + 1}/{max_retries})"
                        )
                        await asyncio.sleep(delay)
                        continue

                    raise ModelAPIError.from_provider_response(
                        provider="google-oauth",
                        response=raw_response,
                    )

                return self.harmonize_response(raw_response, request_start_time)

            except ModelAPIError as e:
                last_error = e
                if e.is_retryable and attempt < max_retries:
                    delay = self._compute_retry_delay(attempt)
                    logger.warning(
                        f"Retryable error from Antigravity API: {e.developer_message}, "
                        f"retrying in {delay:.1f}s (attempt {attempt + 1}/{max_retries})"
                    )
                    await asyncio.sleep(delay)
                    continue
                raise
            except Exception as e:
                raise ModelAPIError.from_provider_response(
                    provider="google-oauth",
                    exception=e
                )

        raise last_error

