"""OAuth 2.1 provider for Crustdata MCP Server.

Implements a lightweight OAuth flow where users either paste their Crustdata
API key in a browser form or sign in with their Crustdata account via Auth0.
The API key is validated against the Crustdata API, then used as the Bearer
token for all subsequent MCP requests.

This enables per-user authentication when Claude.ai connects to the MCP server.
"""

import html as _html
import secrets
import time
from typing import Any
from urllib.parse import quote

import httpx
from starlette.requests import Request
from starlette.responses import HTMLResponse, RedirectResponse

from mcp.server.auth.provider import (
    AccessToken,
    AuthorizationCode,
    AuthorizationParams,
    RefreshToken,
    TokenError,
    construct_redirect_uri,
)
from mcp.shared.auth import OAuthClientInformationFull, OAuthToken

from .config import config

# In-memory stores (single-replica deployment)
_clients: dict[str, OAuthClientInformationFull] = {}
_auth_codes: dict[str, AuthorizationCode] = {}

# Map auth code -> user's Crustdata API key
_code_to_api_key: dict[str, str] = {}

# Revoked tokens: token -> revoked_at timestamp
_revoked_tokens: dict[str, float] = {}

# Pending authorization sessions: session_id -> (params, client, created_at)
_pending_authorizations: dict[str, tuple[AuthorizationParams, OAuthClientInformationFull, float]] = {}

# Rate limiting for login attempts: IP -> list of attempt timestamps
_login_attempts: dict[str, list[float]] = {}

AUTH_CODE_EXPIRY = 300  # 5 minutes
ACCESS_TOKEN_EXPIRY = 86400 * 30  # 30 days
PENDING_SESSION_EXPIRY = 600  # 10 minutes
REVOCATION_RETENTION = 86400 * 30  # keep revocations for 30 days
CLIENT_RETENTION = 86400  # 24 hours
LOGIN_RATE_LIMIT_WINDOW = 60  # 1 minute window
LOGIN_RATE_LIMIT_MAX = 5  # max 5 attempts per IP per window


def _generate_code() -> str:
    return secrets.token_urlsafe(32)


def _get_client_ip(request: Request) -> str:
    """Extract real client IP from the request.

    Behind AWS ALB, the real client IP is the rightmost entry in X-Forwarded-For
    (the ALB appends it). We read that directly instead of relying on uvicorn's
    proxy_headers which trusts the leftmost entry and is spoofable.
    """
    xff = request.headers.get("x-forwarded-for", "")
    if xff:
        return xff.rsplit(",", 1)[-1].strip()
    return request.client.host if request.client else "unknown"


def _check_rate_limit(request: Request) -> bool:
    """Check if the request IP has exceeded the login rate limit. Returns True if blocked."""
    client_ip = _get_client_ip(request)
    now = time.time()
    attempts = _login_attempts.get(client_ip, [])
    recent = [t for t in attempts if now - t < LOGIN_RATE_LIMIT_WINDOW]
    if len(recent) >= LOGIN_RATE_LIMIT_MAX:
        return True
    recent.append(now)
    _login_attempts[client_ip] = recent
    return False


def _cleanup_expired_sessions() -> None:
    """Remove expired pending authorizations and auth codes."""
    now = time.time()
    expired_sessions = [
        sid for sid, (_, _, created_at) in _pending_authorizations.items()
        if now - created_at > PENDING_SESSION_EXPIRY
    ]
    for sid in expired_sessions:
        del _pending_authorizations[sid]

    expired_codes = [
        code for code, ac in _auth_codes.items()
        if ac.expires_at and ac.expires_at < now
    ]
    for code in expired_codes:
        del _auth_codes[code]
        _code_to_api_key.pop(code, None)

    expired_revocations = [
        token for token, revoked_at in _revoked_tokens.items()
        if now - revoked_at > REVOCATION_RETENTION
    ]
    for token in expired_revocations:
        del _revoked_tokens[token]

    # Clean up old login attempt records
    stale_sessions = [
        sid for sid, attempts in _login_attempts.items()
        if not attempts or now - attempts[-1] > LOGIN_RATE_LIMIT_WINDOW
    ]
    for sid in stale_sessions:
        del _login_attempts[sid]

    # Clean up stale client registrations
    stale_clients = [
        cid for cid, c in _clients.items()
        if c.client_id_issued_at and (now - c.client_id_issued_at) > CLIENT_RETENTION
    ]
    for cid in stale_clients:
        del _clients[cid]


