<?php

namespace AskTamMCP\MCP;

defined('ABSPATH') || exit;

/**
 * OAuth 2.0 server with:
 *   - PKCE (S256) verification on token exchange
 *   - redirect_uri allowlist per key
 *   - Short-lived JWT access tokens (not raw API keys)
 *   - authorization_code + client_credentials grants
 */
class OAuthServer {

    // Allowed redirect_uri origins — Claude.ai callback
    private const ALLOWED_REDIRECT_HOSTS = [
        'claude.ai',
        'localhost',
        '127.0.0.1',
    ];

    public function registerRoutes(): void {
        register_rest_route('asktam-mcp/v1', '/oauth/authorize', [
            'methods'             => ['GET', 'POST'],
            'callback'            => [$this, 'handleAuthorize'],
            'permission_callback' => '__return_true',
        ]);

        register_rest_route('asktam-mcp/v1', '/oauth/token', [
            'methods'             => 'POST',
            'callback'            => [$this, 'handleToken'],
            'permission_callback' => '__return_true',
        ]);

        register_rest_route('asktam-mcp/v1', '/oauth/metadata', [
            'methods'             => 'GET',
            'callback'            => [$this, 'handleMetadata'],
            'permission_callback' => '__return_true',
        ]);

        add_action('rest_api_init', function () {
            register_rest_route('', '/.well-known/oauth-authorization-server', [
                'methods'             => 'GET',
                'callback'            => [$this, 'handleMetadata'],
                'permission_callback' => '__return_true',
            ]);
        }, 20);
    }

    // -------------------------------------------------------------------------
    // Authorization endpoint
    // -------------------------------------------------------------------------
    public function handleAuthorize(\WP_REST_Request $request): void {
        $clientId      = sanitize_text_field($request->get_param('client_id')             ?? '');
        $redirectUri   = esc_url_raw($request->get_param('redirect_uri')                  ?? '');
        $state         = sanitize_text_field($request->get_param('state')                 ?? '');
        $codeChallenge = sanitize_text_field($request->get_param('code_challenge')        ?? '');
        $challengeMethod = sanitize_text_field($request->get_param('code_challenge_method') ?? 'S256');
        $approved      = sanitize_text_field($request->get_param('_mcp_approve')          ?? '');

        // 1. Validate redirect_uri host
        if (!$this->isRedirectUriAllowed($redirectUri)) {
            wp_die('Invalid redirect_uri: host not in allowlist.', 'MCP Security Error', ['response' => 400]);
        }

        // 2. Validate client_id
        global $wpdb;
        $table = $wpdb->prefix . 'asktam_mcp_keys';
        $row   = $wpdb->get_row(
            $wpdb->prepare("SELECT * FROM {$table} WHERE api_key = %s AND is_active = 1", $clientId),
            ARRAY_A
        );

        if (!$row) {
            $this->redirectError($redirectUri, 'invalid_client', $state);
            return;
        }

        // 3. PKCE: require code_challenge for authorization_code flow
        if (empty($codeChallenge)) {
            $this->redirectError($redirectUri, 'invalid_request', $state, 'PKCE code_challenge is required');
            return;
        }

        // 4. POST: user approved
        if ($_SERVER['REQUEST_METHOD'] === 'POST' && $approved === '1') {
            $code = 'mcp_code_' . bin2hex(random_bytes(16));
            set_transient('asktam_mcp_code_' . $code, [
                'client_id'        => $clientId,
                'redirect_uri'     => $redirectUri,
                'scopes'           => $row['scopes'],
                'key_id'           => $row['id'],
                'code_challenge'   => $codeChallenge,
                'challenge_method' => $challengeMethod,
            ], 300); // 5 min

            wp_redirect(add_query_arg(['code' => $code, 'state' => $state], $redirectUri));
            exit;
        }

        // 5. GET: render approval page
        $this->renderApprovalPage($row, $clientId, $redirectUri, $state, $codeChallenge, $challengeMethod);
    }

