<?php
/**
 * JWT Authentication Service for Care API
 *
 * Enhanced JWT service with Firebase JWT library integration,
 * 2024 security best practices, and healthcare compliance features.
 *
 * @package Care_API
 * @subpackage Services
 * @version    1.0.0
 * @author     Descomplicar® <dev@descomplicar.pt>
 * @link       https://descomplicar.pt
 * @since      1.0.0
 */

namespace Care_API\Services;

// Exit if accessed directly.
if ( ! defined( 'ABSPATH' ) ) {
    exit;
}

use Firebase\JWT\JWT;
use Firebase\JWT\Key;
use Firebase\JWT\ExpiredException;
use Firebase\JWT\SignatureInvalidException;
use Firebase\JWT\BeforeValidException;
use WP_Error;

/**
 * Class JWT_Service
 *
 * JWT authentication service with healthcare compliance and modern security practices
 *
 * @since 1.0.0
 */
class JWT_Service {

    /**
     * JWT secret key
     *
     * @var string
     */
    private static $secret_key = null;

    /**
     * Access token expiration time (10 minutes in seconds) - 2024 security best practice
     *
     * @var int
     */
    private static $access_token_expiration = 600;

    /**
     * Refresh token expiration time (7 days in seconds)
     *
     * @var int
     */
    private static $refresh_token_expiration = 604800;

    /**
     * Supported JWT algorithms
     *
     * @var array
     */
    private static $supported_algorithms = array( 'HS256', 'RS256' );

    /**
     * Algorithm to use for JWT signing
     *
     * @var string
     */
    private static $algorithm = 'HS256';

    /**
     * Initialize the service
     *
     * @since 1.0.0
     */
    public static function init() {
        self::$secret_key = self::get_secret_key();
        
        // Configure JWT algorithm based on settings
        self::$algorithm = apply_filters( 'kivicare_jwt_algorithm', self::$algorithm );
        
        // Validate algorithm support
        if ( ! in_array( self::$algorithm, self::$supported_algorithms ) ) {
            self::$algorithm = 'HS256';
        }

        // Hook into WordPress authentication
        add_filter( 'determine_current_user', array( self::class, 'determine_current_user' ), 20 );
        
        // Add REST API authentication check
        add_filter( 'rest_authentication_errors', array( self::class, 'rest_authentication_errors' ) );
    }

    /**
     * Get the secret key for JWT signing
     *
     * @return string Secret key
     * @since 1.0.0
     */
    private static function get_secret_key() {
        $key = get_option( 'kivicare_jwt_secret' );
        
        if ( empty( $key ) ) {
            // Generate a cryptographically secure secret key (256-bit for HS256)
            $key = base64_encode( random_bytes( 32 ) );
            update_option( 'kivicare_jwt_secret', $key );
            
            // Log key generation for security audit
            error_log( 'KiviCare JWT: New secret key generated' );
        }
        
        return base64_decode( $key );
    }

