from __future__ import annotations

import asyncio
import secrets
import time
from typing import Any
from urllib.parse import urlencode

import httpx
from fastapi import Depends, FastAPI, HTTPException, Query, Request, Response
from fastapi.responses import JSONResponse, RedirectResponse
from starlette.middleware.sessions import SessionMiddleware

from app.config import Settings, get_settings
from app.crypto import decrypt_token, encrypt_token
from app.db import Database
from app.dcr import SUPPORTED_GRANTS, RateLimiter, client_ip, valid_redirect_uri
from app.pkce import verify_pkce_s256

app = FastAPI(title="Gitea MCP OAuth Bridge")

_db: Database | None = None


def db() -> Database:
    global _db
    if _db is None:
        _db = Database(get_settings().data_dir)
    return _db


def _install_session_middleware() -> None:
    settings = get_settings()
    app.add_middleware(SessionMiddleware, secret_key=settings.session_secret, same_site="lax")


_install_session_middleware()


def authorization_server_metadata(settings: Settings) -> dict[str, Any]:
    base = settings.issuer
    return {
        "issuer": base,
        "authorization_endpoint": f"{base}/authorize",
        "token_endpoint": f"{base}/token",
        "registration_endpoint": f"{base}/register",
        "response_types_supported": ["code"],
        "grant_types_supported": ["authorization_code"],
        "code_challenge_methods_supported": ["S256"],
        "token_endpoint_auth_methods_supported": ["none"],
        "scopes_supported": ["openid"],
    }


def protected_resource_metadata(settings: Settings, resource: str) -> dict[str, Any]:
    return {
        "resource": resource,
        "authorization_servers": [settings.issuer],
        "bearer_methods_supported": ["header"],
    }


def validate_mcp_client(settings: Settings, client_id: str, redirect_uri: str) -> bool:
    allowed = settings.allowed_clients.get(client_id)
    if allowed is not None:
        return redirect_uri in allowed
    if not settings.dcr_enabled:
        return False
    registered = db().get_client_redirect_uris(client_id)
    return registered is not None and redirect_uri in registered


def session_user_id(request: Request) -> int | None:
    raw = request.session.get("user_id")
    if raw is None:
        return None
    try:
        return int(raw)
    except (TypeError, ValueError):
        return None


def mcp_unauthorized_response(settings: Settings) -> JSONResponse:
    resource_meta = protected_resource_metadata(settings, settings.mcp_resource_url)
    www = f'Bearer realm="mcp", resource_metadata="{settings.issuer}/.well-known/oauth-protected-resource/mcp"'
    return JSONResponse(
        status_code=401,
        content={
            "error": "unauthorized",
            "message": "Valid OAuth access token required.",
            "resource_metadata": resource_meta,
        },
        headers={"WWW-Authenticate": www},
    )


@app.get("/health")
async def health() -> dict[str, bool]:
    return {"ok": True}


@app.get("/.well-known/oauth-authorization-server")
@app.get("/mcp/.well-known/oauth-authorization-server")
async def well_known_as(settings: Settings = Depends(get_settings)) -> dict[str, Any]:
    return authorization_server_metadata(settings)


@app.get("/.well-known/oauth-protected-resource")
async def well_known_pr(settings: Settings = Depends(get_settings)) -> dict[str, Any]:
    return protected_resource_metadata(settings, settings.issuer)


@app.get("/.well-known/oauth-protected-resource/mcp")
async def well_known_pr_mcp(settings: Settings = Depends(get_settings)) -> dict[str, Any]:
    return protected_resource_metadata(settings, settings.mcp_resource_url)


_dcr_limiter = RateLimiter(max_events=30, window_seconds=600)


def _oauth_error(error: str, description: str, status: int = 400) -> JSONResponse:
    return JSONResponse(status_code=status, content={"error": error, "error_description": description})


def _client_metadata(settings: Settings, row: Any) -> dict[str, Any]:
    import json

    return {
        "client_id": row["client_id"],
        "client_id_issued_at": int(row["created_at"]),
        "client_name": row["client_name"] or "",
        "redirect_uris": json.loads(row["redirect_uris"]),
        "token_endpoint_auth_method": "none",
        "grant_types": ["authorization_code"],
        "response_types": ["code"],
        "registration_client_uri": f"{settings.issuer}/register/{row['client_id']}",
    }