    // -------------------------------------------------------------------------
    // Token endpoint
    // -------------------------------------------------------------------------
    public function handleToken(\WP_REST_Request $request): \WP_REST_Response {
        $grantType = $request->get_param('grant_type') ?? '';

        return match ($grantType) {
            'authorization_code'  => $this->handleAuthCodeExchange($request),
            'client_credentials'  => $this->handleClientCredentials($request),
            default               => $this->oauthError('unsupported_grant_type', 'Supported: authorization_code, client_credentials', 400),
        };
    }

    private function handleAuthCodeExchange(\WP_REST_Request $request): \WP_REST_Response {
        $code         = $request->get_param('code')          ?? '';
        $codeVerifier = $request->get_param('code_verifier') ?? '';
        $redirectUri  = $request->get_param('redirect_uri')  ?? '';

        if (empty($code)) {
            return $this->oauthError('invalid_request', 'Missing code', 400);
        }

        $codeData = get_transient('asktam_mcp_code_' . $code);
        if (!$codeData) {
            return $this->oauthError('invalid_grant', 'Authorization code expired or invalid', 400);
        }

        // Always delete — one-time use
        delete_transient('asktam_mcp_code_' . $code);

        // PKCE verification
        if (!empty($codeData['code_challenge'])) {
            if (empty($codeVerifier)) {
                return $this->oauthError('invalid_grant', 'PKCE code_verifier is required', 400);
            }

            $method    = $codeData['challenge_method'] ?? 'S256';
            $challenge = $codeData['code_challenge'];

            $valid = match ($method) {
                'S256'  => hash_equals($challenge, rtrim(strtr(base64_encode(hash('sha256', $codeVerifier, true)), '+/', '-_'), '=')),
                'plain' => hash_equals($challenge, $codeVerifier),
                default => false,
            };

            if (!$valid) {
                return $this->oauthError('invalid_grant', 'PKCE verification failed', 400);
            }
        }

        // redirect_uri must match what was used in /authorize
        if (!empty($codeData['redirect_uri']) && $redirectUri !== $codeData['redirect_uri']) {
            return $this->oauthError('invalid_grant', 'redirect_uri mismatch', 400);
        }

        // Fetch key row and issue JWT
        global $wpdb;
        $row = $wpdb->get_row(
            $wpdb->prepare("SELECT * FROM {$wpdb->prefix}asktam_mcp_keys WHERE id = %d AND is_active = 1", $codeData['key_id']),
            ARRAY_A
        );

        if (!$row) {
            return $this->oauthError('invalid_grant', 'Key no longer active', 400);
        }

        $accessToken = \AskTamMCP\Core\TokenManager::issue($row);

        return new \WP_REST_Response([
            'access_token' => $accessToken,
            'token_type'   => 'bearer',
            'expires_in'   => 3600,
            'scope'        => $row['scopes'],
        ], 200);
    }

    private function handleClientCredentials(\WP_REST_Request $request): \WP_REST_Response {
        [$clientId, $clientSecret] = $this->extractCredentials($request);

        if (empty($clientId) || empty($clientSecret)) {
            return $this->oauthError('invalid_client', 'Missing client_id or client_secret', 401);
        }

        global $wpdb;
        $row = $wpdb->get_row(
            $wpdb->prepare("SELECT * FROM {$wpdb->prefix}asktam_mcp_keys WHERE api_key = %s AND is_active = 1", $clientId),
            ARRAY_A
        );

        if (!$row || !hash_equals($clientId, $clientSecret)) {
            return $this->oauthError('invalid_client', 'Invalid credentials', 401);
        }

        $accessToken = \AskTamMCP\Core\TokenManager::issue($row);

        return new \WP_REST_Response([
            'access_token' => $accessToken,
            'token_type'   => 'bearer',
            'expires_in'   => 3600,
            'scope'        => $row['scopes'],
        ], 200);
    }

