"""OAuth2 authentication with AWS Cognito using PKCE flow.

Handles browser-based login, token caching, and automatic token refresh.
"""

from __future__ import annotations

import base64
import errno
import hashlib
import html
import json
import os
import threading
import uuid
import webbrowser
from http.server import BaseHTTPRequestHandler, HTTPServer
from pathlib import Path
from time import time
from typing import Any
from urllib.parse import parse_qs, quote, urlparse

import requests
from oauthlib.oauth2 import WebApplicationClient

from .cached_token import clear_cached_token, get_cached_token, set_cached_token
from .config import (
    CONFIG_DIR,
    get_auth_domain,
    get_web_url,
    write_secure_file,
)
from .config import (
    get_client_id as _config_client_id,
)
from .logging import log_auth

SCOPE: list[str] = []


def _get_cognito_urls() -> tuple[str, str]:
    """Get authorize and token URLs for current environment."""
    domain = get_auth_domain()
    authorize_url = f"https://{domain}/login"
    token_url = f"https://{domain}/oauth2/token"
    return authorize_url, token_url


def _get_client_id() -> str:
    """Get Cognito client ID for current environment."""
    return _config_client_id()


REDIRECT_HOST = "localhost"
REDIRECT_PORT = 6565
REDIRECT_PATH = "/callback"
REDIRECT_URI = f"http://{REDIRECT_HOST}:{REDIRECT_PORT}{REDIRECT_PATH}"

# Path on the Data Portal web app that renders the manual login confirmation
# page. Used as the OAuth redirect target when there is no local callback
# server (``--no-callback`` or browser/port fallback).
CLI_LOGIN_PATH = "/cli-login"


def _manual_redirect_uri() -> str:
    """OAuth redirect URI for the manual flow: the web app's ``/cli-login`` page."""
    return f"{get_web_url().rstrip('/')}{CLI_LOGIN_PATH}"


def _gated_login_url(authorization_url: str) -> str:
    """Wrap the Cognito authorization URL in the web app's ``/cli-login`` gate.

    Instead of sending the browser straight to Cognito's hosted UI, we send it
    to the web app first. That page checks the browser can actually store
    cookies (and that site cookie consent allows it) before forwarding to
    Cognito. The AWS hosted UI hangs silently when cookies are blocked; the gate
    turns that into an actionable message.
    """
    web = get_web_url().rstrip("/")
    return f"{web}{CLI_LOGIN_PATH}?auth={quote(authorization_url, safe='')}"


TOKEN_FILE_PATH = CONFIG_DIR / "token.json"

TOKEN_EXPIRY_BUFFER_SECONDS = 60  # Refresh tokens this many seconds before expiry

# Client type for headers (can be set by CLI to override default "api")
_client_type: str = "api"


def get_token() -> dict[str, Any]:
    """Get the current token, loading from cache or triggering login if needed."""
    token = _load_token()
    if token:
        return token
    login()
    token = get_cached_token(get_auth_domain())
    if token is None:
        raise RuntimeError("Login failed to cache token")
    return token


def set_client_type(client_type: str) -> None:
    """Set the client type for API request headers.

    Args:
        client_type: The client type identifier ("api", "cli", or "webapp").
    """
    global _client_type
    _client_type = client_type


def get_token_header() -> dict[str, str]:
    """Return headers for API requests including auth and client identification.

    Returns:
        Dict with Authorization, X-Client-Type, and X-Client-Version headers.
    """
    from . import __version__

    token = get_token()
    id_token = token.get("id_token")
    if not id_token:
        raise RuntimeError("No id_token available in token response")
    return {
        "Authorization": f"Bearer {id_token}",
        "X-Client-Type": _client_type,
        "X-Client-Version": __version__,
    }


def login(no_callback: bool = False) -> None:
    """Run OAuth2 authorization-code + PKCE flow against Cognito.

    If already logged in with a valid token, this is a no-op.

    Args:
        no_callback: If True, use the manual flow where the user logs in on the
                     web app's ``/cli-login`` page and pastes back the login
                     code it displays. Useful when the client cannot open the
                     callback port.
    """
    token = _load_token()
    if token:
        return

    client_id = _get_client_id()
    authorize_url, token_url = _get_cognito_urls()

    client = WebApplicationClient(client_id)
    expected_state = str(uuid.uuid4())

    # PKCE
    code_verifier = _generate_code_verifier()
    code_challenge = _generate_code_challenge(code_verifier)

    # Each flow builds the authorization URL with its own redirect_uri and
    # returns the redirect_uri actually used, since the token exchange below
    # must present the same value. The browser flow may fall back to the
    # manual flow (and thus a different redirect_uri) at runtime.
    if no_callback:
        code, redirect_uri = _manual_login_flow(
            client, authorize_url, expected_state, code_challenge
        )
    else:
        code, redirect_uri = _browser_login_flow(
            client, authorize_url, expected_state, code_challenge
        )

    # Exchange code for token
    _exchange_code_for_token(client, token_url, code, code_verifier, redirect_uri)


