"""Google Cloud Code Assist (Gemini CLI) OAuth flow."""

import asyncio
import base64
import json
import os
import time
from collections.abc import Awaitable, Callable
from typing import Any
from urllib.parse import parse_qs, urlencode, urlparse

import aiohttp
from aiohttp import web

from otto.log import get_logger

from .pkce import generate_challenge, generate_state, generate_verifier

# === Constants ===

_CLIENT_ID = base64.b64decode(
    "NjgxMjU1ODA5Mzk1LW9vOGZ0Mm9wcmRybnA5ZTNhcWY2YXYzaG1kaWIxMzVqLmFwcHMuZ29vZ2xldXNlcmNvbnRlbnQuY29t"
).decode("ascii")
_CLIENT_SECRET = base64.b64decode("R09DU1BYLTR1SGdNUG0tMW83U2stZ2VWNkN1NWNsWEZzeGw=").decode(
    "ascii"
)

_AUTHORIZE_URL = "https://accounts.google.com/o/oauth2/v2/auth"
_TOKEN_URL = "https://oauth2.googleapis.com/token"
_CALLBACK_PATH = "/oauth2callback"
_CALLBACK_HOST = "127.0.0.1"
_CALLBACK_PORT = 8085
_REDIRECT_URI = f"http://localhost:{_CALLBACK_PORT}{_CALLBACK_PATH}"
_SCOPES = [
    "https://www.googleapis.com/auth/cloud-platform",
    "https://www.googleapis.com/auth/userinfo.email",
    "https://www.googleapis.com/auth/userinfo.profile",
]

_CODE_ASSIST_ENDPOINT = "https://cloudcode-pa.googleapis.com"
_CALLBACK_TIMEOUT_SECONDS = 300
_EXPIRY_BUFFER_SECONDS = 300

_NODE_USER_AGENT = "google-api-nodejs-client/9.15.1"
_X_GOOG_API_CLIENT = "gl-node/22.17.0"

_TIER_FREE = "free-tier"
_PROVISIONING_MAX_ATTEMPTS = 60
_PROVISIONING_INITIAL_DELAY_SECONDS = 1.0
_PROVISIONING_MAX_DELAY_SECONDS = 5.0

log = get_logger(__name__)


def _env_truthy(name: str) -> bool:
    value = os.environ.get(name, "").strip().lower()
    return value in {"1", "true", "yes", "on"}


# === Helpers ===


def _metadata(*, include_project: str | None = None) -> dict[str, str]:
    metadata = {
        "ideType": "IDE_UNSPECIFIED",
        "platform": "PLATFORM_UNSPECIFIED",
        "pluginType": "GEMINI",
    }
    if include_project:
        metadata["duetProject"] = include_project
    return metadata


def _auth_headers(access_token: str) -> dict[str, str]:
    return {
        "Authorization": f"Bearer {access_token}",
        "Content-Type": "application/json",
        "User-Agent": _NODE_USER_AGENT,
        "X-Goog-Api-Client": _X_GOOG_API_CLIENT,
    }


def _env_project_id() -> str | None:
    for key in ("GOOGLE_CLOUD_PROJECT", "GOOGLE_CLOUD_PROJECT_ID"):
        value = os.environ.get(key, "").strip()
        if value:
            return value
    return None


def _is_security_policy_violated(raw_error: str) -> bool:
    if "SECURITY_POLICY_VIOLATED" in raw_error:
        return True

    try:
        payload = json.loads(raw_error)
    except json.JSONDecodeError:
        return False

    error = payload.get("error")
    if not isinstance(error, dict):
        return False

    details = error.get("details")
    if not isinstance(details, list):
        return False

    for detail in details:
        if not isinstance(detail, dict):
            continue
        reason = str(detail.get("reason", "")).upper()
        if reason == "SECURITY_POLICY_VIOLATED":
            return True

    return False


def _parse_auth_input(raw: str) -> dict[str, str | None]:
    value = raw.strip()
    if not value:
        return {}

    parsed = urlparse(value)
    if parsed.query:
        query = parse_qs(parsed.query)
        return {
            "code": query.get("code", [None])[0],
            "state": query.get("state", [None])[0],
        }

    if "#" in value:
        code, state = value.split("#", 1)
        return {"code": code or None, "state": state or None}

    return {"code": value}