    // -------------------------------------------------------------------------
    // Metadata
    // -------------------------------------------------------------------------
    public function handleMetadata(): \WP_REST_Response {
        $base = rest_url('asktam-mcp/v1');
        return new \WP_REST_Response([
            'issuer'                                => get_site_url(),
            'authorization_endpoint'                => $base . '/oauth/authorize',
            'token_endpoint'                        => $base . '/oauth/token',
            'token_endpoint_auth_methods_supported' => ['client_secret_post', 'client_secret_basic', 'none'],
            'grant_types_supported'                 => ['authorization_code', 'client_credentials'],
            'response_types_supported'              => ['code'],
            'code_challenge_methods_supported'      => ['S256', 'plain'],
            'scopes_supported'                      => ['read', 'write', 'admin', 'woocommerce'],
        ]);
    }

    // -------------------------------------------------------------------------
    // Security helpers
    // -------------------------------------------------------------------------

    /**
     * Validate redirect_uri against an allowlist of trusted hosts.
     * Only claude.ai and localhost are allowed by default.
     * Filterable for custom setups.
     */
    private function isRedirectUriAllowed(string $uri): bool {
        if (empty($uri)) return false;

        $host = parse_url($uri, PHP_URL_HOST);
        if (!$host) return false;

        $allowed = apply_filters('asktam_mcp_allowed_redirect_hosts', self::ALLOWED_REDIRECT_HOSTS);

        foreach ($allowed as $allowedHost) {
            // Exact match or subdomain match
            if ($host === $allowedHost || str_ends_with($host, '.' . $allowedHost)) {
                return true;
            }
        }

        return false;
    }

    private function redirectError(string $redirectUri, string $error, string $state, string $description = ''): void {
        if ($redirectUri && $this->isRedirectUriAllowed($redirectUri)) {
            $params = ['error' => $error, 'state' => $state];
            if ($description) $params['error_description'] = $description;
            wp_redirect(add_query_arg($params, $redirectUri));
            exit;
        }
        status_header(400);
        header('Content-Type: text/plain');
        echo "OAuth error: {$error}";
        exit;
    }

    private function extractCredentials(\WP_REST_Request $request): array {
        $clientId     = $request->get_param('client_id');
        $clientSecret = $request->get_param('client_secret');

        if (!empty($clientId) && !empty($clientSecret)) {
            return [$clientId, $clientSecret];
        }

        $auth = $request->get_header('Authorization');
        if ($auth && str_starts_with($auth, 'Basic ')) {
            $decoded = base64_decode(substr($auth, 6));
            if ($decoded && str_contains($decoded, ':')) {
                [$clientId, $clientSecret] = explode(':', $decoded, 2);
                return [trim($clientId), trim($clientSecret)];
            }
        }

        return [$clientId ?? '', ''];
    }

    private function oauthError(string $error, string $description, int $status): \WP_REST_Response {
        return new \WP_REST_Response(['error' => $error, 'error_description' => $description], $status);
    }