def _build_authorization_url(
    client: WebApplicationClient,
    authorize_url: str,
    redirect_uri: str,
    state: str,
    code_challenge: str,
) -> str:
    """Build the Cognito authorization URL with PKCE params for a redirect_uri."""
    return client.prepare_request_uri(
        authorize_url,
        redirect_uri=redirect_uri,
        scope=SCOPE,
        state=state,
        code_challenge=code_challenge,
        code_challenge_method="S256",
    )


def _browser_login_flow(
    client: WebApplicationClient,
    authorize_url: str,
    expected_state: str,
    code_challenge: str,
) -> tuple[str, str]:
    """Run the browser-based login flow with local callback server.

    Falls back to the manual flow if the callback port cannot be opened
    or if no browser is available (e.g., headless environment).

    Returns ``(authorization_code, redirect_uri)``.
    """
    redirect_uri = REDIRECT_URI
    authorization_url = _build_authorization_url(
        client, authorize_url, redirect_uri, expected_state, code_challenge
    )

    # Start local HTTP server (handles 1 request)
    try:
        httpd: HTTPServer = HTTPServer(
            (REDIRECT_HOST, REDIRECT_PORT), _OAuthCallbackHandler
        )
    except OSError as e:
        if e.errno == errno.EADDRINUSE:
            print(
                f"Warning: Cannot open callback port {REDIRECT_PORT} (already in use).\n"
                f"         Falling back to manual login.\n"
            )
            return _manual_login_flow(
                client, authorize_url, expected_state, code_challenge
            )
        raise
    httpd.timeout = 300  # 5 minute timeout for the socket

    def serve() -> None:
        """Handle a single OAuth callback request."""
        httpd.handle_request()

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

    # Route through the web app's /cli-login gate, which checks the browser can
    # store cookies before forwarding to Cognito's hosted UI.
    login_url = _gated_login_url(authorization_url)

    print("Opening browser for authentication...")
    print("If it doesn't open automatically, visit this URL:")
    print(login_url)

    browser_opened = webbrowser.open(login_url, new=1, autoraise=True)

    if not browser_opened:
        httpd.server_close()
        print(
            "\nWarning: Could not open browser.\n"
            "         Falling back to manual login.\n"
        )
        return _manual_login_flow(client, authorize_url, expected_state, code_challenge)

    # Wait for redirect with timeout
    server_thread.join(timeout=300)
    if server_thread.is_alive():
        httpd.shutdown()
        raise RuntimeError("Authentication timed out after 5 minutes")

    auth_error = getattr(httpd, "auth_error", None)
    if auth_error:
        raise RuntimeError(f"OAuth error: {auth_error}")

    code = getattr(httpd, "auth_code", None)
    if not code:
        raise RuntimeError("No authorization code received")

    return code, redirect_uri


def _manual_login_flow(
    client: WebApplicationClient,
    authorize_url: str,
    expected_state: str,
    code_challenge: str,
) -> tuple[str, str]:
    """Run the manual login flow for headless environments.

    The user logs in via a browser elsewhere, lands on the web app's
    ``/cli-login`` confirmation page, and pastes back the login code shown
    there.

    Returns ``(authorization_code, redirect_uri)``.
    """
    redirect_uri = _manual_redirect_uri()
    authorization_url = _build_authorization_url(
        client, authorize_url, redirect_uri, expected_state, code_challenge
    )
    # Route through the web app's /cli-login gate, which checks the browser can
    # store cookies before forwarding to Cognito's hosted UI.
    login_url = _gated_login_url(authorization_url)

    print("Visit this URL in a browser to log in:")
    print()
    print(f"  {login_url}")
    print()
    print("After logging in, the page will display a login code.")
    print("Copy that code and paste it below.")
    print()

    token_input = input("Paste login code: ").strip()

    code, state = _decode_paste_token(token_input)
    if state != expected_state:
        raise RuntimeError("Authorization state mismatch. Please try again.")

    return code, redirect_uri


