"""OAuth 2.0 + PKCE helpers for the luckee CLI."""

from __future__ import annotations

import base64
import hashlib
import json
import os
import secrets
import time
import urllib.error
import urllib.parse
import urllib.request
import webbrowser
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, Optional


_ENV_DEFAULTS: dict[str, dict[str, str]] = {
    "prod": {
        "oauth_base_url": "https://user-api.motse.ai/luckee/user/",
        "frontend_base_url": "https://motse.ai",
        "ws_url": "wss://core-agent-loop-950663559455.us-central1.run.app/api/v2/agent/ws/v3",
        "file_api_url": "https://file-api-main-950663559455.us-central1.run.app",
    },
    "staging": {
        "oauth_base_url": "https://staging-user-api.motse.ai/luckee/user/",
        "frontend_base_url": "https://luckee-frontend-staging-950663559455.us-central1.run.app",
        "ws_url": "wss://core-agent-loop-staging.motse.ai/api/v2/agent/ws",
        "file_api_url": "https://file-api-dev-950663559455.us-central1.run.app",
    },
}


def _resolve_env() -> str:
    return os.getenv("LUCKEE_ENV", "prod").strip().lower()


def _env_default(key: str) -> str:
    env = _resolve_env()
    return _ENV_DEFAULTS.get(env, _ENV_DEFAULTS["prod"])[key]


DEFAULT_OAUTH_CLIENT_ID = "luckee-cli"
DEFAULT_SCOPE = "all"
AUTH_FILE_PATH = Path.home() / ".luckee" / "auth.json"
TOKEN_EXPIRY_SKEW_SECONDS = 60
HTTP_TIMEOUT_SECONDS = 20
LOGIN_POLL_INTERVAL_SECONDS = 3
CODE_VERIFIER_NOT_READY_CODES = {425}
DEFAULT_HTTP_USER_AGENT = "luckee-cli/0.1 (+https://motse.ai) python-urllib"


class InvalidTokenError(Exception):
    """Raised when the auth server explicitly rejects a token."""

    def __init__(self, message: str = "Token is invalid or expired.") -> None:
        super().__init__(message)


@dataclass
class OAuthTokens:
    access_token: str
    refresh_token: str
    expires_in: int
    saved_at: float
    oauth_base_url: str
    client_id: str

    def to_dict(self) -> dict[str, Any]:
        return {
            "accessToken": self.access_token,
            "refreshToken": self.refresh_token,
            "expiresIn": self.expires_in,
            "saved_at": self.saved_at,
            "oauth_base_url": self.oauth_base_url,
            "client_id": self.client_id,
        }


@dataclass
class LoginFlowResult:
    tokens: OAuthTokens
    authorize_url: str
    redirect_uri: str


def resolve_oauth_base_url(override: Optional[str] = None) -> str:
    base_url = (
        override
        or os.getenv("LUCKEE_OAUTH_BASE_URL")
        or os.getenv("USER_API_ROOT")
        or _env_default("oauth_base_url")
    ).strip()
    if not base_url:
        raise RuntimeError("OAuth base URL is empty")
    if not base_url.endswith("/"):
        base_url += "/"
    return base_url


def resolve_frontend_base_url(override: Optional[str] = None) -> str:
    frontend_value = (override or os.getenv("LUCKEE_FRONTEND") or _env_default("frontend_base_url")).strip()
    if not frontend_value:
        raise RuntimeError("Frontend base URL is empty")
    return frontend_value


def resolve_oauth_client_id(override: Optional[str] = None) -> str:
    client_id = (override or os.getenv("LUCKEE_OAUTH_CLIENT_ID") or DEFAULT_OAUTH_CLIENT_ID).strip()
    if not client_id:
        raise RuntimeError("OAuth client_id is empty")
    return client_id


def generate_code_verifier() -> str:
    verifier = base64.urlsafe_b64encode(secrets.token_bytes(64)).decode("ascii").rstrip("=")
    return verifier[:128]