async def _validate_api_key(api_key: str) -> bool | None:
    """Validate a Crustdata API key. Returns True/False for valid/invalid, None on network error."""
    try:
        async with httpx.AsyncClient(timeout=10) as client:
            resp = await client.get(
                f"{config.api_url}/user/credits",
                headers={"Authorization": f"Token {api_key}"},
            )
            return resp.status_code == 200
    except Exception:
        return None


def _render_login_page(session_id: str, error: str = "", auth0_url: str = "") -> str:
    """Render the login page HTML with the given session_id and error."""
    error_display = "block" if error else "none"
    safe_session_id = _html.escape(session_id)
    safe_error = _html.escape(error)
    safe_auth0_url = _html.escape(auth0_url)
    auth0_section = ""
    if auth0_url:
        auth0_section = f"""
        <div style="display:flex;align-items:center;gap:12px;margin-top:24px;margin-bottom:4px">
            <div style="flex:1;height:1px;background:#ddd"></div>
            <span style="color:#999;font-size:13px">or</span>
            <div style="flex:1;height:1px;background:#ddd"></div>
        </div>
        <a href="{safe_auth0_url}" style="display:block;width:100%;padding:12px;background:#1a1a2e;color:white;
           border:none;border-radius:8px;font-size:15px;font-weight:600;cursor:pointer;margin-top:12px;
           text-align:center;text-decoration:none">Sign in with Crustdata account</a>"""
    return f"""<!DOCTYPE html>
<html>
<head>
    <meta charset="utf-8">
    <meta name="viewport" content="width=device-width, initial-scale=1">
    <title>Connect to Crustdata</title>
    <style>
        * {{ margin: 0; padding: 0; box-sizing: border-box; }}
        body {{ font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
               background: #f8f9fa; display: flex; justify-content: center; align-items: center;
               min-height: 100vh; padding: 20px; }}
        .card {{ background: white; border-radius: 12px; box-shadow: 0 2px 10px rgba(0,0,0,0.08);
                padding: 40px; max-width: 440px; width: 100%; }}
        .logo {{ font-size: 24px; font-weight: 700; color: #1a1a2e; margin-bottom: 8px; }}
        .subtitle {{ color: #666; font-size: 14px; margin-bottom: 24px; }}
        label {{ display: block; font-size: 13px; font-weight: 600; color: #333; margin-bottom: 6px; }}
        #api_key:focus {{ border-color: #4f46e5; box-shadow: 0 0 0 3px rgba(79,70,229,0.1); }}
        .btn {{ width: 100%; padding: 12px; background: #4f46e5; color: white; border: none;
               border-radius: 8px; font-size: 15px; font-weight: 600; cursor: pointer; margin-top: 16px; }}
        .btn:hover {{ background: #4338ca; }}
        .btn:disabled {{ background: #9ca3af; cursor: not-allowed; }}
        .error {{ color: #dc2626; font-size: 13px; margin-top: 12px; display: {error_display}; }}
        .hint {{ color: #999; font-size: 12px; margin-top: 8px; }}
    </style>
</head>
<body>
    <div class="card">
        <div class="logo">Crustdata</div>
        <div class="subtitle">Enter your API key to connect with Claude</div>
        <form id="form" method="POST">
            <input type="hidden" name="session_id" value="{safe_session_id}">
            <label for="api_key">API Key</label>
            <div style="position:relative">
                <input type="password" id="api_key" name="api_key" placeholder="Paste your Crustdata API key" required
                       style="width:100%;padding:10px 70px 10px 12px;border:1px solid #ddd;border-radius:8px;font-size:14px;font-family:monospace;outline:none">
                <button type="button" id="toggle" style="position:absolute;right:8px;top:50%;transform:translateY(-50%);background:none;border:1px solid #ddd;border-radius:4px;padding:4px 10px;font-size:12px;color:#666;cursor:pointer">Show</button>
            </div>
            <div class="hint">Find your API key at <a href="https://crustdata.com" target="_blank">crustdata.com</a></div>
            <button type="submit" class="btn" id="submit">Connect</button>
            <div class="error" id="error">{safe_error}</div>
        </form>{auth0_section}
    </div>
    <script>
        document.getElementById('form').addEventListener('submit', function() {{
            document.getElementById('submit').disabled = true;
            document.getElementById('submit').textContent = 'Validating...';
        }});
        document.getElementById('toggle').addEventListener('click', function() {{
            var input = document.getElementById('api_key');
            if (input.type === 'password') {{
                input.type = 'text';
                this.textContent = 'Hide';
            }} else {{
                input.type = 'password';
                this.textContent = 'Show';
            }}
        }});
    </script>
</body>
</html>"""