    /**
     * Generate access and refresh JWT tokens for a user
     *
     * @param int   $user_id      WordPress user ID
     * @param array $extra_claims Additional claims to include
     * @return array|WP_Error Token pair array or error
     * @since 1.0.0
     */
    public static function generate_tokens( $user_id, $extra_claims = array() ) {
        $user = get_user_by( 'id', $user_id );
        
        if ( ! $user ) {
            return new WP_Error( 
                'invalid_user', 
                __( 'User not found', 'kivicare-api' ), 
                array( 'status' => 404 ) 
            );
        }

        // Check if user account is active
        if ( get_user_meta( $user_id, 'kivicare_account_status', true ) === 'inactive' ) {
            return new WP_Error(
                'account_inactive',
                __( 'User account is inactive', 'kivicare-api' ),
                array( 'status' => 403 )
            );
        }

        $current_time = time();
        $jti_access = wp_generate_uuid4(); // Unique token ID
        $jti_refresh = wp_generate_uuid4();
        
        // Get user's primary kivicare role for healthcare context
        $primary_role = Permission_Service::get_primary_kivicare_role( $user );
        $clinic_ids = Permission_Service::get_accessible_clinic_ids( $user );
        
        // Generate session for tracking
        $session_id = Session_Service::create_session( 
            $user_id, 
            self::get_client_ip(), 
            $_SERVER['HTTP_USER_AGENT'] ?? '' 
        );

        // Access token payload
        $access_payload = array_merge(
            array(
                'iss' => get_bloginfo( 'url' ),      // Issuer
                'aud' => get_bloginfo( 'url' ),      // Audience
                'iat' => $current_time,              // Issued at
                'exp' => $current_time + self::$access_token_expiration,  // Expiration
                'nbf' => $current_time,              // Not before
                'jti' => $jti_access,                // JWT ID
                'type' => 'access',                  // Token type
                'user_id' => $user_id,
                'username' => $user->user_login,
                'user_email' => $user->user_email,
                'user_roles' => $user->roles,
                'primary_role' => $primary_role,
                'clinic_ids' => $clinic_ids,
                'session_id' => $session_id,
                'capabilities' => Permission_Service::get_user_permissions( $user ),
                'ip' => self::get_client_ip(),       // IP binding for security
            ),
            $extra_claims
        );

        // Refresh token payload (minimal data)
        $refresh_payload = array(
            'iss' => get_bloginfo( 'url' ),
            'aud' => get_bloginfo( 'url' ),
            'iat' => $current_time,
            'exp' => $current_time + self::$refresh_token_expiration,
            'nbf' => $current_time,
            'jti' => $jti_refresh,
            'type' => 'refresh',
            'user_id' => $user_id,
            'session_id' => $session_id,
            'access_jti' => $jti_access // Link to access token
        );

        try {
            $access_token = JWT::encode( $access_payload, self::$secret_key, self::$algorithm );
            $refresh_token = JWT::encode( $refresh_payload, self::$secret_key, self::$algorithm );

            // Store token metadata for revocation capabilities
            self::store_token_metadata( $jti_access, $user_id, 'access', $current_time + self::$access_token_expiration );
            self::store_token_metadata( $jti_refresh, $user_id, 'refresh', $current_time + self::$refresh_token_expiration );

            // Log successful token generation for healthcare audit
            self::log_token_event( $user_id, 'token_generated', array(
                'access_jti' => $jti_access,
                'refresh_jti' => $jti_refresh,
                'session_id' => $session_id,
                'ip_address' => self::get_client_ip()
            ) );

            return array(
                'access_token' => $access_token,
                'refresh_token' => $refresh_token,
                'token_type' => 'Bearer',
                'expires_in' => self::$access_token_expiration,
                'refresh_expires_in' => self::$refresh_token_expiration,
                'session_id' => $session_id
            );

        } catch ( Exception $e ) {
            return new WP_Error(
                'token_generation_failed',
                __( 'Failed to generate tokens', 'kivicare-api' ),
                array( 'status' => 500 )
            );
        }
    }

    /**
     * Validate JWT token with comprehensive security checks
     *
     * @param string $token JWT token
     * @param string $expected_type Expected token type ('access' or 'refresh')
     * @return array|WP_Error Decoded payload or error
     * @since 1.0.0
     */
    public static function validate_token( $token, $expected_type = 'access' ) {
        if ( empty( $token ) ) {
            return new WP_Error( 
                'missing_token', 
                __( 'Token is required', 'kivicare-api' ), 
                array( 'status' => 401 ) 
            );
        }

        try {
            // Decode using Firebase JWT library
            $payload = (array) JWT::decode( $token, new Key( self::$secret_key, self::$algorithm ) );
            
            // Validate token type
            if ( isset( $payload['type'] ) && $payload['type'] !== $expected_type ) {
                return new WP_Error( 
                    'invalid_token_type', 
                    sprintf( __( 'Expected %s token', 'kivicare-api' ), $expected_type ),
                    array( 'status' => 401 ) 
                );
            }

            // Verify user still exists and is active
            $user = get_user_by( 'id', $payload['user_id'] );
            if ( ! $user ) {
                return new WP_Error( 
                    'invalid_user', 
                    __( 'User no longer exists', 'kivicare-api' ), 
                    array( 'status' => 401 ) 
                );
            }

            // Check if user account is still active
            if ( get_user_meta( $payload['user_id'], 'kivicare_account_status', true ) === 'inactive' ) {
                return new WP_Error(
                    'account_inactive',
                    __( 'User account is inactive', 'kivicare-api' ),
                    array( 'status' => 403 )
                );
            }

            // Check if token is revoked
            if ( self::is_token_revoked( $payload['jti'] ) ) {
                return new WP_Error( 
                    'token_revoked', 
                    __( 'Token has been revoked', 'kivicare-api' ), 
                    array( 'status' => 401 ) 
                );
            }

            // Validate session if present
            if ( isset( $payload['session_id'] ) ) {
                $session = Session_Service::validate_session( $payload['session_id'], $payload['user_id'] );
                if ( ! $session ) {
                    return new WP_Error( 
                        'invalid_session', 
                        __( 'Session is invalid or expired', 'kivicare-api' ), 
                        array( 'status' => 401 ) 
                    );
                }
            }

            // IP binding validation for access tokens (if enabled)
            if ( $expected_type === 'access' && apply_filters( 'kivicare_jwt_ip_binding', false ) ) {
                if ( isset( $payload['ip'] ) && $payload['ip'] !== self::get_client_ip() ) {
                    return new WP_Error( 
                        'ip_mismatch', 
                        __( 'Token IP mismatch', 'kivicare-api' ), 
                        array( 'status' => 401 ) 
                    );
                }
            }

            // Healthcare compliance: Log token usage for audit trail
            self::log_token_event( $payload['user_id'], 'token_validated', array(
                'jti' => $payload['jti'],
                'type' => $expected_type,
                'ip_address' => self::get_client_ip()
            ) );

            return $payload;
            
        } catch ( ExpiredException $e ) {
            return new WP_Error( 
                'token_expired', 
                __( 'Token has expired', 'kivicare-api' ), 
                array( 'status' => 401 ) 
            );
        } catch ( SignatureInvalidException $e ) {
            return new WP_Error( 
                'invalid_signature', 
                __( 'Token signature is invalid', 'kivicare-api' ), 
                array( 'status' => 401 ) 
            );
        } catch ( BeforeValidException $e ) {
            return new WP_Error( 
                'token_not_yet_valid', 
                __( 'Token is not yet valid', 'kivicare-api' ), 
                array( 'status' => 401 ) 
            );
        } catch ( Exception $e ) {
            return new WP_Error( 
                'invalid_token', 
                sprintf( __( 'Invalid token: %s', 'kivicare-api' ), $e->getMessage() ),
                array( 'status' => 401 ) 
            );
        }
    }