def _decode_paste_token(token: str) -> tuple[str, str]:
    """Decode the login code pasted from the ``/cli-login`` page.

    The page encodes ``{"code", "state"}`` as a base64url JSON blob. As a
    convenience, a full redirect URL (containing ``?code=...&state=...``) is
    also accepted.

    Returns ``(code, state)``.
    """
    token = token.strip()

    # Courtesy fallback: a pasted redirect URL with a query string.
    if "?" in token:
        query = parse_qs(urlparse(token).query)
        code = query.get("code", [None])[0]
        state = query.get("state", [None])[0]
        if not code:
            raise RuntimeError(
                "Could not find an authorization code in the pasted URL. "
                "Make sure you copied the entire login code."
            )
        return code, state  # type: ignore[return-value]

    try:
        padding = "=" * (-len(token) % 4)
        raw = base64.urlsafe_b64decode(token + padding)
        data = json.loads(raw)
        return data["code"], data["state"]
    except Exception as e:
        raise RuntimeError(
            "Could not decode the login code. Make sure you copied the "
            "entire code shown on the login page."
        ) from e


def _exchange_code_for_token(
    client: WebApplicationClient,
    token_url: str,
    code: str,
    code_verifier: str,
    redirect_uri: str,
) -> None:
    """Exchange an authorization code for tokens and save them.

    ``redirect_uri`` must match the one used to obtain ``code`` (the local
    callback server for the browser flow, or the web app's ``/cli-login`` page
    for the manual flow).
    """
    final_token_url, headers, body = client.prepare_token_request(
        token_url,
        redirect_url=redirect_uri,
        code=code,
        include_client_id=True,
        code_verifier=code_verifier,
    )

    headers["Content-Type"] = "application/x-www-form-urlencoded"

    token_response = requests.post(
        final_token_url,
        headers=headers,
        data=body,
        timeout=30,
    )
    token_response.raise_for_status()

    client.parse_request_body_response(token_response.text)

    _save_token(client.token)

    log_auth("login", success=True)

    print("\nLogged in successfully.")


def logout() -> None:
    """Clear cached token for the current auth domain."""
    auth_domain = get_auth_domain()
    clear_cached_token(auth_domain)

    if TOKEN_FILE_PATH.exists():
        with open(TOKEN_FILE_PATH) as f:
            stored = json.load(f)

        if auth_domain in stored:
            del stored[auth_domain]
            if stored:
                write_secure_file(TOKEN_FILE_PATH, stored)
            else:
                TOKEN_FILE_PATH.unlink()
            log_auth("logout", success=True)
            print("Logged out successfully.")
        else:
            print("No active session.")
    else:
        print("No active session.")


def _generate_code_verifier(length: int = 64) -> str:
    """Generate a high-entropy PKCE code_verifier (43-128 chars)."""
    verifier = base64.urlsafe_b64encode(os.urandom(length)).decode("ascii")
    return verifier.rstrip("=")


def _generate_code_challenge(verifier: str) -> str:
    """Generate an S256 PKCE code_challenge from a verifier."""
    digest = hashlib.sha256(verifier.encode("ascii")).digest()
    challenge = base64.urlsafe_b64encode(digest).decode("ascii")
    return challenge.rstrip("=")


PAGES_DIR = Path(__file__).parent / "pages"


def _load_image_base64(filename: str) -> str:
    """Load an image file and return its base64-encoded contents."""
    path = PAGES_DIR / "images" / filename
    if path.exists():
        return base64.b64encode(path.read_bytes()).decode("ascii")
    return ""


def _load_svg_content(filename: str) -> str:
    """Load an SVG file and return its contents."""
    path = PAGES_DIR / "images" / filename
    if path.exists():
        return path.read_text()
    return ""


def _render_page(name: str, **kwargs: str) -> str:
    """Load an HTML template and substitute {{key}} placeholders with escaped values.

    Newlines in values are converted to <br> tags after escaping.
    """
    template = (PAGES_DIR / name).read_text()
    template = template.replace("{{logo_base64}}", _load_image_base64("logo.png"))
    template = template.replace("{{favicon_base64}}", _load_image_base64("favicon.png"))

    # Add status icon based on status
    status = kwargs.get("status", "")
    if status == "success":
        template = template.replace("{{status_icon}}", _load_svg_content("success.svg"))
    elif status == "error":
        template = template.replace("{{status_icon}}", _load_svg_content("error.svg"))

    for key, value in kwargs.items():
        escaped = html.escape(value).replace("\n", "<br>")
        template = template.replace("{{" + key + "}}", escaped)

    return template


