<?php
/**
 * OAuth2 Server Implementation
 * 
 * Implements secure OAuth2 authorization server functionality
 * with support for authorization code flow with PKCE
 * 
 * @package WordPress OAuth2 PoC
 */

namespace WPOAuth2Server\Core;

defined('ABSPATH') or die('Direct access forbidden.');

class OAuth2Server {
    
    /**
     * WordPress database instance
     * @var \wpdb
     */
    private \wpdb $wpdb;
    
    /**
     * Plugin settings
     * @var array
     */
    private array $settings;
    
    /**
     * Constructor
     */
    public function __construct(array $settings) {
        global $wpdb;
        $this->wpdb = $wpdb;
        $this->settings = $settings;
    }
    
    /**
     * Create OAuth2 database tables
     */
    public function create_tables(): void {
        $charset_collate = $this->wpdb->get_charset_collate();
        
        // OAuth2 Clients table
        $clients_table = $this->wpdb->prefix . 'oauth2_clients';
        $sql_clients = "CREATE TABLE $clients_table (
            client_id varchar(80) NOT NULL,
            client_secret varchar(255) NOT NULL,
            client_name varchar(255) NOT NULL,
            redirect_uri text NOT NULL,
            grant_types varchar(255) DEFAULT 'authorization_code',
            scope text,
            user_id bigint(20) NOT NULL,
            is_public tinyint(1) DEFAULT 0,
            created_at datetime DEFAULT CURRENT_TIMESTAMP,
            updated_at datetime DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
            PRIMARY KEY (client_id),
            KEY user_id (user_id)
        ) $charset_collate;";
        