# === Local Callback ===


async def _run_callback_server(
    expected_state: str | None = None,
) -> tuple[web.AppRunner, asyncio.Future[dict[str, str] | None]]:
    loop = asyncio.get_running_loop()
    result_future: asyncio.Future[dict[str, str] | None] = loop.create_future()

    async def _handle_callback(request: web.Request) -> web.Response:
        error = request.query.get("error")
        if error:
            if not result_future.done():
                result_future.set_result(None)
            return web.Response(
                text=f"<h1>Authentication failed</h1><p>{error}</p>",
                content_type="text/html",
            )

        code = request.query.get("code")
        state = request.query.get("state")
        if expected_state is not None and (not state or state != expected_state):
            return web.Response(status=400, text="State mismatch")
        if not code:
            return web.Response(status=400, text="Missing authorization code")

        if not result_future.done():
            result_future.set_result({"code": code, "state": state})
        return web.Response(
            text="<h1>Authentication successful</h1><p>You can close this tab.</p>",
            content_type="text/html",
        )

    app = web.Application()
    app.router.add_get(_CALLBACK_PATH, _handle_callback)

    runner = web.AppRunner(app)
    await runner.setup()
    site = web.TCPSite(runner, _CALLBACK_HOST, _CALLBACK_PORT)
    await site.start()
    return runner, result_future


# === Project Discovery + Provisioning ===


async def _poll_operation(
    session: aiohttp.ClientSession,
    *,
    operation_name: str,
    headers: dict[str, str],
    on_progress: Callable[[str], None] | None,
) -> dict[str, Any]:
    delay_seconds = _PROVISIONING_INITIAL_DELAY_SECONDS

    for attempt in range(_PROVISIONING_MAX_ATTEMPTS):
        if on_progress:
            on_progress(
                f"Waiting for provisioning (attempt {attempt + 1}/{_PROVISIONING_MAX_ATTEMPTS})..."
            )

        await asyncio.sleep(delay_seconds)

        async with session.get(
            f"{_CODE_ASSIST_ENDPOINT}/v1internal/{operation_name}",
            headers=headers,
        ) as response:
            if response.status != 200:
                raise RuntimeError(f"Provisioning poll failed ({response.status})")
            operation = await response.json()

        if operation.get("done"):
            return operation

        delay_seconds = min(delay_seconds * 2, _PROVISIONING_MAX_DELAY_SECONDS)

    raise RuntimeError("Provisioning did not complete in time.")