class _OAuthCallbackHandler(BaseHTTPRequestHandler):
    """Handle the redirect back to REDIRECT_URI.

    Stores the auth code + state (or an error) on the HTTPServer instance.
    """

    def do_GET(self) -> None:
        """Process the OAuth callback request and extract query parameters."""
        parsed = urlparse(self.path)

        if parsed.path != REDIRECT_PATH:
            self.send_response(404)
            self.end_headers()
            self.wfile.write(b"Not Found")
            return

        query = parse_qs(parsed.query)

        # Attach data to the server instance
        self.server.auth_code = query.get("code", [None])[0]  # type: ignore[attr-defined]
        self.server.auth_error = query.get("error", [None])[0]  # type: ignore[attr-defined]
        self.server.auth_state = query.get("state", [None])[0]  # type: ignore[attr-defined]

        self.send_response(200)
        self.send_header("Content-Type", "text/html; charset=utf-8")
        self.end_headers()

        if self.server.auth_error:  # type: ignore[attr-defined]
            page = _render_page(
                "auth_complete.html",
                title="Authentication Failed",
                message="Authentication failed.\nYou may close this window.",
                status="error",
            )
        else:
            page = _render_page(
                "auth_complete.html",
                title="Authentication Complete",
                message="Authentication complete.\nYou may close this window.",
                status="success",
            )
        self.wfile.write(page.encode("utf-8"))

    def log_message(self, format: str, *args: Any) -> None:
        """Silence the default HTTP request logging."""


def _load_token() -> dict[str, Any] | None:
    """Load a token from disk and refresh it if expired (or about to expire)."""
    auth_domain = get_auth_domain()

    # Check in-memory cache first
    cached = get_cached_token(auth_domain)
    if cached:
        expires_at_raw = cached.get("expires_at")
        if expires_at_raw is not None:
            expires_at = int(expires_at_raw)
            if float(expires_at - TOKEN_EXPIRY_BUFFER_SECONDS) >= time():
                return cached

    # Load from disk
    token: dict[str, Any] | None = None
    if TOKEN_FILE_PATH.exists():
        with open(TOKEN_FILE_PATH) as f:
            stored = json.load(f)

        token = stored.get(auth_domain)

        if token:
            expires_at_raw = token.get("expires_at")
            if expires_at_raw is not None:
                expires_at = int(expires_at_raw)
                # Refresh if expired or expiring within buffer period
                if float(expires_at - TOKEN_EXPIRY_BUFFER_SECONDS) < time():
                    token = _refresh(token)
            else:
                # If we don't have an expires_at, assume refresh is needed if possible
                token = _refresh(token)

    if token:
        set_cached_token(auth_domain, token)
    return token


def _save_token(token: dict[str, Any]) -> None:
    """Persist the token to disk and update the in-memory cache."""
    auth_domain = get_auth_domain()

    # Ensure we have an absolute expiry timestamp for local caching.
    if "expires_at" not in token:
        expires_in = token.get("expires_in")
        if expires_in is not None:
            token["expires_at"] = int(time()) + int(expires_in)

    set_cached_token(auth_domain, token)

    # Merge into existing token file (preserving other domains' tokens)
    stored: dict[str, Any] = {}
    if TOKEN_FILE_PATH.exists():
        with open(TOKEN_FILE_PATH) as f:
            stored = json.load(f)
    stored[auth_domain] = token
    write_secure_file(TOKEN_FILE_PATH, stored)


def _refresh(old_token: dict[str, Any]) -> dict[str, Any] | None:
    """Refresh an existing token using its refresh_token, if present.

    Returns None if refresh fails (e.g., token expired, revoked, or no refresh_token).
    The caller should trigger a new login flow when this returns None.
    """
    refresh_token = old_token.get("refresh_token")
    if not refresh_token:
        log_auth("refresh", success=False, error="no refresh_token")
        return None

    _, token_url = _get_cognito_urls()
    client_id = _get_client_id()

    try:
        res = requests.post(
            token_url,
            headers={"Content-Type": "application/x-www-form-urlencoded"},
            data={
                "grant_type": "refresh_token",
                "client_id": client_id,
                "refresh_token": refresh_token,
            },
            timeout=30,
        )
    except requests.RequestException as e:
        log_auth("refresh", success=False, error=str(e))
        print(f"Token refresh failed: {e}. You may need to log in again.")
        return None

    if not res.ok:
        log_auth("refresh", success=False, error=f"status={res.status_code}")
        print(
            f"Token refresh failed (status {res.status_code}). You may need to log in again."
        )
        return None

    new_token: dict[str, Any] = res.json()

    # Cognito may omit refresh_token on refresh; carry it forward.
    if "refresh_token" not in new_token:
        new_token["refresh_token"] = refresh_token

    # Compute expires_at for local caching.
    if "expires_in" in new_token:
        new_token["expires_at"] = int(time()) + int(new_token["expires_in"])

    _save_token(new_token)
    log_auth("refresh", success=True)
    return new_token