@app.post("/register")
async def register_client(request: Request, settings: Settings = Depends(get_settings)) -> JSONResponse:
    """RFC 7591 dynamic client registration for public PKCE clients."""
    if not settings.dcr_enabled:
        return _oauth_error(
            "registration_not_supported",
            "Register MCP clients via ALLOWED_MCP_CLIENTS; dynamic registration is disabled.",
        )
    if _dcr_limiter.hit(client_ip(request)):
        return _oauth_error("too_many_requests", "Rate limit exceeded for client registration", 429)
    try:
        body = await request.json()
    except Exception:
        body = None
    if not isinstance(body, dict):
        return _oauth_error("invalid_client_metadata", "Request body must be a JSON object")

    uris = body.get("redirect_uris")
    if not isinstance(uris, list) or not uris or len(uris) > 20:
        return _oauth_error("invalid_redirect_uri", "redirect_uris must be a non-empty list")
    bad = [u for u in uris if not valid_redirect_uri(u)]
    if bad:
        return _oauth_error(
            "invalid_redirect_uri",
            "Allowed redirect URIs: https://..., http://localhost or 127.0.0.1, or an app scheme "
            f"such as cursor://... Rejected: {bad[:3]}",
        )
    grants = body.get("grant_types") or ["authorization_code"]
    if not isinstance(grants, list) or "authorization_code" not in grants or not set(grants) <= SUPPORTED_GRANTS:
        return _oauth_error("invalid_client_metadata", "grant_types must include authorization_code")
    rtypes = body.get("response_types") or ["code"]
    if rtypes != ["code"]:
        return _oauth_error("invalid_client_metadata", 'Only response_types ["code"] is supported')
    # Public clients only: any requested auth method is answered with "none"
    # (RFC 7591 lets the server replace requested values); PKCE is mandatory.
    name = str(body.get("client_name") or "")[:200]

    created = db().register_client(name, uris, settings.dcr_max_clients)
    if created is None:
        return _oauth_error("invalid_client_metadata", "Client registration limit reached")
    client_id, reg_token, issued_at = created
    content = _client_metadata(settings, db().get_client(client_id))
    content["registration_access_token"] = reg_token
    return JSONResponse(status_code=201, content=content, headers={"Cache-Control": "no-store"})


def _registration_auth(request: Request, client_id: str) -> bool:
    auth = request.headers.get("authorization", "")
    token = auth[7:].strip() if auth.lower().startswith("bearer ") else ""
    return db().check_registration_token(client_id, token)


@app.get("/register/{client_id}")
async def read_client(client_id: str, request: Request, settings: Settings = Depends(get_settings)) -> JSONResponse:
    if not _registration_auth(request, client_id):
        return _oauth_error("invalid_token", "Invalid registration access token", 401)
    return JSONResponse(content=_client_metadata(settings, db().get_client(client_id)))


@app.delete("/register/{client_id}")
async def delete_client(client_id: str, request: Request) -> Response:
    if not _registration_auth(request, client_id):
        return _oauth_error("invalid_token", "Invalid registration access token", 401)
    db().delete_client(client_id)
    return Response(status_code=204)


@app.get("/authorize")
async def authorize(
    request: Request,
    response_type: str = Query(...),
    client_id: str = Query(...),
    redirect_uri: str = Query(...),
    code_challenge: str = Query(...),
    code_challenge_method: str = Query("S256"),
    state: str | None = Query(None),
    settings: Settings = Depends(get_settings),
) -> Response:
    if response_type != "code":
        raise HTTPException(status_code=400, detail="unsupported response_type")
    if code_challenge_method != "S256":
        raise HTTPException(status_code=400, detail="PKCE S256 required")
    if not validate_mcp_client(settings, client_id, redirect_uri):
        raise HTTPException(status_code=400, detail="invalid client_id or redirect_uri")

    user_id = session_user_id(request)
    if user_id is None:
        query = request.url.query
        next_url = f"/authorize?{query}" if query else "/authorize"
        return RedirectResponse(url=f"/oidc/login?{urlencode({'next': next_url})}", status_code=302)

    user = db().get_user_by_id(user_id)
    if user is None:
        raise HTTPException(status_code=401, detail="Login required")
    tokens = db().get_upstream_tokens(user.id)
    if tokens is None or not tokens["upstream_access"]:
        query = request.url.query
        next_url = f"/authorize?{query}" if query else "/authorize"
        return RedirectResponse(url=f"/oidc/login?{urlencode({'next': next_url})}", status_code=302)

    code = db().create_oauth_code(
        client_id=client_id,
        redirect_uri=redirect_uri,
        code_challenge=code_challenge,
        code_challenge_method=code_challenge_method,
        user_id=user.id,
    )
    params = {"code": code}
    if state:
        params["state"] = state
    location = f"{redirect_uri}?{urlencode(params)}"
    return RedirectResponse(url=location, status_code=302)