def generate_code_challenge(code_verifier: str) -> str:
    digest = hashlib.sha256(code_verifier.encode("utf-8")).digest()
    return base64.urlsafe_b64encode(digest).decode("ascii").rstrip("=")


def generate_state() -> str:
    return secrets.token_urlsafe(24)


def _join_url(base_url: str, path: str) -> str:
    return base_url.rstrip("/") + "/" + path.lstrip("/")


def _safe_parse_json(raw: str) -> dict[str, Any]:
    try:
        parsed = json.loads(raw)
    except json.JSONDecodeError as exc:
        raise RuntimeError(f"Invalid JSON response: {exc}") from exc
    if not isinstance(parsed, dict):
        raise RuntimeError("Unexpected response JSON type")
    return parsed


def _extract_error_message(payload: dict[str, Any]) -> str:
    message = payload.get("message")
    if isinstance(message, str) and message.strip():
        return message.strip()
    code = payload.get("code")
    if code is not None:
        return f"OAuth API error code={code}"
    return "OAuth API request failed"


def _extract_token_data(payload: dict[str, Any]) -> dict[str, Any]:
    if isinstance(payload.get("data"), dict):
        data = payload["data"]
        if "accessToken" in data:
            return data
    if "accessToken" in payload:
        return payload
    raise RuntimeError(_extract_error_message(payload))


def _get_json(
    url: str,
    timeout_seconds: float = HTTP_TIMEOUT_SECONDS,
    extra_headers: Optional[dict[str, str]] = None,
) -> dict[str, Any]:
    headers = {
        "Accept": "*/*",
        "User-Agent": os.getenv("LUCKEE_HTTP_USER_AGENT", DEFAULT_HTTP_USER_AGENT),
    }
    if extra_headers:
        headers.update(extra_headers)
    request = urllib.request.Request(
        url,
        headers=headers,
        method="GET",
    )
    try:
        with urllib.request.urlopen(request, timeout=timeout_seconds) as response:
            raw = response.read().decode("utf-8")
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"HTTP {exc.code} calling {url}: {detail}") from exc
    except urllib.error.URLError as exc:
        raise RuntimeError(f"Network error calling {url}: {exc}") from exc
    return _safe_parse_json(raw)


def fetch_lingxing_account(oauth_base_url: str, user_id: str, user_token: Optional[str] = None) -> Optional[str]:
    """Fetch the lingxing account name for *user_id* from the user-info endpoint.

    Returns the ``account`` field from the response data, or ``None`` on any
    error so callers can proceed gracefully when the service is unavailable.
    """
    if not user_id:
        return None
    url = _join_url(oauth_base_url, f"api/deductScore/getUserInfo.do?userId={urllib.parse.quote(user_id)}")
    extra_headers: Optional[dict[str, str]] = None
    if user_token:
        extra_headers = {"Luckee-Authorization": user_token}
    try:
        body = _get_json(url, extra_headers=extra_headers)
    except Exception:
        return None
    if str(body.get("code")) != "200":
        return None
    data = body.get("data")
    if not isinstance(data, dict):
        return None
    account = str(data.get("account") or "").strip()
    return account or None


def _post_json(
    url: str,
    body: dict[str, Any],
    timeout_seconds: float = HTTP_TIMEOUT_SECONDS,
    extra_headers: Optional[dict[str, str]] = None,
) -> dict[str, Any]:
    payload = json.dumps(body).encode("utf-8")
    headers = {
        "Content-Type": "application/json",
        "Accept": "application/json",
        "User-Agent": os.getenv("LUCKEE_HTTP_USER_AGENT", DEFAULT_HTTP_USER_AGENT),
    }
    if extra_headers:
        headers.update(extra_headers)
    request = urllib.request.Request(
        url,
        data=payload,
        headers=headers,
        method="POST",
    )
    try:
        with urllib.request.urlopen(request, timeout=timeout_seconds) as response:
            raw = response.read().decode("utf-8")
    except urllib.error.HTTPError as exc:
        detail = exc.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"HTTP {exc.code} calling {url}: {detail}") from exc
    except urllib.error.URLError as exc:
        raise RuntimeError(f"Network error calling {url}: {exc}") from exc
    return _safe_parse_json(raw)