        // Authorization Codes table
        $auth_codes_table = $this->wpdb->prefix . 'oauth2_authorization_codes';
        $sql_auth_codes = "CREATE TABLE $auth_codes_table (
            authorization_code varchar(255) NOT NULL,
            client_id varchar(80) NOT NULL,
            user_id bigint(20) NOT NULL,
            redirect_uri text,
            expires datetime NOT NULL,
            scope text,
            code_challenge varchar(255),
            code_challenge_method varchar(20),
            created_at datetime DEFAULT CURRENT_TIMESTAMP,
            PRIMARY KEY (authorization_code),
            KEY client_id (client_id),
            KEY user_id (user_id),
            KEY expires (expires)
        ) $charset_collate;";
        
        // Access Tokens table  
        $access_tokens_table = $this->wpdb->prefix . 'oauth2_access_tokens';
        $sql_access_tokens = "CREATE TABLE $access_tokens_table (
            access_token varchar(255) NOT NULL,
            client_id varchar(80) NOT NULL,
            user_id bigint(20) NOT NULL,
            expires datetime NOT NULL,
            scope text,
            created_at datetime DEFAULT CURRENT_TIMESTAMP,
            PRIMARY KEY (access_token),
            KEY client_id (client_id),
            KEY user_id (user_id),
            KEY expires (expires)
        ) $charset_collate;";
        
        // Refresh Tokens table
        $refresh_tokens_table = $this->wpdb->prefix . 'oauth2_refresh_tokens';
        $sql_refresh_tokens = "CREATE TABLE $refresh_tokens_table (
            refresh_token varchar(255) NOT NULL,
            client_id varchar(80) NOT NULL,
            user_id bigint(20) NOT NULL,
            expires datetime NOT NULL,
            scope text,
            created_at datetime DEFAULT CURRENT_TIMESTAMP,
            PRIMARY KEY (refresh_token),
            KEY client_id (client_id),
            KEY user_id (user_id),
            KEY expires (expires)
        ) $charset_collate;";
        
        require_once ABSPATH . 'wp-admin/includes/upgrade.php';
        dbDelta($sql_clients);
        dbDelta($sql_auth_codes);
        dbDelta($sql_access_tokens);
        dbDelta($sql_refresh_tokens);
    }
    
    /**
     * Handle authorization endpoint request
     */
    public function handle_authorization_request(): void {
        // Validate request method
        if ($_SERVER['REQUEST_METHOD'] !== 'GET') {
            $this->send_error_response(405, 'invalid_request', 'Authorization endpoint requires GET method');
            return;
        }
        
        // Extract and validate parameters
        $client_id = sanitize_text_field($_GET['client_id'] ?? '');
        $redirect_uri = esc_url_raw($_GET['redirect_uri'] ?? '');
        $response_type = sanitize_text_field($_GET['response_type'] ?? '');
        $scope = sanitize_text_field($_GET['scope'] ?? '');
        $state = sanitize_text_field($_GET['state'] ?? '');
        $code_challenge = sanitize_text_field($_GET['code_challenge'] ?? '');
        $code_challenge_method = sanitize_text_field($_GET['code_challenge_method'] ?? '');
        
        // Validate required parameters
        if (empty($client_id) || empty($redirect_uri) || empty($response_type)) {
            $this->send_error_response(400, 'invalid_request', 'Missing required parameters');
            return;
        }
        
        // Validate response type
        if ($response_type !== 'code') {
            $this->send_error_response(400, 'unsupported_response_type', 'Only authorization code flow is supported');
            return;
        }
        
        // Validate client
        $client = $this->get_client($client_id);
        if (!$client) {
            $this->send_error_response(400, 'invalid_client', 'Invalid client_id');
            return;
        }
        
        // Validate redirect URI
        if (!$this->validate_redirect_uri($client, $redirect_uri)) {
            $this->send_error_response(400, 'invalid_redirect_uri', 'Invalid redirect_uri');
            return;
        }
        
        // Validate PKCE (required for public clients)
        if ($client['is_public'] && (empty($code_challenge) || $code_challenge_method !== 'S256')) {
            $this->send_authorization_error($redirect_uri, 'invalid_request', 'PKCE required for public clients', $state);
            return;
        }
        
        // Check if user is logged in
        if (!is_user_logged_in()) {
            // Redirect to login with return URL
            $login_url = wp_login_url($_SERVER['REQUEST_URI']);
            wp_redirect($login_url);
            exit;
        }
        
        // Show authorization consent form
        $this->show_authorization_form([
            'client' => $client,
            'redirect_uri' => $redirect_uri,
            'scope' => $scope,
            'state' => $state,
            'code_challenge' => $code_challenge,
            'code_challenge_method' => $code_challenge_method,
        ]);
    }
    
    /**
     * Handle token endpoint request
     */
    public function handle_token_request(): void {
        // Validate request method
        if ($_SERVER['REQUEST_METHOD'] !== 'POST') {
            $this->send_error_response(405, 'invalid_request', 'Token endpoint requires POST method');
            return;
        }
        
        // Get grant type
        $grant_type = sanitize_text_field($_POST['grant_type'] ?? '');
        
        switch ($grant_type) {
            case 'authorization_code':
                $this->handle_authorization_code_grant();
                break;
            case 'refresh_token':
                $this->handle_refresh_token_grant();
                break;
            case 'client_credentials':
                $this->handle_client_credentials_grant();
                break;
            default:
                $this->send_error_response(400, 'unsupported_grant_type', 'Unsupported grant type');
        }
    }
    
    /**
     * Handle authorization code grant
     */
    private function handle_authorization_code_grant(): void {
        // Extract parameters
        $code = sanitize_text_field($_POST['code'] ?? '');
        $client_id = sanitize_text_field($_POST['client_id'] ?? '');
        $redirect_uri = esc_url_raw($_POST['redirect_uri'] ?? '');
        $code_verifier = sanitize_text_field($_POST['code_verifier'] ?? '');
        
        // Get client credentials from Authorization header or POST
        $client_credentials = $this->extract_client_credentials();
        if ($client_credentials) {
            $client_id = $client_credentials['client_id'];
            $client_secret = $client_credentials['client_secret'];
        } else {
            $client_secret = sanitize_text_field($_POST['client_secret'] ?? '');
        }
        
        // Validate required parameters
        if (empty($code) || empty($client_id)) {
            $this->send_error_response(400, 'invalid_request', 'Missing required parameters');
            return;
        }
        
        // Get and validate authorization code
        $auth_code = $this->get_authorization_code($code);
        if (!$auth_code) {
            $this->send_error_response(400, 'invalid_grant', 'Invalid authorization code');
            return;
        }
        
        // Check if code has expired
        if (strtotime($auth_code['expires']) < time()) {
            $this->delete_authorization_code($code);
            $this->send_error_response(400, 'invalid_grant', 'Authorization code expired');
            return;
        }
        
        // Validate client
        if ($auth_code['client_id'] !== $client_id) {
            $this->send_error_response(400, 'invalid_client', 'Client mismatch');
            return;
        }
        
        // Get client details
        $client = $this->get_client($client_id);
        if (!$client) {
            $this->send_error_response(400, 'invalid_client', 'Invalid client');
            return;
        }
        
        // Validate client secret for confidential clients
        if (!$client['is_public']) {
            if (empty($client_secret) || !password_verify($client_secret, $client['client_secret'])) {
                $this->send_error_response(401, 'invalid_client', 'Invalid client credentials');
                return;
            }
        }
        
        // Validate redirect URI
        if (!empty($redirect_uri) && $auth_code['redirect_uri'] !== $redirect_uri) {
            $this->send_error_response(400, 'invalid_grant', 'Redirect URI mismatch');
            return;
        }
        
        // Validate PKCE for public clients or if code challenge was provided
        if (!empty($auth_code['code_challenge'])) {
            if (empty($code_verifier)) {
                $this->send_error_response(400, 'invalid_request', 'Code verifier required');
                return;
            }
            
            if (!$this->verify_pkce($code_verifier, $auth_code['code_challenge'], $auth_code['code_challenge_method'])) {
                $this->send_error_response(400, 'invalid_grant', 'Invalid code verifier');
                return;
            }
        }
        
        // Generate tokens
        $access_token = $this->generate_access_token($client_id, $auth_code['user_id'], $auth_code['scope']);
        $refresh_token = $this->generate_refresh_token($client_id, $auth_code['user_id'], $auth_code['scope']);
        
        // Delete authorization code (one-time use)
        $this->delete_authorization_code($code);
        
        // Send token response
        $response = [
            'access_token' => $access_token,
            'token_type' => 'Bearer',
            'expires_in' => $this->settings['access_token_lifetime'] ?? 3600,
            'refresh_token' => $refresh_token,
            'scope' => $auth_code['scope'],
        ];
        
        $this->send_json_response($response);
    }
    
    /**
     * Generate authorization code
     */
    public function generate_authorization_code(array $params): string {
        $code = bin2hex(random_bytes(32));
        $expires = date('Y-m-d H:i:s', time() + 600); // 10 minutes
        
        $this->wpdb->insert(
            $this->wpdb->prefix . 'oauth2_authorization_codes',
            [
                'authorization_code' => $code,
                'client_id' => $params['client_id'],
                'user_id' => get_current_user_id(),
                'redirect_uri' => $params['redirect_uri'],
                'expires' => $expires,
                'scope' => $params['scope'] ?? '',
                'code_challenge' => $params['code_challenge'] ?? '',
                'code_challenge_method' => $params['code_challenge_method'] ?? '',
            ]
        );
        
        return $code;
    }
    
    /**
     * Get authorization code from database
     */
    private function get_authorization_code(string $code): ?array {
        $result = $this->wpdb->get_row(
            $this->wpdb->prepare(
                "SELECT * FROM {$this->wpdb->prefix}oauth2_authorization_codes WHERE authorization_code = %s",
                $code
            ),
            ARRAY_A
        );
        
        return $result ?: null;
    }
    
    /**
     * Delete authorization code
     */
    private function delete_authorization_code(string $code): void {
        $this->wpdb->delete(
            $this->wpdb->prefix . 'oauth2_authorization_codes',
            ['authorization_code' => $code]
        );
    }
    
    /**
     * Generate access token
     */
    private function generate_access_token(string $client_id, int $user_id, string $scope): string {
        $token = bin2hex(random_bytes(32));
        $expires = date('Y-m-d H:i:s', time() + ($this->settings['access_token_lifetime'] ?? 3600));
        
        $this->wpdb->insert(
            $this->wpdb->prefix . 'oauth2_access_tokens',
            [
                'access_token' => $token,
                'client_id' => $client_id,
                'user_id' => $user_id,
                'expires' => $expires,
                'scope' => $scope,
            ]
        );
        
        return $token;
    }
    
    /**
     * Generate refresh token
     */
    private function generate_refresh_token(string $client_id, int $user_id, string $scope): string {
        $token = bin2hex(random_bytes(32));
        $expires = date('Y-m-d H:i:s', time() + ($this->settings['refresh_token_lifetime'] ?? 86400 * 30)); // 30 days
        
        $this->wpdb->insert(
            $this->wpdb->prefix . 'oauth2_refresh_tokens',
            [
                'refresh_token' => $token,
                'client_id' => $client_id,
                'user_id' => $user_id,
                'expires' => $expires,
                'scope' => $scope,
            ]
        );
        
        return $token;
    }
    
    /**
     * Validate access token
     */
    public function validate_access_token(string $token): ?array {
        $result = $this->wpdb->get_row(
            $this->wpdb->prepare(
                "SELECT * FROM {$this->wpdb->prefix}oauth2_access_tokens 
                 WHERE access_token = %s AND expires > NOW()",
                $token
            ),
            ARRAY_A
        );
        
        return $result ?: null;
    }
    
    /**
     * Get client by client_id
     */
    private function get_client(string $client_id): ?array {
        // Query WordPress posts to find OAuth2 client
        $posts = get_posts([
            'post_type' => 'oauth2_client',
            'meta_key' => 'client_id',
            'meta_value' => $client_id,
            'posts_per_page' => 1,
            'post_status' => 'publish'
        ]);
        
        if (empty($posts)) {
            return null;
        }
        
        $post = $posts[0];
        $client_data = [
            'client_id' => get_post_meta($post->ID, 'client_id', true),
            'client_secret' => get_post_meta($post->ID, 'client_secret', true),
            'redirect_uris' => get_post_meta($post->ID, 'redirect_uris', true),
            'grant_types' => get_post_meta($post->ID, 'grant_types', true),
            'scope' => get_post_meta($post->ID, 'scope', true),
            'is_public' => (bool)get_post_meta($post->ID, 'is_public', true),
            'name' => $post->post_title,
            'client_name' => $post->post_title  // Add client_name for display
        ];
        
        return $client_data;
    }
    
    /**
     * Validate redirect URI
     */
    private function validate_redirect_uri(array $client, string $redirect_uri): bool {
        $redirect_uris = $client['redirect_uris'] ?? '';
        if (empty($redirect_uris)) {
            return false;
        }
        
        $allowed_uris = explode(',', $redirect_uris);
        return in_array($redirect_uri, array_map('trim', $allowed_uris));
    }
    
    /**
     * Verify PKCE code verifier
     */
    private function verify_pkce(string $code_verifier, string $code_challenge, string $method): bool {
        if ($method !== 'S256') {
            return false;
        }
        
        $computed_challenge = rtrim(strtr(base64_encode(hash('sha256', $code_verifier, true)), '+/', '-_'), '=');
        return hash_equals($code_challenge, $computed_challenge);
    }
    
    /**
     * Extract client credentials from Authorization header
     */
    private function extract_client_credentials(): ?array {
        $auth_header = $_SERVER['HTTP_AUTHORIZATION'] ?? '';
        
        if (strpos($auth_header, 'Basic ') === 0) {
            $credentials = base64_decode(substr($auth_header, 6));
            $parts = explode(':', $credentials, 2);
            
            if (count($parts) === 2) {
                return [
                    'client_id' => $parts[0],
                    'client_secret' => $parts[1],
                ];
            }
        }
        
        return null;
    }
    
    /**
     * Show authorization consent form
     */
    private function show_authorization_form(array $params): void {
        // Set content type
        header('Content-Type: text/html; charset=utf-8');
        
        // Simple authorization form
        ?>
        <!DOCTYPE html>
        <html <?php language_attributes(); ?>>
        <head>
            <meta charset="<?php bloginfo('charset'); ?>">
            <meta name="viewport" content="width=device-width, initial-scale=1">
            <title><?php _e('Authorize Application', 'wp-oauth2-poc'); ?></title>
            <?php wp_head(); ?>
        </head>
        <body>
            <div style="max-width: 500px; margin: 50px auto; padding: 20px; border: 1px solid #ddd; border-radius: 5px;">
                <h2><?php _e('Authorize Application', 'wp-oauth2-poc'); ?></h2>
                
                <p><?php printf(__('The application "%s" wants to access your account.', 'wp-oauth2-poc'), esc_html($params['client']['client_name'])); ?></p>
                
                <?php if (!empty($params['scope'])): ?>
                <p><strong><?php _e('Requested permissions:', 'wp-oauth2-poc'); ?></strong></p>
                <ul>
                    <?php foreach (explode(' ', $params['scope']) as $scope): ?>
                        <li><?php echo esc_html($scope); ?></li>
                    <?php endforeach; ?>
                </ul>
                <?php endif; ?>
                
                <form method="post" action="<?php echo esc_url(home_url('oauth/authorize')); ?>">
                    <?php wp_nonce_field('oauth2_authorize', 'oauth2_nonce'); ?>
                    <input type="hidden" name="client_id" value="<?php echo esc_attr($params['client']['client_id']); ?>">
                    <input type="hidden" name="redirect_uri" value="<?php echo esc_attr($params['redirect_uri']); ?>">
                    <input type="hidden" name="scope" value="<?php echo esc_attr($params['scope']); ?>">
                    <input type="hidden" name="state" value="<?php echo esc_attr($params['state']); ?>">
                    <input type="hidden" name="code_challenge" value="<?php echo esc_attr($params['code_challenge']); ?>">
                    <input type="hidden" name="code_challenge_method" value="<?php echo esc_attr($params['code_challenge_method']); ?>">
                    
                    <p>
                        <button type="submit" name="authorize" value="yes" class="button button-primary">
                            <?php _e('Authorize', 'wp-oauth2-poc'); ?>
                        </button>
                        <button type="submit" name="authorize" value="no" class="button">
                            <?php _e('Deny', 'wp-oauth2-poc'); ?>
                        </button>
                    </p>
                </form>
            </div>
            <?php wp_footer(); ?>
        </body>
        </html>
        <?php
        exit;
    }
    
    /**
     * Process authorization form submission
     */
    public function process_authorization_form(): void {
        if ($_SERVER['REQUEST_METHOD'] !== 'POST') {
            return;
        }
        
        // Verify nonce
        if (!wp_verify_nonce($_POST['oauth2_nonce'] ?? '', 'oauth2_authorize')) {
            wp_die('Security check failed');
        }
        
        // Check if user denied authorization
        if (($_POST['authorize'] ?? '') !== 'yes') {
            $redirect_uri = esc_url_raw($_POST['redirect_uri'] ?? '');
            $state = sanitize_text_field($_POST['state'] ?? '');
            $this->send_authorization_error($redirect_uri, 'access_denied', 'User denied authorization', $state);
            return;
        }
        
        // Generate authorization code
        $code = $this->generate_authorization_code([
            'client_id' => sanitize_text_field($_POST['client_id'] ?? ''),
            'redirect_uri' => esc_url_raw($_POST['redirect_uri'] ?? ''),
            'scope' => sanitize_text_field($_POST['scope'] ?? ''),
            'code_challenge' => sanitize_text_field($_POST['code_challenge'] ?? ''),
            'code_challenge_method' => sanitize_text_field($_POST['code_challenge_method'] ?? ''),
        ]);
        
        // Redirect back to client with authorization code
        $redirect_uri = esc_url_raw($_POST['redirect_uri'] ?? '');
        $state = sanitize_text_field($_POST['state'] ?? '');
        
        $params = [
            'code' => $code,
        ];
        
        if (!empty($state)) {
            $params['state'] = $state;
        }
        
        $redirect_url = add_query_arg($params, $redirect_uri);
        wp_redirect($redirect_url);
        exit;
    }
    
    /**
     * Send authorization error response
     */
    private function send_authorization_error(string $redirect_uri, string $error, string $description, string $state = ''): void {
        $params = [
            'error' => $error,
            'error_description' => $description,
        ];
        
        if (!empty($state)) {
            $params['state'] = $state;
        }
        
        $redirect_url = add_query_arg($params, $redirect_uri);
        wp_redirect($redirect_url);
        exit;
    }
    
    /**
     * Send JSON response
     */
    private function send_json_response(array $data, int $status_code = 200): void {
        status_header($status_code);
        header('Content-Type: application/json; charset=utf-8');
        header('Cache-Control: no-cache, no-store, must-revalidate');
        header('Pragma: no-cache');
        header('Expires: 0');
        
        echo json_encode($data, JSON_UNESCAPED_SLASHES | JSON_UNESCAPED_UNICODE);
    }
    
    /**
     * Send error response
     */
    private function send_error_response(int $status_code, string $error, string $description): void {
        $error_data = [
            'error' => $error,
            'error_description' => $description,
        ];
        
        $this->send_json_response($error_data, $status_code);
    }
    
    /**
     * Handle refresh token grant (placeholder)
     */
    private function handle_refresh_token_grant(): void {
        $this->send_error_response(501, 'not_implemented', 'Refresh token grant not yet implemented');
    }
    
    /**
     * Handle client credentials grant (placeholder)
     */
    private function handle_client_credentials_grant(): void {
        $this->send_error_response(501, 'not_implemented', 'Client credentials grant not yet implemented');
    }
}