    /**
     * Refresh access token using refresh token
     *
     * @param string $refresh_token Refresh token
     * @return array|WP_Error New token pair or error
     * @since 1.0.0
     */
    public static function refresh_token( $refresh_token ) {
        // Validate refresh token
        $payload = self::validate_token( $refresh_token, 'refresh' );
        
        if ( is_wp_error( $payload ) ) {
            return $payload;
        }

        // Revoke the used refresh token
        self::revoke_token( $payload['jti'] );

        // Generate new token pair
        $new_tokens = self::generate_tokens( $payload['user_id'] );
        
        if ( is_wp_error( $new_tokens ) ) {
            return $new_tokens;
        }

        // Log token refresh for audit
        self::log_token_event( $payload['user_id'], 'token_refreshed', array(
            'old_refresh_jti' => $payload['jti'],
            'new_access_jti' => $new_tokens['access_token'],
            'ip_address' => self::get_client_ip()
        ) );

        return $new_tokens;
    }

    /**
     * Revoke a token by JTI
     *
     * @param string $jti Token JTI
     * @return bool Success status
     * @since 1.0.0
     */
    public static function revoke_token( $jti ) {
        global $wpdb;

        $result = $wpdb->update(
            $wpdb->prefix . 'kivicare_jwt_tokens',
            array( 'is_revoked' => 1, 'revoked_at' => current_time( 'mysql' ) ),
            array( 'jti' => $jti ),
            array( '%d', '%s' ),
            array( '%s' )
        );

        return $result !== false;
    }

    /**
     * Revoke all tokens for a user
     *
     * @param int $user_id User ID
     * @return bool Success status
     * @since 1.0.0
     */
    public static function revoke_user_tokens( $user_id ) {
        global $wpdb;

        $result = $wpdb->update(
            $wpdb->prefix . 'kivicare_jwt_tokens',
            array( 'is_revoked' => 1, 'revoked_at' => current_time( 'mysql' ) ),
            array( 'user_id' => $user_id, 'is_revoked' => 0 ),
            array( '%d', '%s' ),
            array( '%d', '%d' )
        );

        // Also expire user sessions
        Session_Service::expire_user_sessions( $user_id );

        return $result !== false;
    }

    /**
     * Extract token from Authorization header
     * 
     * @param string $authorization_header Authorization header value
     * @return string|null Token or null if not found
     * @since 1.0.0
     */
    public static function extract_token_from_header( $authorization_header ) {
        if ( empty( $authorization_header ) ) {
            return null;
        }
        
        // Remove "Bearer " prefix
        if ( strpos( $authorization_header, 'Bearer ' ) === 0 ) {
            return trim( substr( $authorization_header, 7 ) );
        }
        
        return null; // Only accept Bearer tokens
    }