def build_authorize_url(
    frontend_base_url: str,
    client_id: str,
    redirect_uri: str,
    state: str,
    code_challenge: str,
    scope: str = DEFAULT_SCOPE,
) -> str:
    # Support both forms:
    # 1) frontend root, e.g. https://example.com
    # 2) explicit authorize page, e.g. https://example.com/authLogin
    stripped = frontend_base_url.rstrip("/")
    if stripped.endswith("/authLogin"):
        authorize_url = stripped
    else:
        authorize_url = _join_url(frontend_base_url, "authLogin")
    query = urllib.parse.urlencode(
        {
            "client_id": client_id,
            "response_type": "code",
            "redirect_uri": redirect_uri,
            "state": state,
            "scope": scope,
            "code_challenge": code_challenge,
            "code_challenge_method": "S256",
        }
    )
    return f"{authorize_url}?{query}"


def exchange_code_for_tokens(
    oauth_base_url: str,
    client_id: str,
    code: str,
    redirect_uri: str,
    code_verifier: str,
) -> OAuthTokens:
    token_url = _join_url(oauth_base_url, "web/oauth/token.do")
    payload = _post_json(
        token_url,
        {
            "clientId": client_id,
            "code": code,
            "redirectUri": redirect_uri,
            "codeVerifier": code_verifier,
        },
    )
    token_data = _extract_token_data(payload)
    return OAuthTokens(
        access_token=str(token_data.get("accessToken") or ""),
        refresh_token=str(token_data.get("refreshToken") or ""),
        expires_in=int(token_data.get("expiresIn") or 3600),
        saved_at=time.time(),
        oauth_base_url=oauth_base_url,
        client_id=client_id,
    )


class _CodeVerifierNotReady(Exception):
    """Raised when getTokenByCodeVerifier returns 425 (auth not yet completed)."""


def get_token_by_code_verifier(
    oauth_base_url: str,
    client_id: str,
    code_verifier: str,
    on_response: Optional[Callable[[dict], None]] = None,
) -> OAuthTokens:
    """Exchange a PKCE code_verifier for tokens via the server-side callback flow.

    Raises ``_CodeVerifierNotReady`` when the server returns 425 (as int or
    string), indicating the user hasn't completed browser authorization yet.
    """
    url = _join_url(oauth_base_url, "web/oauth/getTokenByCodeVerifier.do")
    payload = _post_json(url, {"codeVerifier": code_verifier})
    if on_response is not None:
        on_response(payload)
    code = payload.get("code")
    if str(code) in {str(c) for c in CODE_VERIFIER_NOT_READY_CODES}:
        raise _CodeVerifierNotReady(_extract_error_message(payload))
    token_data = _extract_token_data(payload)
    return OAuthTokens(
        access_token=str(token_data.get("accessToken") or ""),
        refresh_token=str(token_data.get("refreshToken") or ""),
        expires_in=int(token_data.get("expiresIn") or 3600),
        saved_at=time.time(),
        oauth_base_url=oauth_base_url,
        client_id=client_id,
    )


def refresh_tokens(oauth_base_url: str, client_id: str, refresh_token: str) -> OAuthTokens:
    refresh_url = _join_url(oauth_base_url, "web/oauth/refresh.do")
    payload = _post_json(
        refresh_url,
        {
            "clientId": client_id,
            "refreshToken": refresh_token,
        },
    )
    token_data = _extract_token_data(payload)
    next_refresh_token = str(token_data.get("refreshToken") or refresh_token)
    return OAuthTokens(
        access_token=str(token_data.get("accessToken") or ""),
        refresh_token=next_refresh_token,
        expires_in=int(token_data.get("expiresIn") or 3600),
        saved_at=time.time(),
        oauth_base_url=oauth_base_url,
        client_id=client_id,
    )