    // -------------------------------------------------------------------------
    // Approval page UI
    // -------------------------------------------------------------------------
    private function renderApprovalPage(array $row, string $clientId, string $redirectUri, string $state, string $codeChallenge, string $challengeMethod): void {
        $siteName  = get_bloginfo('name');
        $keyName   = esc_html($row['name']);
        $scopes    = array_map('trim', explode(',', $row['scopes']));
        $actionUrl = rest_url('asktam-mcp/v1/oauth/authorize');

        while (ob_get_level()) ob_end_clean();
        status_header(200);
        header('Content-Type: text/html; charset=UTF-8');
        header('X-Robots-Tag: noindex');
        ?>
<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1">
    <title>Authorize — <?php echo esc_html($siteName); ?></title>
    <style>
        * { box-sizing: border-box; margin: 0; padding: 0; }
        body { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif; background: #f0f2f5; display: flex; align-items: center; justify-content: center; min-height: 100vh; padding: 20px; }
        .card { background: #fff; border-radius: 16px; box-shadow: 0 4px 32px rgba(0,0,0,.12); max-width: 420px; width: 100%; padding: 40px 36px; }
        .logo { display: flex; align-items: center; gap: 12px; margin-bottom: 32px; }
        .logo strong { font-size: 17px; color: #1a1a1a; }
        h2 { font-size: 22px; color: #1a1a1a; margin-bottom: 10px; }
        p { color: #666; font-size: 14px; line-height: 1.65; margin-bottom: 22px; }
        .key-box { background: #f8f9ff; border: 1.5px solid #d0d9ff; border-radius: 10px; padding: 16px 18px; margin-bottom: 28px; }
        .key-box .lbl { font-size: 11px; font-weight: 600; text-transform: uppercase; letter-spacing: .6px; color: #888; margin-bottom: 5px; }
        .key-box .val { font-size: 15px; font-weight: 700; color: #1a1a1a; margin-bottom: 10px; }
        .scopes { display: flex; flex-wrap: wrap; gap: 6px; }
        .scope { background: #ede9fe; color: #5b21b6; border-radius: 5px; padding: 3px 10px; font-size: 12px; font-weight: 600; }
        .security-note { background: #fff8ed; border: 1px solid #f5d87a; border-radius: 8px; padding: 10px 14px; margin-bottom: 20px; font-size: 12px; color: #7a5c00; }
        .btns { display: flex; gap: 10px; }
        .btn { flex: 1; padding: 13px; border-radius: 9px; border: none; font-size: 15px; font-weight: 700; cursor: pointer; }
        .approve { background: #2271b1; color: #fff; }
        .deny { background: #f5f5f5; color: #444; border: 1.5px solid #ddd; }
        .note { margin-top: 20px; font-size: 12px; color: #aaa; text-align: center; line-height: 1.5; }
    </style>
</head>
<body>
<div class="card">
    <div class="logo">
        <svg width="36" height="36" viewBox="0 0 36 36" fill="none"><circle cx="18" cy="18" r="18" fill="#2271b1"/><path d="M11 26L18 11L25 26" stroke="white" stroke-width="2.5" stroke-linecap="round" stroke-linejoin="round"/><path d="M13.5 21H22.5" stroke="white" stroke-width="2.5" stroke-linecap="round"/></svg>
        <strong>AskTam MCP &mdash; <?php echo esc_html($siteName); ?></strong>
    </div>

    <h2>Authorize Access</h2>
    <p>An AI client is requesting access to your WordPress site using this API key:</p>

    <div class="key-box">
        <div class="lbl">API Key</div>
        <div class="val"><?php echo $keyName; ?></div>
        <div class="scopes">
            <?php foreach ($scopes as $s): ?>
                <span class="scope"><?php echo esc_html($s); ?></span>
            <?php endforeach; ?>
        </div>
    </div>

    <div class="security-note">
        🔒 Access token expires in <strong>1 hour</strong>. Your API key is never shared with the client.
    </div>

    <form method="post" action="<?php echo esc_url($actionUrl); ?>">
        <input type="hidden" name="client_id"          value="<?php echo esc_attr($clientId); ?>">
        <input type="hidden" name="redirect_uri"        value="<?php echo esc_attr($redirectUri); ?>">
        <input type="hidden" name="state"               value="<?php echo esc_attr($state); ?>">
        <input type="hidden" name="code_challenge"      value="<?php echo esc_attr($codeChallenge); ?>">
        <input type="hidden" name="code_challenge_method" value="<?php echo esc_attr($challengeMethod); ?>">
        <input type="hidden" name="_mcp_approve"        value="1">
        <div class="btns">
            <button type="submit" class="btn approve">Approve</button>
            <button type="button" class="btn deny" onclick="window.location='<?php echo esc_js(add_query_arg(['error' => 'access_denied', 'state' => $state], $redirectUri)); ?>'">Deny</button>
        </div>
    </form>

    <div class="note">Only approve if you initiated this from Claude.ai or a trusted MCP client.</div>
</div>
</body>
</html>
        <?php
        exit;
    }
}