    /**
     * Get current user ID from JWT token in request
     * 
     * @return int|null User ID or null if not authenticated
     * @since 1.0.0
     */
    public static function get_current_user_from_token() {
        $authorization = self::get_authorization_header();
        
        if ( empty( $authorization ) ) {
            return null;
        }
        
        $token = self::extract_token_from_header( $authorization );
        if ( empty( $token ) ) {
            return null;
        }
        
        $payload = self::validate_token( $token );
        
        if ( is_wp_error( $payload ) ) {
            return null;
        }
        
        return $payload['user_id'] ?? null;
    }

    /**
     * Check if current request has valid JWT authentication
     * 
     * @return bool|WP_Error True if authenticated, WP_Error if not
     * @since 1.0.0
     */
    public static function check_jwt_authentication() {
        $authorization = self::get_authorization_header();
        
        if ( empty( $authorization ) ) {
            return new WP_Error( 
                'missing_authorization', 
                __( 'Authorization header is missing', 'kivicare-api' ),
                array( 'status' => 401 ) 
            );
        }
        
        $token = self::extract_token_from_header( $authorization );
        if ( empty( $token ) ) {
            return new WP_Error( 
                'invalid_authorization_format', 
                __( 'Invalid authorization format. Expected: Bearer <token>', 'kivicare-api' ),
                array( 'status' => 401 ) 
            );
        }
        
        $payload = self::validate_token( $token );
        
        if ( is_wp_error( $payload ) ) {
            return $payload;
        }
        
        // Set current user for WordPress context
        wp_set_current_user( $payload['user_id'] );
        
        // Update session activity
        if ( isset( $payload['session_id'] ) ) {
            Session_Service::update_session_activity( $payload['session_id'] );
        }
        
        return true;
    }

    /**
     * WordPress hook: Determine current user from JWT
     *
     * @param int $user_id Current user ID
     * @return int User ID
     * @since 1.0.0
     */
    public static function determine_current_user( $user_id ) {
        // Skip if user already determined
        if ( $user_id ) {
            return $user_id;
        }

        // Only for REST API requests
        if ( ! defined( 'REST_REQUEST' ) || ! REST_REQUEST ) {
            return $user_id;
        }

        return self::get_current_user_from_token() ?: $user_id;
    }

    /**
     * WordPress hook: REST API authentication errors
     *
     * @param WP_Error|null|true $error Authentication error
     * @return WP_Error|null|true Authentication result
     * @since 1.0.0
     */
    public static function rest_authentication_errors( $error ) {
        // Pass through if already authenticated or has error
        if ( $error !== null ) {
            return $error;
        }

        // Check for JWT authentication
        $authorization = self::get_authorization_header();
        
        if ( ! empty( $authorization ) ) {
            $auth_result = self::check_jwt_authentication();
            
            if ( is_wp_error( $auth_result ) ) {
                return $auth_result;
            }
        }

        return null;
    }

    /**
     * Get authorization header from request
     *
     * @return string|null Authorization header value
     * @since 1.0.0
     */
    private static function get_authorization_header() {
        $headers = array();
        
        // Try different methods to get headers
        if ( function_exists( 'getallheaders' ) ) {
            $headers = getallheaders();
        } elseif ( function_exists( 'apache_request_headers' ) ) {
            $headers = apache_request_headers();
        } else {
            // Fallback to $_SERVER
            foreach ( $_SERVER as $key => $value ) {
                if ( strpos( $key, 'HTTP_' ) === 0 ) {
                    $header_key = str_replace( ' ', '-', ucwords( strtolower( str_replace( '_', ' ', substr( $key, 5 ) ) ) ) );
                    $headers[ $header_key ] = $value;
                }
            }
        }

        return $headers['Authorization'] ?? $_SERVER['HTTP_AUTHORIZATION'] ?? null;
    }

    /**
     * Store token metadata for revocation tracking
     *
     * @param string $jti     Token JTI
     * @param int    $user_id User ID
     * @param string $type    Token type
     * @param int    $expires Expiration timestamp
     * @return bool Success status
     * @since 1.0.0
     */
    private static function store_token_metadata( $jti, $user_id, $type, $expires ) {
        global $wpdb;

        // Create table if needed
        self::create_tokens_table();

        $result = $wpdb->insert(
            $wpdb->prefix . 'kivicare_jwt_tokens',
            array(
                'jti' => $jti,
                'user_id' => $user_id,
                'token_type' => $type,
                'expires_at' => date( 'Y-m-d H:i:s', $expires ),
                'created_at' => current_time( 'mysql' ),
                'is_revoked' => 0
            ),
            array( '%s', '%d', '%s', '%s', '%s', '%d' )
        );

        return $result !== false;
    }