def _ensure_auth_dir() -> None:
    AUTH_FILE_PATH.parent.mkdir(mode=0o700, parents=True, exist_ok=True)


def save_tokens(tokens: OAuthTokens) -> None:
    _ensure_auth_dir()
    tmp_path = AUTH_FILE_PATH.with_suffix(".tmp")
    tmp_path.write_text(json.dumps(tokens.to_dict(), ensure_ascii=True, indent=2), encoding="utf-8")
    os.chmod(tmp_path, 0o600)
    os.replace(tmp_path, AUTH_FILE_PATH)
    os.chmod(AUTH_FILE_PATH, 0o600)


def load_saved_tokens() -> Optional[OAuthTokens]:
    if not AUTH_FILE_PATH.exists():
        return None
    try:
        payload = _safe_parse_json(AUTH_FILE_PATH.read_text(encoding="utf-8"))
    except Exception:
        return None
    access_token = str(payload.get("accessToken") or "")
    refresh_token = str(payload.get("refreshToken") or "")
    if not access_token and not refresh_token:
        return None
    expires_in = int(payload.get("expiresIn") or 3600)
    saved_at = float(payload.get("saved_at") or 0)
    oauth_base_url = str(payload.get("oauth_base_url") or _env_default("oauth_base_url"))
    client_id = str(payload.get("client_id") or DEFAULT_OAUTH_CLIENT_ID)
    return OAuthTokens(
        access_token=access_token,
        refresh_token=refresh_token,
        expires_in=expires_in,
        saved_at=saved_at,
        oauth_base_url=resolve_oauth_base_url(oauth_base_url),
        client_id=client_id.strip() or DEFAULT_OAUTH_CLIENT_ID,
    )


def clear_saved_tokens() -> bool:
    """Delete locally cached OAuth tokens. Returns True if anything was removed."""
    removed = False
    if AUTH_FILE_PATH.exists():
        AUTH_FILE_PATH.unlink()
        removed = True

    tmp_path = AUTH_FILE_PATH.with_suffix(".tmp")
    if tmp_path.exists():
        tmp_path.unlink()
        removed = True

    auth_dir = AUTH_FILE_PATH.parent
    if auth_dir.exists():
        try:
            auth_dir.rmdir()
            removed = True
        except OSError:
            # Directory is not empty or cannot be removed; keep remaining files.
            pass
    return removed


def _token_is_still_valid(tokens: OAuthTokens) -> bool:
    if not tokens.access_token:
        return False
    if not tokens.saved_at or tokens.expires_in <= 0:
        return True
    return (tokens.saved_at + tokens.expires_in - TOKEN_EXPIRY_SKEW_SECONDS) > time.time()


def fetch_user_id_from_token(oauth_base_url: str, user_token: str) -> Optional[str]:
    """Validate a user token with the user service and return the effective user_id.

    Mirrors the server-side logic in UserService.validate_user_token:
    POST <oauth_base_url>/api/user/validUserToken.do  {"luckeeToken": token}
    """
    if not user_token:
        return None

    url = _join_url(oauth_base_url, "api/user/validUserToken.do")
    try:
        body = _post_json(
            url,
            {"luckeeToken": user_token},
            extra_headers={"Luckee-Authorization": user_token},
        )
    except Exception as exc:
        msg = str(exc).lower()
        if "401" in msg or "403" in msg or "unauthorized" in msg:
            raise InvalidTokenError(
                f"Token rejected by auth server. {exc}".strip()
            ) from exc
        return None

    code = body.get("code")
    if str(code) != "200":
        server_msg = body.get("msg") or body.get("message") or ""
        raise InvalidTokenError(
            f"Token rejected by auth server (code={code}). {server_msg}".strip()
        )

    data = body.get("data")
    if isinstance(data, dict):
        user_id = (
            str(data.get("userId") or data.get("user_id") or data.get("id") or "")
        ).strip()
    elif isinstance(data, bool):
        user_id = ""
    else:
        user_id = str(data or "").strip()

    if not user_id:
        return None
    if user_id == "0":
        return "luckee_evaluation_user"
    return user_id


