"""OpenAI / Codex OAuth login and refresh 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 .pkce import generate_challenge, generate_state, generate_verifier

_CLIENT_ID = "app_EMoamEEZ73f0CkXaXp7hrann"
_AUTHORIZE_URL = "https://auth.openai.com/oauth/authorize"
_TOKEN_URL = "https://auth.openai.com/oauth/token"
_CALLBACK_PATH = "/auth/callback"
_CALLBACK_HOST = "127.0.0.1"
_CALLBACK_PORT = 1455
_REDIRECT_URI = f"http://localhost:{_CALLBACK_PORT}{_CALLBACK_PATH}"
_SCOPE = "openid profile email offline_access"
_JWT_CLAIM_PATH = "https://api.openai.com/auth"


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


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


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


def _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}


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

    async def _handle_callback(request: web.Request) -> web.Response:
        state = request.query.get("state")
        code = request.query.get("code")
        if 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="<p>Authentication successful. Return to your terminal.</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


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

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

    try:
        if force_manual:
            if on_progress is not None:
                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 is not None:
                    on_progress(
                        f"Listening for OAuth callback on http://localhost:{_CALLBACK_PORT}{_CALLBACK_PATH}"
                    )
            except OSError:
                runner = None
                callback_result = None
                if on_progress is not None:
                    on_progress(
                        f"Could not bind callback port {_CALLBACK_PORT}. Falling back to manual code entry."
                    )

        params = urlencode(
            {
                "response_type": "code",
                "client_id": _CLIENT_ID,
                "redirect_uri": _REDIRECT_URI,
                "scope": _SCOPE,
                "code_challenge": challenge,
                "code_challenge_method": "S256",
                "state": state,
                "codex_cli_simplified_flow": "true",
                "originator": "otto",
            }
        )
        on_auth_url(f"{_AUTHORIZE_URL}?{params}")

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

        if not code:
            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("State mismatch")
            parsed_code = parsed.get("code")
            if isinstance(parsed_code, str) and parsed_code:
                code = parsed_code

        if not code:
            raise RuntimeError("Missing authorization code")

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

        access = payload.get("access_token")
        refresh = payload.get("refresh_token")
        expires_in = payload.get("expires_in")
        try:
            expires_in_value = float(expires_in)
        except (TypeError, ValueError):
            raise RuntimeError("Token response missing required fields") from None

        if not isinstance(access, str) or not access:
            raise RuntimeError("Token response missing required fields")
        if not isinstance(refresh, str) or not refresh:
            raise RuntimeError("Token response missing required fields")

        account_id = _extract_account_id(access)
        if not account_id:
            raise RuntimeError("Failed to extract accountId from token")

        return {
            "provider": "openai",
            "type": "oauth",
            "refresh": refresh,
            "access": access,
            "expires": time.time() + expires_in_value,
            "accountId": account_id,
        }
    finally:
        if runner is not None:
            await runner.cleanup()


async def refresh_token(refresh: str) -> dict[str, Any]:
    """Refresh an OpenAI/Codex OAuth access token."""
    async with aiohttp.ClientSession() as session:
        async with session.post(
            _TOKEN_URL,
            data={
                "grant_type": "refresh_token",
                "refresh_token": refresh,
                "client_id": _CLIENT_ID,
            },
            headers={"Content-Type": "application/x-www-form-urlencoded"},
        ) as response:
            if response.status != 200:
                body = await response.text()
                raise RuntimeError(f"OpenAI token refresh failed ({response.status}): {body}")
            payload = await response.json()

    access = payload.get("access_token")
    new_refresh = payload.get("refresh_token")
    expires_in = payload.get("expires_in")
    try:
        expires_in_value = float(expires_in)
    except (TypeError, ValueError):
        raise RuntimeError("Token refresh response missing fields") from None

    if not isinstance(access, str) or not access:
        raise RuntimeError("Token refresh response missing fields")
    if not isinstance(new_refresh, str) or not new_refresh:
        raise RuntimeError("Token refresh response missing fields")

    account_id = _extract_account_id(access)
    if not account_id:
        raise RuntimeError("Failed to extract accountId from refreshed token")

    return {
        "provider": "openai",
        "type": "oauth",
        "refresh": new_refresh,
        "access": access,
        "expires": time.time() + expires_in_value,
        "accountId": account_id,
    }


def get_api_key(credentials: dict[str, Any]) -> str:
    """Extract usable API key from OpenAI OAuth credentials."""
    return str(credentials["access"])
