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

from __future__ import annotations

import base64
import html
import hashlib
import json
import os
import secrets
import threading
import time
import urllib.error
import urllib.parse
import urllib.request
import webbrowser
from dataclasses import dataclass
from http.server import BaseHTTPRequestHandler, HTTPServer
from pathlib import Path
from typing import Any, Callable, Optional


DEFAULT_OAUTH_BASE_URL = "https://staging-user-api.motse.ai/luckee/user/"
DEFAULT_FRONTEND_BASE_URL_STAGING = "https://luckee-frontend-staging-950663559455.us-central1.run.app"
DEFAULT_FRONTEND_BASE_URL_PROD = "https://motse.ai"
DEFAULT_FRONTEND_BASE_URL = DEFAULT_FRONTEND_BASE_URL_PROD
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
DEFAULT_HTTP_USER_AGENT = "luckee-cli/0.1 (+https://motse.ai) python-urllib"


@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 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 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 _post_json(url: str, body: dict[str, Any], timeout_seconds: float = HTTP_TIMEOUT_SECONDS) -> dict[str, Any]:
    payload = json.dumps(body).encode("utf-8")
    request = urllib.request.Request(
        url,
        data=payload,
        headers={
            "Content-Type": "application/json",
            "Accept": "application/json",
            "User-Agent": os.getenv("LUCKEE_HTTP_USER_AGENT", DEFAULT_HTTP_USER_AGENT),
        },
        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,
    )


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 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 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 run_login_flow(
    oauth_base_url: str,
    client_id: str,
    frontend_base_url: Optional[str] = None,
    preferred_port: int = 0,
    timeout_seconds: int = 180,
    open_browser: bool = True,
    on_authorize_url: Optional[Callable[[str], None]] = None,
) -> LoginFlowResult:
    code_verifier = generate_code_verifier()
    code_challenge = generate_code_challenge(code_verifier)
    expected_state = generate_state()
    callback_event = threading.Event()
    callback_data: dict[str, str] = {}

    class CallbackHandler(BaseHTTPRequestHandler):
        def log_message(self, format: str, *args: Any) -> None:
            return

        def do_GET(self) -> None:  # noqa: N802
            parsed = urllib.parse.urlparse(self.path)
            if parsed.path != "/callback":
                self.send_response(404)
                self.end_headers()
                self.wfile.write(b"Not Found")
                return
            query = urllib.parse.parse_qs(parsed.query)
            code = (query.get("code") or [None])[0]
            state = (query.get("state") or [None])[0]
            error = (query.get("error") or [None])[0]
            error_description = (query.get("error_description") or [None])[0]
            if code:
                callback_data["code"] = code
            if state:
                callback_data["state"] = state
            if error:
                callback_data["error"] = error
            if error_description:
                callback_data["error_description"] = error_description
            callback_event.set()

            status_code = 200
            title = "Luckee login successful."
            detail = "You can close this window and return to the CLI."

            if error:
                status_code = 400
                title = "Luckee login failed."
                reason = error_description or error
                detail = f"Authorization failed: {html.escape(reason)}"
            elif not code:
                status_code = 400
                title = "Luckee login failed."
                detail = "Authorization callback did not include a code."
            elif state != expected_state:
                status_code = 400
                title = "Luckee login failed."
                detail = "Authorization state mismatch. Please retry login from the CLI."

            self.send_response(status_code)
            self.send_header("Content-Type", "text/html; charset=utf-8")
            self.end_headers()
            self.wfile.write(
                (
                    "<html><body>"
                    f"<h3>{title}</h3>"
                    f"<p>{detail}</p>"
                    "<p>You can close this window and return to the CLI.</p>"
                    "</body></html>"
                ).encode("utf-8")
            )

    if preferred_port < 0 or preferred_port > 65535:
        raise RuntimeError("Invalid callback port")
    try:
        server = HTTPServer(("127.0.0.1", preferred_port), CallbackHandler)
    except OSError as exc:
        if preferred_port == 0:
            raise RuntimeError(f"Failed to bind local callback server: {exc}") from exc
        raise RuntimeError(
            f"Port {preferred_port} is unavailable for OAuth callback. "
            "Retry with a different --port or omit --port to auto-pick a free port."
        ) from exc

    callback_port = int(server.server_address[1])
    redirect_uri = f"http://localhost:{callback_port}/callback"
    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=expected_state,
        code_challenge=code_challenge,
    )

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

    server_thread = threading.Thread(target=server.serve_forever, daemon=True)
    server_thread.start()

    try:
        if not callback_event.wait(timeout=timeout_seconds):
            raise RuntimeError(
                f"Login timed out after {timeout_seconds}s while waiting for browser callback."
            )
        if callback_data.get("error"):
            raise RuntimeError(f"Authorization failed: {callback_data['error']}")
        code = callback_data.get("code")
        returned_state = callback_data.get("state")
        if not code:
            raise RuntimeError("Authorization callback missing code")
        if returned_state != expected_state:
            raise RuntimeError("Authorization state mismatch, login aborted")
        tokens = exchange_code_for_tokens(
            oauth_base_url=oauth_base_url,
            client_id=client_id,
            code=code,
            redirect_uri=redirect_uri,
            code_verifier=code_verifier,
        )
        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)
    finally:
        server.shutdown()
        server.server_close()
        server_thread.join(timeout=2.0)