class CrustdataOAuthProvider:
    """OAuth 2.1 provider that lets users authenticate with their Crustdata API key
    or sign in with their Crustdata account via Auth0."""

    def __init__(self, base_url: str = "") -> None:
        self.base_url = base_url

    def _auth0_login_url(self, session_id: str) -> str:
        """Build the URL that redirects the user to Crustdata's Auth0 login.

        After Auth0 authentication, the Crustdata callback sets a `token` cookie
        on .crustdata.com and redirects the browser to our /oauth/auth0-callback.
        """
        callback = f"{self.base_url}/oauth/auth0-callback?session_id={session_id}"
        return f"{config.api_url}/auth0/login?redirect_url={quote(callback, safe='')}"

    def _complete_authorization(self, session_id: str, api_key: str) -> RedirectResponse | None:
        """Generate auth code for a validated API key and redirect to the client.

        Returns None if the session was already consumed (e.g. double-submit).
        """
        session = _pending_authorizations.pop(session_id, None)
        if not session:
            return None
        params, client, _ = session

        code = _generate_code()
        auth_code = AuthorizationCode(
            code=code,
            scopes=params.scopes or [],
            expires_at=time.time() + AUTH_CODE_EXPIRY,
            client_id=client.client_id or "",
            code_challenge=params.code_challenge,
            redirect_uri=params.redirect_uri,
            redirect_uri_provided_explicitly=params.redirect_uri_provided_explicitly,
            resource=params.resource,
        )
        _auth_codes[code] = auth_code
        _code_to_api_key[code] = api_key

        redirect_uri = construct_redirect_uri(
            str(params.redirect_uri),
            code=code,
            state=params.state,
        )
        return RedirectResponse(redirect_uri, status_code=302)

    async def get_client(self, client_id: str) -> OAuthClientInformationFull | None:
        return _clients.get(client_id)

    async def register_client(self, client_info: OAuthClientInformationFull) -> None:
        if not client_info.client_id:
            client_info.client_id = secrets.token_urlsafe(16)
        if not client_info.client_secret:
            client_info.client_secret = secrets.token_urlsafe(32)
        client_info.client_id_issued_at = int(time.time())
        _clients[client_info.client_id] = client_info

    async def authorize(
        self, client: OAuthClientInformationFull, params: AuthorizationParams
    ) -> str:
        # Clean up expired sessions before creating a new one
        _cleanup_expired_sessions()
        session_id = secrets.token_urlsafe(24)
        _pending_authorizations[session_id] = (params, client, time.time())
        return f"/oauth/login?session_id={session_id}"

    async def handle_login_get(self, request: Request) -> HTMLResponse:
        """Render the API key input form (with optional Auth0 sign-in button)."""
        session_id = request.query_params.get("session_id", "")
        if session_id not in _pending_authorizations:
            return HTMLResponse("<h1>Invalid or expired session</h1>", status_code=400)
        auth0_url = self._auth0_login_url(session_id) if self.base_url else ""
        return HTMLResponse(_render_login_page(session_id, auth0_url=auth0_url))

    async def handle_login_post(self, request: Request) -> HTMLResponse | RedirectResponse:
        """Handle API key submission."""
        form_data = await request.form()
        session_id = str(form_data.get("session_id", ""))
        api_key = str(form_data.get("api_key", "")).strip()

        if session_id not in _pending_authorizations:
            return HTMLResponse("<h1>Invalid or expired session</h1>", status_code=400)

        auth0_url = self._auth0_login_url(session_id) if self.base_url else ""

        if _check_rate_limit(request):
            return HTMLResponse(_render_login_page(session_id, "Too many attempts. Please wait a minute and try again.", auth0_url=auth0_url), status_code=429)

        if not api_key:
            return HTMLResponse(_render_login_page(session_id, "API key is required", auth0_url=auth0_url))

        # Validate the API key against Crustdata
        valid = await _validate_api_key(api_key)
        if valid is None:
            return HTMLResponse(_render_login_page(session_id, "Service temporarily unavailable. Please try again shortly.", auth0_url=auth0_url))
        if not valid:
            return HTMLResponse(_render_login_page(session_id, "Invalid API key. Please check and try again.", auth0_url=auth0_url))

        return self._complete_authorization(session_id, api_key) or HTMLResponse("<h1>Session already used</h1>", status_code=409)

    async def handle_auth0_callback(self, request: Request) -> HTMLResponse | RedirectResponse:
        """Handle redirect from Crustdata's Auth0 login.

        After Auth0 authentication, the Crustdata API callback redirects here
        and sets a `token` cookie (on .crustdata.com) containing the user's
        API key. We read that cookie and complete the OAuth authorization flow.
        """
        session_id = request.query_params.get("session_id", "")

        if session_id not in _pending_authorizations:
            return HTMLResponse("<h1>Invalid or expired session</h1>", status_code=400)

        auth0_url = self._auth0_login_url(session_id) if self.base_url else ""

        if _check_rate_limit(request):
            return HTMLResponse(_render_login_page(session_id, "Too many attempts. Please wait a minute and try again.", auth0_url=auth0_url), status_code=429)

        # The Crustdata Auth0 callback sets the user's API token as a cookie
        api_key = request.cookies.get("token", "").strip()

        if not api_key:
            return HTMLResponse(_render_login_page(
                session_id,
                error="Sign-in completed but no token received. Please try using your API key instead.",
                auth0_url=auth0_url,
            ))

        # Validate the token (guards against stale or injected cookies)
        valid = await _validate_api_key(api_key)
        if valid is None:
            return HTMLResponse(_render_login_page(
                session_id,
                error="Service temporarily unavailable. Please try again shortly.",
                auth0_url=auth0_url,
            ))
        if not valid:
            return HTMLResponse(_render_login_page(
                session_id,
                error="Sign-in succeeded but your API token appears invalid. Please try using your API key directly.",
                auth0_url=auth0_url,
            ))

        return self._complete_authorization(session_id, api_key) or HTMLResponse("<h1>Session already used</h1>", status_code=409)

    async def load_authorization_code(
        self, client: OAuthClientInformationFull, authorization_code: str
    ) -> AuthorizationCode | None:
        code = _auth_codes.get(authorization_code)
        if not code:
            return None
        if code.expires_at < time.time():
            del _auth_codes[authorization_code]
            _code_to_api_key.pop(authorization_code, None)
            return None
        if code.client_id != client.client_id:
            return None
        return code

    async def exchange_authorization_code(
        self, client: OAuthClientInformationFull, authorization_code: AuthorizationCode
    ) -> OAuthToken:
        api_key = _code_to_api_key.pop(authorization_code.code, None)
        if not api_key:
            raise TokenError(error="invalid_grant", error_description="Authorization code not found")

        # Clean up auth code
        _auth_codes.pop(authorization_code.code, None)

        # The access token IS the user's API key — this way it flows through
        # as the Bearer token and gets picked up by the existing middleware
        return OAuthToken(
            access_token=api_key,
            token_type="Bearer",
            expires_in=ACCESS_TOKEN_EXPIRY,
            scope=" ".join(authorization_code.scopes) if authorization_code.scopes else None,
        )

    async def load_refresh_token(
        self, client: OAuthClientInformationFull, refresh_token: str
    ) -> RefreshToken | None:
        return None

    async def exchange_refresh_token(
        self,
        client: OAuthClientInformationFull,
        refresh_token: RefreshToken,
        scopes: list[str],
    ) -> OAuthToken:
        raise TokenError(error="unsupported_grant_type", error_description="Refresh tokens not supported")

    async def load_access_token(self, token: str) -> AccessToken | None:
        # Reject empty, too-short, or revoked tokens
        if not token or len(token) < 10:
            return None
        if token in _revoked_tokens:
            return None
        # Accept any valid-looking token — it's the user's Crustdata API key.
        # Real validation happens when the Crustdata API is called.
        return AccessToken(
            token=token,
            client_id="crustdata-user",
            scopes=["read", "write"],
            expires_at=None,
        )

    async def revoke_token(self, token: Any) -> None:
        token_str = getattr(token, "token", None) or (token if isinstance(token, str) else None)
        if token_str:
            _revoked_tokens[token_str] = time.time()
