"""OAuth 2.1 with PKCE for the MCP connector.

Designed for internal single-user use. Implements:
- Dynamic Client Registration (RFC 7591)
- Authorization Code flow with PKCE (S256)
- Bearer token issuance
- A setup-secret gate so only the intended user can complete the flow
"""

import os
import json
import time
import sqlite3
import secrets
import hashlib
import base64
import logging
from urllib.parse import urlencode
from flask import Blueprint, request, jsonify, redirect, abort

logger = logging.getLogger('experiment-analyzer-mcp.oauth')

DB_PATH = os.path.join(os.path.dirname(__file__), "oauth.db")
USERS_PATH = os.path.join(os.path.dirname(__file__), "users.json")
ISSUER = "https://mcp.analizaweb.pl"


def _lookup_user_by_secret(setup_secret: str):
    """Return (user_id, user_info) if the secret matches a configured user, else (None, None)."""
    if not setup_secret:
        return None, None
    try:
        with open(USERS_PATH) as f:
            users = json.load(f)
    except FileNotFoundError:
        logger.error(f"users.json not found at {USERS_PATH}")
        return None, None
    except json.JSONDecodeError as e:
        logger.error(f"users.json is invalid JSON: {e}")
        return None, None

    for user_id, info in users.items():
        if info.get("setup_secret") == setup_secret:
            return user_id, info
    return None, None

# Token lifetime: 30 days. Refresh tokens not supported in this minimal impl.
TOKEN_LIFETIME_SECONDS = 30 * 24 * 3600
AUTH_CODE_LIFETIME_SECONDS = 600  # 10 minutes

oauth_bp = Blueprint("oauth", __name__)


def _db():
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    return conn


def _b64url(data: bytes) -> str:
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode()


def _verify_pkce(code_verifier: str, stored_challenge: str, method: str) -> bool:
    """Verify the PKCE code_verifier matches the stored code_challenge."""
    if method != "S256":
        return False
    digest = hashlib.sha256(code_verifier.encode()).digest()
    return _b64url(digest) == stored_challenge


# ---------- OAuth discovery endpoints ----------

@oauth_bp.route("/.well-known/oauth-authorization-server")
def oauth_metadata():
    """RFC 8414 discovery document."""
    return jsonify({
        "issuer": ISSUER,
        "authorization_endpoint": f"{ISSUER}/oauth/authorize",
        "token_endpoint": f"{ISSUER}/oauth/token",
        "registration_endpoint": f"{ISSUER}/oauth/register",
        "response_types_supported": ["code"],
        "grant_types_supported": ["authorization_code"],
        "code_challenge_methods_supported": ["S256"],
        "token_endpoint_auth_methods_supported": ["none"],
        "scopes_supported": ["mcp"],
    })


@oauth_bp.route("/.well-known/oauth-protected-resource")
def protected_resource_metadata():
    """RFC 9728 protected resource metadata (for /mcp)."""
    return jsonify({
        "resource": f"{ISSUER}/mcp",
        "authorization_servers": [ISSUER],
        "bearer_methods_supported": ["header"],
    })


# ---------- Dynamic Client Registration ----------

@oauth_bp.route("/oauth/register", methods=["POST"])
def register_client():
    """RFC 7591 dynamic client registration. Claude.ai calls this."""
    try:
        body = request.get_json(force=True) or {}
    except Exception:
        return jsonify({"error": "invalid_client_metadata"}), 400

    redirect_uris = body.get("redirect_uris") or []
    if not redirect_uris:
        return jsonify({"error": "invalid_redirect_uri"}), 400

    client_id = secrets.token_urlsafe(16)
    client_name = body.get("client_name", "unknown")

    with _db() as conn:
        conn.execute(
            "INSERT INTO oauth_clients (client_id, client_secret, redirect_uris, "
            "client_name, registered_at) VALUES (?, ?, ?, ?, ?)",
            (client_id, None, json.dumps(redirect_uris), client_name, int(time.time())),
        )

    logger.info(f"Registered OAuth client: {client_id} name={client_name}")

    return jsonify({
        "client_id": client_id,
        "client_id_issued_at": int(time.time()),
        "redirect_uris": redirect_uris,
        "client_name": client_name,
        "token_endpoint_auth_method": "none",
        "grant_types": ["authorization_code"],
        "response_types": ["code"],
    }), 201


# ---------- Authorization endpoint ----------