@app.post("/token")
async def token(
    request: Request,
    settings: Settings = Depends(get_settings),
) -> JSONResponse:
    form = await request.form()
    grant_type = form.get("grant_type")
    if grant_type != "authorization_code":
        raise HTTPException(status_code=400, detail="unsupported grant_type")

    code = form.get("code")
    redirect_uri = form.get("redirect_uri")
    client_id = form.get("client_id")
    code_verifier = form.get("code_verifier")
    if not all(isinstance(x, str) for x in (code, redirect_uri, client_id, code_verifier)):
        raise HTTPException(status_code=400, detail="invalid request")

    assert isinstance(code, str)
    assert isinstance(redirect_uri, str)
    assert isinstance(client_id, str)
    assert isinstance(code_verifier, str)

    if not validate_mcp_client(settings, client_id, redirect_uri):
        raise HTTPException(status_code=400, detail="invalid client_id or redirect_uri")

    row = db().consume_oauth_code(code)
    if row is None:
        raise HTTPException(status_code=400, detail="invalid or expired code")
    if row["client_id"] != client_id or row["redirect_uri"] != redirect_uri:
        raise HTTPException(status_code=400, detail="client or redirect mismatch")
    if not verify_pkce_s256(code_verifier, row["code_challenge"]):
        raise HTTPException(status_code=400, detail="invalid code_verifier")

    access_token = db().create_access_token(client_id=client_id, user_id=int(row["user_id"]))
    return JSONResponse(
        content={
            "access_token": access_token,
            "token_type": "Bearer",
            "expires_in": 604800,
        }
    )


def _oidc_discovery(settings: Settings) -> dict[str, Any]:
    issuer = settings.oidc_issuer.rstrip("/")
    with httpx.Client(timeout=15.0) as client:
        resp = client.get(f"{issuer}/.well-known/openid-configuration")
        resp.raise_for_status()
        return resp.json()


@app.get("/oidc/login")
async def oidc_login(
    request: Request,
    next: str = Query("/"),
    settings: Settings = Depends(get_settings),
) -> RedirectResponse:
    state = secrets.token_urlsafe(24)
    db().save_oidc_state(state, next)
    discovery = _oidc_discovery(settings)
    auth_endpoint = discovery["authorization_endpoint"]
    redirect_uri = f"{settings.issuer}/oidc/callback"
    params = {
        "client_id": settings.oidc_client_id,
        "response_type": "code",
        "redirect_uri": redirect_uri,
        "scope": settings.oidc_scopes,
        "state": state,
    }
    return RedirectResponse(url=f"{auth_endpoint}?{urlencode(params)}", status_code=302)


@app.get("/oidc/callback")
async def oidc_callback(
    request: Request,
    code: str | None = Query(None),
    state: str | None = Query(None),
    settings: Settings = Depends(get_settings),
) -> RedirectResponse:
    if not code or not state:
        raise HTTPException(status_code=400, detail="missing code or state")
    next_url = db().pop_oidc_state(state)
    if next_url is None:
        raise HTTPException(status_code=400, detail="invalid or expired state")

    discovery = _oidc_discovery(settings)
    token_endpoint = discovery["token_endpoint"]
    redirect_uri = f"{settings.issuer}/oidc/callback"
    data = {
        "grant_type": "authorization_code",
        "code": code,
        "redirect_uri": redirect_uri,
        "client_id": settings.oidc_client_id,
        "client_secret": settings.oidc_client_secret,
    }
    with httpx.Client(timeout=15.0) as client:
        token_resp = client.post(token_endpoint, data=data)
        token_resp.raise_for_status()
        tokens = token_resp.json()

    id_token = tokens.get("id_token")
    access = tokens.get("access_token")
    sub = None
    email = None
    if id_token:
        import base64
        import json

        payload_b64 = id_token.split(".")[1]
        padding = "=" * (-len(payload_b64) % 4)
        payload = json.loads(base64.urlsafe_b64decode(payload_b64 + padding))
        sub = payload.get("sub")
        email = payload.get("email")
    if not sub and access:
        userinfo = discovery.get("userinfo_endpoint")
        if userinfo:
            with httpx.Client(timeout=15.0) as client:
                ui = client.get(userinfo, headers={"Authorization": f"Bearer {access}"})
                if ui.status_code == 200:
                    body = ui.json()
                    sub = body.get("sub")
                    email = body.get("email")

    if not sub:
        raise HTTPException(status_code=400, detail="could not resolve OIDC subject")

    user = db().upsert_user(oidc_sub=str(sub), email=email)
    request.session["user_id"] = user.id

    if not access:
        raise HTTPException(status_code=400, detail="upstream returned no access token")
    _store_upstream_tokens(settings, user.id, tokens)
    return RedirectResponse(url=next_url, status_code=302)