def get_valid_token(oauth_base_url: str, client_id: str) -> Optional[str]:
    saved = load_saved_tokens()
    if saved is None:
        return None
    if saved.oauth_base_url != oauth_base_url or saved.client_id != client_id:
        return None
    if _token_is_still_valid(saved):
        return saved.access_token
    if not saved.refresh_token:
        return None
    try:
        refreshed = refresh_tokens(oauth_base_url=oauth_base_url, client_id=client_id, refresh_token=saved.refresh_token)
    except Exception:
        return None
    if not refreshed.access_token:
        return None
    save_tokens(refreshed)
    return refreshed.access_token


def force_refresh_token(oauth_base_url: str, client_id: str) -> Optional[str]:
    """Force-refresh the saved token, ignoring local expiry.

    Unlike ``get_valid_token`` this always calls the refresh endpoint even when
    the local expiry check says the token is still valid.  Use this when the
    server has already rejected the access token.
    """
    saved = load_saved_tokens()
    if saved is None or not saved.refresh_token:
        return None
    if saved.oauth_base_url != oauth_base_url or saved.client_id != client_id:
        return None
    try:
        refreshed = refresh_tokens(
            oauth_base_url=oauth_base_url,
            client_id=client_id,
            refresh_token=saved.refresh_token,
        )
    except Exception:
        return None
    if not refreshed.access_token:
        return None
    save_tokens(refreshed)
    return refreshed.access_token


def run_login_flow(
    oauth_base_url: str,
    client_id: str,
    frontend_base_url: Optional[str] = None,
    timeout_seconds: int = 180,
    open_browser: bool = True,
    on_authorize_url: Optional[Callable[[str], None]] = None,
    on_poll_response: Optional[Callable[[dict], None]] = None,
) -> LoginFlowResult:
    """Run the browser-based OAuth + PKCE login flow.

    Instead of starting a local HTTP server, this uses the server-side
    callback URL and polls ``getTokenByCodeVerifier.do`` until the user
    completes authorization in the browser (or the timeout expires).
    """
    code_verifier = generate_code_verifier()
    code_challenge = generate_code_challenge(code_verifier)
    state = generate_state()

    redirect_uri = _join_url(oauth_base_url, "web/oauth/callback.do")
    resolved_frontend_base_url = resolve_frontend_base_url(frontend_base_url)
    authorize_url = build_authorize_url(
        frontend_base_url=resolved_frontend_base_url,
        client_id=client_id,
        redirect_uri=redirect_uri,
        state=state,
        code_challenge=code_challenge,
    )

    if on_authorize_url is not None:
        on_authorize_url(authorize_url)
    if open_browser:
        webbrowser.open(authorize_url)

    deadline = time.time() + timeout_seconds
    while True:
        remaining = deadline - time.time()
        if remaining <= 0:
            raise RuntimeError(
                f"Login timed out after {timeout_seconds}s waiting for "
                "browser authorization. Please try again."
            )
        try:
            tokens = get_token_by_code_verifier(
                oauth_base_url=oauth_base_url,
                client_id=client_id,
                code_verifier=code_verifier,
                on_response=on_poll_response,
            )
        except _CodeVerifierNotReady:
            time.sleep(min(LOGIN_POLL_INTERVAL_SECONDS, remaining))
            continue
        except Exception:
            # Transient network error or unexpected server response — keep polling.
            time.sleep(min(LOGIN_POLL_INTERVAL_SECONDS, remaining))
            continue

        if not tokens.access_token:
            raise RuntimeError("Token endpoint returned empty access token")
        save_tokens(tokens)
        return LoginFlowResult(
            tokens=tokens,
            authorize_url=authorize_url,
            redirect_uri=redirect_uri,
        )
