"""OAuth Authorization Server provider for FastMCP.

This enables Codex/MCP clients to authenticate via standard OAuth endpoints
served by FastMCP when `auth_server_provider` is configured:
- `/.well-known/oauth-authorization-server`
- `/authorize`
- `/token`
- `/register` (dynamic client registration)

Implementation notes:
- Dynamic clients are cached in-memory and persisted in auth storage.
- Authorization codes are short-lived, in-memory, single-use.
- Access tokens are Sibyl JWT access tokens (Bearer).
- Refresh tokens are JWT refresh tokens tracked through the active auth runtime.
"""

from __future__ import annotations

import secrets
import time
from dataclasses import dataclass
from datetime import UTC, datetime, timedelta
from html import escape
from typing import Protocol, cast
from urllib.parse import urlencode, urlsplit, urlunsplit
from uuid import UUID, uuid4

import jwt
from mcp.server.auth.provider import (
    AccessToken,
    AuthorizationCode,
    AuthorizationParams,
    OAuthAuthorizationServerProvider,
    RefreshToken,
    TokenError,
)
from mcp.shared.auth import OAuthClientInformationFull, OAuthToken
from starlette.requests import Request
from starlette.responses import HTMLResponse, RedirectResponse, Response

from sibyl import config as config_module
from sibyl.auth.api_key_common import ApiKeyAuth
from sibyl.auth.jwt import JwtError, create_access_token, verify_access_token
from sibyl.persistence.auth_runtime import (
    authenticate_api_key,
    authenticate_local_user,
    create_session_record,
    ensure_personal_organization,
    get_user_by_id,
    list_user_organizations,
    load_oauth_client_registration,
    load_refresh_session_record,
    revoke_refresh_session_record,
    rotate_refresh_session_record,
    save_oauth_client_registration,
    validate_access_session,
)

type JwtClaims = dict[str, object]

OAUTH_SCOPE = "mcp"


def _require_jwt_secret() -> str:
    secret = config_module.settings.jwt_secret.get_secret_value()
    if not secret:
        raise JwtError("JWT secret is not configured (set SIBYL_JWT_SECRET)")
    return secret


def _jwt_encode(payload: JwtClaims) -> str:
    secret = _require_jwt_secret()
    return jwt.encode(payload, secret, algorithm=config_module.settings.jwt_algorithm)


def _jwt_decode(token: str) -> JwtClaims:
    secret = _require_jwt_secret()
    return cast(
        "JwtClaims",
        jwt.decode(
            token,
            secret,
            algorithms=[config_module.settings.jwt_algorithm],
            options={"require": ["sub", "iat", "exp"]},
        ),
    )


def _parse_scopes_from_claims(claims: JwtClaims) -> list[str]:
    scopes = claims.get("scopes")
    if isinstance(scopes, list):
        parsed_scopes = [item for item in scopes if isinstance(item, str)]
        if len(parsed_scopes) == len(scopes):
            return parsed_scopes
    scope = claims.get("scope")
    if isinstance(scope, str) and scope.strip():
        return scope.split()
    return [OAUTH_SCOPE]


def _add_query_params(url: str, params: dict[str, str]) -> str:
    parts = urlsplit(url)
    return urlunsplit((parts.scheme, parts.netloc, parts.path, urlencode(params), parts.fragment))


def _create_refresh_token(
    *,
    user_id: UUID,
    organization_id: UUID | None,
    client_id: str,
    scopes: list[str],
    session_id: UUID | None = None,
    expires_in: timedelta = timedelta(days=30),
) -> tuple[str, datetime]:
    now = datetime.now(UTC)
    expires_at = now + expires_in
    payload: JwtClaims = {
        "sub": str(user_id),
        "typ": "refresh",
        "cid": client_id,
        "iat": int(now.timestamp()),
        "exp": int(expires_at.timestamp()),
        "scope": " ".join(scopes),
    }
    if organization_id is not None:
        payload["org"] = str(organization_id)
    if session_id is not None:
        payload["sid"] = str(session_id)
    return _jwt_encode(payload), expires_at


class SibylAuthorizationCode(AuthorizationCode):
    user_id: str
    organization_id: str | None = None