def _store_upstream_tokens(settings: Settings, user_id: int, tokens: dict[str, Any]) -> None:
    key = settings.token_encryption_key
    refresh = tokens.get("refresh_token")
    expires_in = float(tokens.get("expires_in") or 3600)
    db().set_upstream_tokens(
        user_id,
        encrypt_token(key, tokens["access_token"]),
        encrypt_token(key, refresh) if refresh else None,
        time.time() + expires_in,
    )


_refresh_locks: dict[int, asyncio.Lock] = {}


async def _upstream_access_token(settings: Settings, user_id: int) -> str | None:
    """Current upstream access token for the user, refreshing it when it is about to expire."""
    key = settings.token_encryption_key

    def fresh() -> str | None:
        row = db().get_upstream_tokens(user_id)
        if row is None or not row["upstream_access"]:
            return None
        if (row["upstream_expires_at"] or 0) > time.time() + 60:
            return decrypt_token(key, row["upstream_access"])
        return None

    token = fresh()
    if token:
        return token
    lock = _refresh_locks.setdefault(user_id, asyncio.Lock())
    async with lock:
        token = fresh()
        if token:
            return token
        row = db().get_upstream_tokens(user_id)
        refresh = decrypt_token(key, row["upstream_refresh"]) if row is not None and row["upstream_refresh"] else None
        if not refresh:
            return None
        discovery = _oidc_discovery(settings)
        async with httpx.AsyncClient(timeout=15.0) as client:
            resp = await client.post(
                discovery["token_endpoint"],
                data={
                    "grant_type": "refresh_token",
                    "refresh_token": refresh,
                    "client_id": settings.oidc_client_id,
                    "client_secret": settings.oidc_client_secret,
                },
            )
        if resp.status_code != 200:
            return None
        tokens = resp.json()
        if not tokens.get("access_token"):
            return None
        _store_upstream_tokens(settings, user_id, tokens)
        return tokens["access_token"]


_FORWARD_REQUEST_HEADERS = ("mcp-session-id", "mcp-protocol-version", "last-event-id")
_FORWARD_RESPONSE_HEADERS = ("mcp-session-id",)


async def _proxy_mcp(request: Request, settings: Settings) -> Response:
    auth = request.headers.get("authorization", "")
    if not auth.lower().startswith("bearer "):
        return mcp_unauthorized_response(settings)
    token = auth[7:].strip()
    row = db().get_access_token(token)

    if row is None:
        if not (settings.passthrough_unknown_tokens and token):
            return mcp_unauthorized_response(settings)
        upstream_auth = f"Bearer {token}"
    else:
        upstream_token = await _upstream_access_token(settings, int(row["user_id"]))
        if not upstream_token:
            return mcp_unauthorized_response(settings)
        upstream_auth = f"Bearer {upstream_token}"

    body = await request.body()
    headers = {
        "Authorization": upstream_auth,
        "Content-Type": request.headers.get("content-type", "application/json"),
        "Accept": request.headers.get("accept", "application/json"),
    }
    for name in _FORWARD_REQUEST_HEADERS:
        value = request.headers.get(name)
        if value:
            headers[name] = value
    if settings.upstream_host_header:
        headers["Host"] = settings.upstream_host_header
    async with httpx.AsyncClient(timeout=120.0) as client:
        upstream = await client.request(
            method=request.method,
            url=settings.upstream_mcp_url,
            content=body,
            headers=headers,
        )
    if upstream.status_code == 401:
        return mcp_unauthorized_response(settings)
    out_headers = {"Content-Type": upstream.headers.get("content-type", "application/json")}
    for name in _FORWARD_RESPONSE_HEADERS:
        value = upstream.headers.get(name)
        if isinstance(value, str) and value:
            out_headers[name] = value
    return Response(content=upstream.content, status_code=upstream.status_code, headers=out_headers)


@app.api_route("/mcp", methods=["GET", "POST", "DELETE"])
async def mcp_endpoint(request: Request, settings: Settings = Depends(get_settings)) -> Response:
    if request.method == "GET":
        return Response(status_code=405, headers={"Allow": "POST, DELETE"})
    return await _proxy_mcp(request, settings)