async def _discover_project(
    access_token: str,
    on_progress: Callable[[str], None] | None = None,
) -> str:
    env_project = _env_project_id()
    headers = _auth_headers(access_token)

    if on_progress:
        on_progress("Checking Cloud Code Assist project...")

    load_request: dict[str, Any] = {
        "metadata": _metadata(include_project=env_project),
    }
    if env_project:
        load_request["cloudaicompanionProject"] = env_project

    async with aiohttp.ClientSession() as session:
        async with session.post(
            f"{_CODE_ASSIST_ENDPOINT}/v1internal:loadCodeAssist",
            headers=headers,
            json=load_request,
        ) as response:
            if response.status != 200:
                error_text = await response.text()
                if _is_security_policy_violated(error_text):
                    if env_project:
                        return env_project
                    raise RuntimeError(
                        "Google account is restricted by VPC-SC. "
                        "Set GOOGLE_CLOUD_PROJECT or GOOGLE_CLOUD_PROJECT_ID and retry."
                    )

                raise RuntimeError(f"loadCodeAssist failed ({response.status}): {error_text}")

            load_data = await response.json()

        current_tier = load_data.get("currentTier")
        if current_tier:
            project_from_api = load_data.get("cloudaicompanionProject")
            if isinstance(project_from_api, str) and project_from_api.strip():
                return project_from_api
            if env_project:
                return env_project
            raise RuntimeError(
                "This account requires GOOGLE_CLOUD_PROJECT or GOOGLE_CLOUD_PROJECT_ID."
            )

        allowed_tiers = load_data.get("allowedTiers")
        default_tier: dict[str, Any] | None = None
        if isinstance(allowed_tiers, list):
            default_tier = next(
                (
                    tier
                    for tier in allowed_tiers
                    if isinstance(tier, dict) and bool(tier.get("isDefault"))
                ),
                None,
            )

        tier_id = str((default_tier or {}).get("id") or _TIER_FREE)
        if tier_id != _TIER_FREE and not env_project:
            raise RuntimeError(
                "This account requires GOOGLE_CLOUD_PROJECT or GOOGLE_CLOUD_PROJECT_ID."
            )

        if on_progress:
            on_progress("Provisioning Cloud Code Assist project...")

        onboard_request: dict[str, Any] = {
            "tierId": tier_id,
            "metadata": _metadata(include_project=env_project if tier_id != _TIER_FREE else None),
        }
        if tier_id != _TIER_FREE and env_project:
            onboard_request["cloudaicompanionProject"] = env_project

        async with session.post(
            f"{_CODE_ASSIST_ENDPOINT}/v1internal:onboardUser",
            headers=headers,
            json=onboard_request,
        ) as response:
            if response.status != 200:
                error_text = await response.text()
                if _is_security_policy_violated(error_text) and env_project:
                    return env_project
                raise RuntimeError(f"onboardUser failed ({response.status}): {error_text}")

            operation = await response.json()

        if not operation.get("done"):
            operation_name = str(operation.get("name", "")).strip()
            if not operation_name:
                raise RuntimeError("Provisioning response missing operation name")

            operation = await _poll_operation(
                session,
                operation_name=operation_name,
                headers=headers,
                on_progress=on_progress,
            )

    response_payload = operation.get("response")
    if isinstance(response_payload, dict):
        cloud_project = response_payload.get("cloudaicompanionProject")
        if isinstance(cloud_project, dict):
            project_id = cloud_project.get("id")
            if isinstance(project_id, str) and project_id.strip():
                return project_id

    if env_project:
        return env_project

    raise RuntimeError("Could not discover or provision a Google Cloud project.")


# === Public API ===


async def login(
    on_auth_url: Callable[[str], None],
    on_prompt_code: Callable[[], Awaitable[Any]] | None = None,
    on_progress: Callable[[str], None] | None = None,
    no_browser: bool = False,
    **_: Any,
) -> dict[str, Any]:
    """Run Google OAuth Authorization Code flow with PKCE."""
    verifier = generate_verifier()
    challenge = generate_challenge(verifier)
    state = generate_state()

    runner: web.AppRunner | None = None
    callback_result: asyncio.Future[dict[str, str] | None] | None = None
    force_manual = no_browser or _env_truthy("OTTO_OAUTH_NO_LOCAL_CALLBACK")

    try:
        if force_manual:
            if on_progress:
                on_progress(
                    "Local OAuth callback disabled by OTTO_OAUTH_NO_LOCAL_CALLBACK. Using manual code entry."
                )
        else:
            try:
                runner, callback_result = await _run_callback_server(state)
                if on_progress:
                    on_progress(
                        f"Listening for OAuth callback on http://localhost:{_CALLBACK_PORT}{_CALLBACK_PATH}"
                    )
            except OSError:
                runner = None
                callback_result = None
                if on_progress:
                    on_progress(
                        f"Could not bind callback port {_CALLBACK_PORT}. Falling back to manual code entry."
                    )

        params = urlencode(
            {
                "client_id": _CLIENT_ID,
                "response_type": "code",
                "redirect_uri": _REDIRECT_URI,
                "scope": " ".join(_SCOPES),
                "code_challenge": challenge,
                "code_challenge_method": "S256",
                "state": state,
                "access_type": "offline",
                "prompt": "consent",
            }
        )
        on_auth_url(f"{_AUTHORIZE_URL}?{params}")

        code: str | None = None
        if callback_result is not None:
            try:
                result = await asyncio.wait_for(callback_result, timeout=_CALLBACK_TIMEOUT_SECONDS)
                if result:
                    maybe_code = result.get("code")
                    if isinstance(maybe_code, str) and maybe_code:
                        code = maybe_code
            except TimeoutError:
                if on_progress:
                    on_progress(
                        "Callback timed out. Paste the authorization code or callback URL from the browser."
                    )

        if not code:
            if on_prompt_code is None:
                raise RuntimeError("No authorization code received")
            raw_input = await on_prompt_code()
            parsed = _parse_auth_input(str(raw_input))
            parsed_state = parsed.get("state")
            if parsed_state and parsed_state != state:
                raise RuntimeError("OAuth state mismatch")
            parsed_code = parsed.get("code")
            if isinstance(parsed_code, str) and parsed_code:
                code = parsed_code

        if not code:
            raise RuntimeError("No authorization code received")

        if on_progress:
            on_progress("Exchanging authorization code...")

        async with aiohttp.ClientSession() as session:
            async with session.post(
                _TOKEN_URL,
                data={
                    "client_id": _CLIENT_ID,
                    "client_secret": _CLIENT_SECRET,
                    "code": code,
                    "grant_type": "authorization_code",
                    "redirect_uri": _REDIRECT_URI,
                    "code_verifier": verifier,
                },
                headers={"Content-Type": "application/x-www-form-urlencoded"},
            ) as response:
                if response.status != 200:
                    error_text = await response.text()
                    raise RuntimeError(
                        f"Google token exchange failed ({response.status}): {error_text}"
                    )
                token_data = await response.json()

        access_token = str(token_data.get("access_token") or "").strip()
        refresh_value = str(token_data.get("refresh_token") or "").strip()

        expires_in = token_data.get("expires_in")
        try:
            expires_seconds = float(expires_in)
        except (TypeError, ValueError):
            raise RuntimeError("Google token exchange response missing fields") from None

        if not access_token or not refresh_value:
            raise RuntimeError("Google token exchange response missing fields")

        project_id = await _discover_project(access_token, on_progress=on_progress)

        credentials = {
            "provider": "google",
            "type": "oauth",
            "refresh": refresh_value,
            "access": access_token,
            "expires": time.time() + expires_seconds - _EXPIRY_BUFFER_SECONDS,
            "projectId": project_id,
        }
        log.debug("google oauth login succeeded", project_id=project_id)
        return credentials
    finally:
        if runner is not None:
            await runner.cleanup()