@dataclass(frozen=True)
class _PendingAuth:
    client_id: str
    expires_at: float
    params: AuthorizationParams


@dataclass(frozen=True)
class _AuthedUser:
    user_id: UUID
    expires_at: float


class _OAuthUser(Protocol):
    id: UUID


class _OAuthOrg(Protocol):
    id: UUID
    name: str
    is_personal: bool


class _RefreshSessionRecord(Protocol):
    id: UUID
    user_id: UUID
    organization_id: UUID | None


class SibylMcpOAuthProvider(
    OAuthAuthorizationServerProvider[SibylAuthorizationCode, RefreshToken, AccessToken]
):
    """OAuth provider for FastMCP auth routes."""

    def __init__(self) -> None:
        self._clients: dict[str, OAuthClientInformationFull] = {}
        self._pending: dict[str, _PendingAuth] = {}
        self._authed: dict[str, _AuthedUser] = {}
        self._codes: dict[str, SibylAuthorizationCode] = {}

    async def get_client(self, client_id: str) -> OAuthClientInformationFull | None:
        cached = self._clients.get(client_id)
        if cached is not None:
            return cached
        stored = await self._load_oauth_client_registration(client_id)
        if stored is None:
            return None
        try:
            client = OAuthClientInformationFull.model_validate(stored)
        except Exception:
            return None
        if not client.client_id:
            return None
        self._clients[str(client.client_id)] = client
        return client

    async def register_client(self, client_info: OAuthClientInformationFull) -> None:
        if not client_info.client_id:
            return
        client_id = str(client_info.client_id)
        self._clients[client_id] = client_info
        await self._save_oauth_client_registration(
            client_id=client_id,
            client_info=client_info.model_dump(mode="json", exclude_none=True),
        )

    async def authorize(
        self, client: OAuthClientInformationFull, params: AuthorizationParams
    ) -> str:
        request_id = secrets.token_urlsafe(24)
        self._pending[request_id] = _PendingAuth(
            client_id=str(client.client_id),
            expires_at=time.time() + 10 * 60,
            params=params,
        )
        issuer = str(config_module.settings.server_url).rstrip("/")
        return _add_query_params(f"{issuer}/_oauth/login", {"req": request_id})

    async def load_authorization_code(
        self, client: OAuthClientInformationFull, authorization_code: str
    ) -> SibylAuthorizationCode | None:
        code = self._codes.get(authorization_code)
        if code is None:
            return None
        if str(code.client_id) != str(client.client_id):
            return None
        return code

    async def exchange_authorization_code(
        self, client: OAuthClientInformationFull, authorization_code: SibylAuthorizationCode
    ) -> OAuthToken:
        self._codes.pop(authorization_code.code, None)

        user_id = UUID(authorization_code.user_id)
        org_id = (
            UUID(authorization_code.organization_id) if authorization_code.organization_id else None
        )
        scopes = authorization_code.scopes or [OAUTH_SCOPE]

        session_id = uuid4()
        access = create_access_token(
            user_id=user_id,
            organization_id=org_id,
            session_id=session_id,
            extra_claims={"scope": " ".join(scopes)},
        )
        refresh, refresh_expires_at = _create_refresh_token(
            user_id=user_id,
            organization_id=org_id,
            client_id=str(client.client_id),
            scopes=scopes,
            session_id=session_id,
        )

        access_expires_at = datetime.now(UTC) + timedelta(
            minutes=config_module.settings.access_token_expire_minutes
        )
        await self._create_session_record(
            user_id=user_id,
            token=access,
            expires_at=access_expires_at,
            session_id=session_id,
            organization_id=org_id,
            refresh_token=refresh,
            refresh_token_expires_at=refresh_expires_at,
            device_name="mcp_oauth",
            device_type="mcp",
        )
        return OAuthToken(
            access_token=access,
            refresh_token=refresh,
            expires_in=int(
                timedelta(
                    minutes=config_module.settings.access_token_expire_minutes
                ).total_seconds()
            ),
            scope=" ".join(scopes),
        )

    async def load_refresh_token(
        self, client: OAuthClientInformationFull, refresh_token: str
    ) -> RefreshToken | None:
        try:
            claims = _jwt_decode(refresh_token)
        except Exception:
            return None
        if claims.get("typ") != "refresh":
            return None
        if str(claims.get("cid", "")) != str(client.client_id):
            return None

        sub = claims.get("sub")
        if not isinstance(sub, str) or not sub:
            return None
        try:
            user_id = UUID(sub)
        except ValueError:
            return None

        org_raw = claims.get("org")
        org_id = None
        if org_raw:
            try:
                org_id = UUID(str(org_raw))
            except ValueError:
                return None

        existing = await self._load_refresh_session_record(refresh_token)
        if existing is None:
            return None
        if existing.user_id != user_id:
            return None
        if org_id != existing.organization_id:
            return None

        scopes = _parse_scopes_from_claims(claims)
        exp = claims.get("exp")
        expires_at = int(exp) if isinstance(exp, int) else None
        return RefreshToken(
            token=refresh_token,
            client_id=str(client.client_id),
            scopes=scopes,
            expires_at=expires_at,
        )

    async def exchange_refresh_token(
        self,
        client: OAuthClientInformationFull,
        refresh_token: RefreshToken,
        scopes: list[str],
    ) -> OAuthToken:
        claims = _jwt_decode(refresh_token.token)
        try:
            user_id = UUID(str(claims["sub"]))
            org_raw = claims.get("org")
            org_id = UUID(str(org_raw)) if org_raw else None
        except (KeyError, ValueError) as exc:
            raise TokenError(
                error="invalid_grant", error_description="invalid refresh token claims"
            ) from exc

        allowed_scopes = set(refresh_token.scopes or [])
        requested_scopes = set(scopes or [])
        if requested_scopes and not requested_scopes.issubset(allowed_scopes):
            scopes = list(allowed_scopes)

        existing = await self._load_refresh_session_record(refresh_token.token)
        if existing is None:
            raise TokenError(
                error="invalid_grant", error_description="refresh token does not exist"
            )

        access = create_access_token(
            user_id=user_id,
            organization_id=org_id,
            session_id=existing.id,
            extra_claims={"scope": " ".join(scopes)},
        )
        new_refresh, new_refresh_expires_at = _create_refresh_token(
            user_id=user_id,
            organization_id=org_id,
            client_id=str(client.client_id),
            scopes=scopes,
            session_id=existing.id,
        )

        access_expires_at = datetime.now(UTC) + timedelta(
            minutes=config_module.settings.access_token_expire_minutes
        )
        rotated = await self._rotate_refresh_session_record(
            refresh_token.token,
            new_access_token=access,
            new_access_expires_at=access_expires_at,
            new_refresh_token=new_refresh,
            new_refresh_expires_at=new_refresh_expires_at,
        )
        if rotated is None:
            raise TokenError(
                error="invalid_grant", error_description="refresh token does not exist"
            )
        return OAuthToken(
            access_token=access,
            refresh_token=new_refresh,
            expires_in=int(
                timedelta(
                    minutes=config_module.settings.access_token_expire_minutes
                ).total_seconds()
            ),
            scope=" ".join(scopes),
        )

    async def load_access_token(self, token: str) -> AccessToken | None:
        if token.startswith("sk_"):
            auth = await self._authenticate_api_key(token)
            if auth is None:
                return None
            scopes = list(auth.scopes or []) or [OAUTH_SCOPE]
            if scopes and OAUTH_SCOPE not in scopes:
                return None
            return AccessToken(token=token, client_id=f"api_key:{auth.api_key_id}", scopes=scopes)

        try:
            claims = verify_access_token(token)
        except JwtError:
            return None
        try:
            if not await validate_access_session(token):
                return None
        except TimeoutError:
            return None

        sub = claims.get("sub")
        if not isinstance(sub, str) or not sub:
            return None
        try:
            user_id = UUID(sub)
        except ValueError:
            return None

        exp = claims.get("exp")
        expires_at = exp if isinstance(exp, int) else None
        scopes = _parse_scopes_from_claims(claims)
        return AccessToken(
            token=token, client_id=f"user:{user_id}", scopes=scopes, expires_at=expires_at
        )

    async def revoke_token(self, token: AccessToken | RefreshToken) -> None:
        if isinstance(token, RefreshToken):
            await self._revoke_refresh_session_record(token.token)

    # ---------------------------------------------------------------------
    # UI helpers (custom routes)
    # ---------------------------------------------------------------------

    def _get_pending(self, request_id: str) -> _PendingAuth | None:
        pending = self._pending.get(request_id)
        if pending is None:
            return None
        if pending.expires_at < time.time():
            self._pending.pop(request_id, None)
            self._authed.pop(request_id, None)
            return None
        return pending

    def _get_authed_user(self, request_id: str) -> _AuthedUser | None:
        authed = self._authed.get(request_id)
        if authed is None:
            return None
        if authed.expires_at < time.time():
            self._authed.pop(request_id, None)
            return None
        return authed

    async def _create_session_record(self, **kwargs: object) -> object:
        return await create_session_record(**kwargs)

    async def _load_oauth_client_registration(self, client_id: str) -> dict[str, object] | None:
        return cast(
            "dict[str, object] | None",
            await load_oauth_client_registration(client_id),
        )

    async def _save_oauth_client_registration(
        self,
        *,
        client_id: str,
        client_info: dict[str, object],
    ) -> None:
        await save_oauth_client_registration(client_id=client_id, client_info=client_info)

    async def _load_refresh_session_record(
        self, refresh_token: str
    ) -> _RefreshSessionRecord | None:
        return cast(
            "_RefreshSessionRecord | None", await load_refresh_session_record(refresh_token)
        )

    async def _rotate_refresh_session_record(
        self, refresh_token: str, **kwargs: object
    ) -> _RefreshSessionRecord | None:
        return cast(
            "_RefreshSessionRecord | None",
            await rotate_refresh_session_record(refresh_token, **kwargs),
        )

    async def _revoke_refresh_session_record(self, refresh_token: str) -> None:
        await revoke_refresh_session_record(refresh_token)

    async def _authenticate_api_key(self, raw_key: str) -> ApiKeyAuth | None:
        return cast("ApiKeyAuth | None", await authenticate_api_key(raw_key))

    async def _authenticate_local_user(self, *, email: str, password: str) -> _OAuthUser | None:
        return cast(
            "_OAuthUser | None",
            await authenticate_local_user(email=email, password=password),
        )

    async def _list_user_orgs(self, *, user_id: UUID) -> list[_OAuthOrg]:
        return cast("list[_OAuthOrg]", await list_user_organizations(user_id=user_id))

    async def _ensure_personal_org(self, *, user_id: UUID) -> _OAuthOrg | None:
        return cast("_OAuthOrg | None", await ensure_personal_organization(user_id=user_id))

    async def _get_user(self, user_id: UUID) -> _OAuthUser | None:
        return cast("_OAuthUser | None", await get_user_by_id(user_id))

    async def ui_login_get(self, request: Request) -> Response:
        request_id = (request.query_params.get("req") or "").strip()
        pending = self._get_pending(request_id)
        if pending is None:
            return HTMLResponse(
                "<h1>OAuth Login</h1><p>Invalid or expired login request.</p>",
                status_code=400,
            )

        client = await self.get_client(pending.client_id)
        client_name = escape((client.client_name if client else None) or "MCP Client", quote=True)

        html = f"""
<!doctype html>
<html lang="en">
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1" />
  <title>Sibyl Login</title>
  <style>
    :root {{ color-scheme: dark; }}
    body {{ font-family: system-ui, -apple-system, Segoe UI, sans-serif; background: #0b0b10; color: #e8e8f0; margin: 0; }}
    .wrap {{ max-width: 520px; margin: 8vh auto; padding: 24px; background: #12121a; border: 1px solid #2a2a3a; border-radius: 14px; }}
    h1 {{ margin: 0 0 8px; font-size: 22px; }}
    .sub {{ color: #a7a7c7; margin-bottom: 18px; }}
    label {{ display: block; margin: 12px 0 6px; color: #cfcfe9; }}
    input {{ width: 100%; padding: 10px 12px; border-radius: 10px; border: 1px solid #2a2a3a; background: #0f0f16; color: #fff; }}
    button {{ margin-top: 16px; width: 100%; padding: 10px 12px; border-radius: 10px; border: 1px solid #3a2a6a; background: #5b2bff; color: #fff; font-weight: 600; cursor: pointer; }}
    .hint {{ margin-top: 12px; color: #a7a7c7; font-size: 13px; }}
    code {{ color: #80ffea; }}
  </style>
</head>
<body>
  <div class="wrap">
    <h1>Login to Sibyl</h1>
    <div class="sub">Authorize <strong>{client_name}</strong> to access your MCP tools.</div>
    <form method="post" action="/_oauth/login">
      <input type="hidden" name="req" value="{escape(request_id, quote=True)}" />
      <label>Email</label>
      <input name="email" type="email" autocomplete="username" required />
      <label>Password</label>
      <input name="password" type="password" autocomplete="current-password" required />
      <button type="submit">Continue</button>
    </form>
    <div class="hint">No local user yet? Create one at <code>/api/auth/local/signup</code>.</div>
  </div>
</body>
</html>
"""
        return HTMLResponse(html, status_code=200)

    async def ui_login_post(self, request: Request) -> Response:
        form = await request.form()
        request_id = str(form.get("req", "")).strip()
        email = str(form.get("email", "")).strip()
        password = str(form.get("password", "")).strip()

        pending = self._get_pending(request_id)
        if pending is None:
            return HTMLResponse(
                "<h1>OAuth Login</h1><p>Invalid or expired login request.</p>",
                status_code=400,
            )

        user = await self._authenticate_local_user(email=email, password=password)
        if user is None:
            return RedirectResponse(
                url=_add_query_params("/_oauth/login", {"req": request_id}), status_code=302
            )

        orgs = await self._list_user_orgs(user_id=user.id)

        if not orgs:
            org = await self._ensure_personal_org(user_id=user.id)
            if org is not None:
                orgs = [org]

        if len(orgs) == 1:
            org = orgs[0]
            code = secrets.token_urlsafe(32)
            auth_code = SibylAuthorizationCode(
                code=code,
                client_id=pending.client_id,
                expires_at=time.time() + 10 * 60,
                scopes=pending.params.scopes or [OAUTH_SCOPE],
                code_challenge=pending.params.code_challenge,
                redirect_uri=pending.params.redirect_uri,
                redirect_uri_provided_explicitly=pending.params.redirect_uri_provided_explicitly,
                resource=pending.params.resource,
                user_id=str(user.id),
                organization_id=str(org.id),
            )
            self._codes[code] = auth_code
            self._pending.pop(request_id, None)

            params: dict[str, str] = {"code": code}
            if pending.params.state:
                params["state"] = pending.params.state
            return RedirectResponse(
                url=_add_query_params(str(pending.params.redirect_uri), params), status_code=302
            )

        self._authed[request_id] = _AuthedUser(
            user_id=user.id,
            expires_at=time.time() + 5 * 60,
        )

        return RedirectResponse(
            url=_add_query_params("/_oauth/org", {"req": request_id}), status_code=302
        )

    async def ui_org_get(self, request: Request) -> Response:
        request_id = (request.query_params.get("req") or "").strip()
        pending = self._get_pending(request_id)
        if pending is None:
            return HTMLResponse(
                "<h1>OAuth Login</h1><p>Invalid or expired login request.</p>",
                status_code=400,
            )

        authed = self._get_authed_user(request_id)
        if authed is None:
            return RedirectResponse(
                url=_add_query_params("/_oauth/login", {"req": request_id}), status_code=302
            )

        orgs = await self._list_user_orgs(user_id=authed.user_id)

        if not orgs:
            return RedirectResponse(
                url=_add_query_params("/_oauth/login", {"req": request_id}), status_code=302
            )

        options = "\n".join(
            (
                '<label style="display:block;margin:10px 0;padding:10px;border:1px solid #2a2a3a;border-radius:10px;">'
                f'<input type="radio" name="org_id" value="{escape(str(org.id), quote=True)}" required style="margin-right:10px;" />'
                f"<strong>{escape(org.name)}</strong>"
                + (' <span style="color:#a7a7c7">(personal)</span>' if org.is_personal else "")
                + "</label>"
            )
            for org in orgs
        )

        html = f"""
<!doctype html>
<html lang="en">
<head>
  <meta charset="utf-8" />
  <meta name="viewport" content="width=device-width, initial-scale=1" />
  <title>Select Organization</title>
  <style>
    :root {{ color-scheme: dark; }}
    body {{ font-family: system-ui, -apple-system, Segoe UI, sans-serif; background: #0b0b10; color: #e8e8f0; margin: 0; }}
    .wrap {{ max-width: 640px; margin: 8vh auto; padding: 24px; background: #12121a; border: 1px solid #2a2a3a; border-radius: 14px; }}
    h1 {{ margin: 0 0 8px; font-size: 22px; }}
    .sub {{ color: #a7a7c7; margin-bottom: 18px; }}
    button {{ margin-top: 16px; width: 100%; padding: 10px 12px; border-radius: 10px; border: 1px solid #3a2a6a; background: #5b2bff; color: #fff; font-weight: 600; cursor: pointer; }}
    .secondary {{ margin-top: 10px; background: transparent; border-color: #2a2a3a; color: #e8e8f0; }}
  </style>
</head>
<body>
  <div class="wrap">
    <h1>Select an organization</h1>
    <div class="sub">Choose which org to use for this MCP session.</div>
    <form method="post" action="/_oauth/org">
      <input type="hidden" name="req" value="{escape(request_id, quote=True)}" />
      {options}
      <button type="submit">Continue</button>
    </form>
    <form method="post" action="/_oauth/org">
      <input type="hidden" name="req" value="{escape(request_id, quote=True)}" />
      <input type="hidden" name="create_personal" value="1" />
      <button class="secondary" type="submit">Create / use personal org</button>
    </form>
  </div>
</body>
</html>
"""
        return HTMLResponse(html, status_code=200)

    async def ui_org_post(self, request: Request) -> Response:
        form = await request.form()
        request_id = str(form.get("req", "")).strip()
        selected_org_id = str(form.get("org_id", "")).strip()
        create_personal = str(form.get("create_personal", "")).strip() == "1"

        pending = self._get_pending(request_id)
        if pending is None:
            return HTMLResponse(
                "<h1>OAuth Login</h1><p>Invalid or expired login request.</p>",
                status_code=400,
            )

        authed = self._get_authed_user(request_id)
        if authed is None:
            return RedirectResponse(
                url=_add_query_params("/_oauth/login", {"req": request_id}), status_code=302
            )

        if create_personal:
            org = await self._ensure_personal_org(user_id=authed.user_id)
            if org is None:
                return RedirectResponse(
                    url=_add_query_params("/_oauth/login", {"req": request_id}), status_code=302
                )
        else:
            orgs = await self._list_user_orgs(user_id=authed.user_id)
            org = next((o for o in orgs if str(o.id) == selected_org_id), None)
            if org is None:
                return HTMLResponse(
                    "<h1>OAuth Login</h1><p>Invalid organization selection.</p>",
                    status_code=400,
                )

        # Issue code outside the DB session; pending/auth mappings are in-memory.
        code = secrets.token_urlsafe(32)
        auth_code = SibylAuthorizationCode(
            code=code,
            client_id=pending.client_id,
            expires_at=time.time() + 10 * 60,
            scopes=pending.params.scopes or [OAUTH_SCOPE],
            code_challenge=pending.params.code_challenge,
            redirect_uri=pending.params.redirect_uri,
            redirect_uri_provided_explicitly=pending.params.redirect_uri_provided_explicitly,
            resource=pending.params.resource,
            user_id=str(authed.user_id),
            organization_id=str(org.id),
        )
        self._codes[code] = auth_code
        self._pending.pop(request_id, None)
        self._authed.pop(request_id, None)

        params: dict[str, str] = {"code": code}
        if pending.params.state:
            params["state"] = pending.params.state
        return RedirectResponse(
            url=_add_query_params(str(pending.params.redirect_uri), params), status_code=302
        )