    /**
     * Check if token is revoked
     *
     * @param string $jti Token JTI
     * @return bool True if revoked
     * @since 1.0.0
     */
    private static function is_token_revoked( $jti ) {
        global $wpdb;

        $is_revoked = $wpdb->get_var(
            $wpdb->prepare(
                "SELECT is_revoked FROM {$wpdb->prefix}kivicare_jwt_tokens WHERE jti = %s",
                $jti
            )
        );

        return (bool) $is_revoked;
    }

    /**
     * Log token-related events for healthcare audit trail
     *
     * @param int    $user_id User ID
     * @param string $event   Event type
     * @param array  $data    Event data
     * @since 1.0.0
     */
    private static function log_token_event( $user_id, $event, $data = array() ) {
        // Use session service logging for consistency
        if ( class_exists( 'Care_API\Services\Session_Service' ) ) {
            // This would ideally call a logging method, but session service has private method
            // So we'll create our own minimal logging
        }

        // Log to WordPress error log for development
        if ( defined( 'WP_DEBUG' ) && WP_DEBUG ) {
            error_log( sprintf(
                'KiviCare JWT Event: %s for user %d - %s',
                $event,
                $user_id,
                wp_json_encode( $data )
            ) );
        }
    }

    /**
     * Get client IP address
     *
     * @return string IP address
     * @since 1.0.0
     */
    private static function get_client_ip() {
        $ip_headers = array(
            'HTTP_CF_CONNECTING_IP',
            'HTTP_X_FORWARDED_FOR',
            'HTTP_X_FORWARDED',
            'HTTP_X_CLUSTER_CLIENT_IP',
            'HTTP_FORWARDED_FOR',
            'HTTP_FORWARDED',
            'REMOTE_ADDR'
        );

        foreach ( $ip_headers as $header ) {
            if ( ! empty( $_SERVER[ $header ] ) ) {
                $ips = explode( ',', $_SERVER[ $header ] );
                $ip = trim( $ips[0] );
                
                if ( filter_var( $ip, FILTER_VALIDATE_IP, FILTER_FLAG_NO_PRIV_RANGE | FILTER_FLAG_NO_RES_RANGE ) ) {
                    return $ip;
                }
            }
        }

        return $_SERVER['REMOTE_ADDR'] ?? '0.0.0.0';
    }

    /**
     * Create JWT tokens table for tracking and revocation
     *
     * @since 1.0.0
     */
    private static function create_tokens_table() {
        global $wpdb;

        $table_name = $wpdb->prefix . 'kivicare_jwt_tokens';
        
        $charset_collate = $wpdb->get_charset_collate();

        $sql = "CREATE TABLE IF NOT EXISTS {$table_name} (
            id bigint(20) unsigned NOT NULL AUTO_INCREMENT,
            jti varchar(36) NOT NULL,
            user_id bigint(20) unsigned NOT NULL,
            token_type varchar(10) NOT NULL,
            created_at datetime NOT NULL,
            expires_at datetime NOT NULL,
            revoked_at datetime NULL,
            is_revoked tinyint(1) NOT NULL DEFAULT 0,
            PRIMARY KEY (id),
            UNIQUE KEY jti (jti),
            KEY user_id (user_id),
            KEY token_type (token_type),
            KEY expires_at (expires_at),
            KEY is_revoked (is_revoked)
        ) {$charset_collate};";

        require_once( ABSPATH . 'wp-admin/includes/upgrade.php' );
        dbDelta( $sql );
    }

    /**
     * Clean up expired and revoked tokens
     *
     * @since 1.0.0
     */
    public static function cleanup_expired_tokens() {
        global $wpdb;

        // Delete expired tokens
        $wpdb->query(
            "DELETE FROM {$wpdb->prefix}kivicare_jwt_tokens 
             WHERE expires_at < NOW() OR (is_revoked = 1 AND revoked_at < DATE_SUB(NOW(), INTERVAL 30 DAY))"
        );
    }
}

// Initialize the service
add_action( 'init', array( 'Care_API\Services\JWT_Service', 'init' ) );