@oauth_bp.route("/oauth/authorize")
def authorize():
    """
    Auto-approves the request and redirects back with a code.
    Gated by ?setup_secret=... so only the intended user can complete the flow.
    """
    client_id = request.args.get("client_id")
    redirect_uri = request.args.get("redirect_uri")
    response_type = request.args.get("response_type")
    code_challenge = request.args.get("code_challenge")
    code_challenge_method = request.args.get("code_challenge_method", "S256")
    state = request.args.get("state", "")
    setup_secret = request.args.get("setup_secret", "")

    # Check setup secret first — this is the single-user gate
    if not SETUP_SECRET or setup_secret != SETUP_SECRET:
        logger.warning(f"authorize: missing or invalid setup_secret from {request.remote_addr}")
        return (
            "<h1>Authorization required</h1>"
            "<p>This MCP server is single-user. To complete the connection, "
            "append <code>&setup_secret=YOUR_SECRET</code> to the URL.</p>"
            "<p>If you are not the intended user, you cannot proceed.</p>",
            403,
        )

    # Validate basics
    if response_type != "code":
        return jsonify({"error": "unsupported_response_type"}), 400
    if not code_challenge or code_challenge_method != "S256":
        return jsonify({"error": "invalid_request",
                        "error_description": "PKCE S256 required"}), 400
    if not client_id or not redirect_uri:
        return jsonify({"error": "invalid_request"}), 400

    # Validate client + redirect_uri
    with _db() as conn:
        row = conn.execute(
            "SELECT redirect_uris FROM oauth_clients WHERE client_id = ?",
            (client_id,),
        ).fetchone()
    if not row:
        return jsonify({"error": "invalid_client"}), 400
    allowed_redirects = json.loads(row["redirect_uris"])
    if redirect_uri not in allowed_redirects:
        return jsonify({"error": "invalid_redirect_uri"}), 400

    # Issue authorization code
    code = secrets.token_urlsafe(32)
    expires_at = int(time.time()) + AUTH_CODE_LIFETIME_SECONDS

    with _db() as conn:
        conn.execute(
            "INSERT INTO authorization_codes (code, client_id, redirect_uri, "
            "code_challenge, code_challenge_method, expires_at) "
            "VALUES (?, ?, ?, ?, ?, ?)",
            (code, client_id, redirect_uri, code_challenge,
             code_challenge_method, expires_at),
        )

    logger.info(f"Issued auth code for client_id={client_id}")

    # Redirect back to client
    params = {"code": code}
    if state:
        params["state"] = state
    return redirect(f"{redirect_uri}?{urlencode(params)}")


# ---------- Token endpoint ----------

@oauth_bp.route("/oauth/token", methods=["POST"])
def token():
    """Exchange authorization code for access token."""
    grant_type = request.form.get("grant_type")
    code = request.form.get("code")
    redirect_uri = request.form.get("redirect_uri")
    client_id = request.form.get("client_id")
    code_verifier = request.form.get("code_verifier")

    if grant_type != "authorization_code":
        return jsonify({"error": "unsupported_grant_type"}), 400
    if not all([code, redirect_uri, client_id, code_verifier]):
        return jsonify({"error": "invalid_request"}), 400

    # Look up the code
    with _db() as conn:
        row = conn.execute(
            "SELECT client_id, redirect_uri, code_challenge, code_challenge_method, "
            "expires_at, used FROM authorization_codes WHERE code = ?",
            (code,),
        ).fetchone()

        if not row:
            return jsonify({"error": "invalid_grant"}), 400
        if row["used"]:
            return jsonify({"error": "invalid_grant",
                            "error_description": "code reused"}), 400
        if row["expires_at"] < time.time():
            return jsonify({"error": "invalid_grant",
                            "error_description": "code expired"}), 400
        if row["client_id"] != client_id:
            return jsonify({"error": "invalid_grant"}), 400
        if row["redirect_uri"] != redirect_uri:
            return jsonify({"error": "invalid_grant"}), 400

        # Verify PKCE
        if not _verify_pkce(code_verifier, row["code_challenge"],
                             row["code_challenge_method"]):
            return jsonify({"error": "invalid_grant",
                            "error_description": "PKCE verification failed"}), 400

        # Mark code as used
        conn.execute("UPDATE authorization_codes SET used = 1 WHERE code = ?", (code,))

        # Issue token
        access_token = secrets.token_urlsafe(32)
        expires_at = int(time.time()) + TOKEN_LIFETIME_SECONDS
        conn.execute(
            "INSERT INTO access_tokens (token, client_id, expires_at) VALUES (?, ?, ?)",
            (access_token, client_id, expires_at),
        )

    logger.info(f"Issued access token for client_id={client_id}")

    return jsonify({
        "access_token": access_token,
        "token_type": "Bearer",
        "expires_in": TOKEN_LIFETIME_SECONDS,
        "scope": "mcp",
    })


# ---------- Token validation (used by /mcp) ----------

def validate_token(token_str: str) -> bool:
    if not token_str:
        return False
    with _db() as conn:
        row = conn.execute(
            "SELECT expires_at, revoked FROM access_tokens WHERE token = ?",
            (token_str,),
        ).fetchone()
    if not row:
        return False
    if row["revoked"]:
        return False
    if row["expires_at"] < time.time():
        return False
    return True