async def refresh_token(
    refresh: str,
    project_id: str | None = None,
    **kwargs: Any,
) -> dict[str, Any]:
    """Refresh Google OAuth credentials while preserving projectId."""
    project_id_alias = kwargs.get("projectId")
    if isinstance(project_id_alias, str) and project_id_alias.strip():
        project_id = project_id_alias

    if not isinstance(project_id, str) or not project_id.strip():
        raise RuntimeError("Missing Google Cloud project ID")

    refresh_value = refresh.strip()
    if not refresh_value:
        raise RuntimeError("Missing refresh token")

    async with aiohttp.ClientSession() as session:
        async with session.post(
            _TOKEN_URL,
            data={
                "client_id": _CLIENT_ID,
                "client_secret": _CLIENT_SECRET,
                "refresh_token": refresh_value,
                "grant_type": "refresh_token",
            },
            headers={"Content-Type": "application/x-www-form-urlencoded"},
        ) as response:
            if response.status != 200:
                error_text = await response.text()
                raise RuntimeError(f"Google token refresh failed ({response.status}): {error_text}")
            token_data = await response.json()

    access_token = str(token_data.get("access_token") or "").strip()
    next_refresh = str(token_data.get("refresh_token") or refresh_value).strip()

    expires_in = token_data.get("expires_in")
    try:
        expires_seconds = float(expires_in)
    except (TypeError, ValueError):
        raise RuntimeError("Google token refresh response missing fields") from None

    if not access_token or not next_refresh:
        raise RuntimeError("Google token refresh response missing fields")

    return {
        "provider": "google",
        "type": "oauth",
        "refresh": next_refresh,
        "access": access_token,
        "expires": time.time() + expires_seconds - _EXPIRY_BUFFER_SECONDS,
        "projectId": project_id,
    }


def get_api_key(credentials: dict[str, Any]) -> str:
    """Return JSON payload required by Google Cloud Code Assist clients."""
    return json.dumps(
        {
            "token": str(credentials.get("access") or ""),
            "projectId": str(credentials.get("projectId") or ""),
        }
    )
