<?php

if ( ! defined( 'ABSPATH' ) ) {
    exit;
}

class Niwa_AI_Auth_Manager implements Niwa_AI_Auth_Provider_Interface {
    const OPTION_KEY = Niwa_AI_Constants::OPT_AUTH;
    const OAUTH_ISSUER = 'https://auth.openai.com';
    const OAUTH_CLIENT_ID = 'app_EMoamEEZ73f0CkXaXp7hrann';
    const OAUTH_SCOPE = 'openid profile email offline_access';
    const OAUTH_ORIGINATOR = 'codex_cli_rs';
    const OAUTH_STATE_TRANSIENT_PREFIX = 'niwa_ai_codex_oauth_state_';
    const OAUTH_STATE_TTL_SECONDS = 900;
    const OAUTH_REFRESH_SKEW_SECONDS = 60;
    const OAUTH_DEVICE_TRANSIENT_PREFIX = 'niwa_ai_codex_device_';
    const OAUTH_DEVICE_VERIFY_URL = 'https://auth.openai.com/codex/device';
    const OAUTH_DEVICE_USERCODE_ENDPOINT = 'https://auth.openai.com/api/accounts/deviceauth/usercode';
    const OAUTH_DEVICE_TOKEN_ENDPOINT = 'https://auth.openai.com/api/accounts/deviceauth/token';
    const OAUTH_DEVICE_CALLBACK_URI = 'https://auth.openai.com/deviceauth/callback';
    const OAUTH_DEVICE_DEFAULT_INTERVAL = 5;
    const OAUTH_DEVICE_MAX_INTERVAL = 20;
    const OAUTH_STATUS_WARN_WINDOW_SECONDS = 600;
    const OAUTH_REFRESH_LOCK_OPTION = Niwa_AI_Constants::OPT_OAUTH_REFRESH_LOCK;
    const OAUTH_REFRESH_LOCK_TTL_SECONDS = 45;
    const OAUTH_REFRESH_LOCK_WAIT_SECONDS = 8;
    const AUTH_CACHE_TRANSIENT_PREFIX = 'niwa_ai_auth_cache_';

    private $lane = 'default';
    private $cached_auth_opts = null;

    public function __construct( $lane = 'default' ) {
        $this->lane = $this->sanitize_lane( $lane );
    }

    public function get_mode() {
        $runtime = $this->resolve_runtime_auth();
        return isset( $runtime['mode'] ) ? (string) $runtime['mode'] : 'api_key';
    }

    public function get_api_key() {
        if ( defined( 'NIWA_AI_OPENAI_API_KEY' ) ) {
            return trim( (string) NIWA_AI_OPENAI_API_KEY );
        }

        $opts = $this->get_auth_options();
        return isset( $opts['api_key'] ) ? trim( (string) $opts['api_key'] ) : '';
    }

    public function get_credentials() {
        $runtime = $this->resolve_runtime_auth();
        return isset( $runtime['creds'] ) && is_array( $runtime['creds'] ) ? $runtime['creds'] : array();
    }

    public function is_oauth_connected() {
        $status = $this->get_oauth_status();
        return ! empty( $status['connected'] );
    }

    public function get_oauth_status() {
        $opts = $this->get_auth_options();
        $stored_access  = isset( $opts['oauth_access_token'] ) ? trim( (string) $opts['oauth_access_token'] ) : '';
        $stored_refresh = isset( $opts['oauth_refresh_token'] ) ? trim( (string) $opts['oauth_refresh_token'] ) : '';
        $stored_expires = isset( $opts['oauth_expires_at'] ) ? intval( $opts['oauth_expires_at'] ) : 0;
        if (
            ! defined( 'NIWA_AI_OAUTH_ACCESS_TOKEN' )
            && '' !== $stored_refresh
            && (
                '' === $stored_access
                || ( $stored_expires > 0 && $stored_expires <= ( time() + self::OAUTH_REFRESH_SKEW_SECONDS ) )
            )
        ) {
            $opts = $this->maybe_refresh_oauth_tokens( $opts );
        }

        $access = defined( 'NIWA_AI_OAUTH_ACCESS_TOKEN' ) ? trim( (string) NIWA_AI_OAUTH_ACCESS_TOKEN ) : ( isset( $opts['oauth_access_token'] ) ? trim( (string) $opts['oauth_access_token'] ) : '' );
        $refresh = defined( 'NIWA_AI_OAUTH_REFRESH_TOKEN' ) ? trim( (string) NIWA_AI_OAUTH_REFRESH_TOKEN ) : ( isset( $opts['oauth_refresh_token'] ) ? trim( (string) $opts['oauth_refresh_token'] ) : '' );
        $expires_at = isset( $opts['oauth_expires_at'] ) ? intval( $opts['oauth_expires_at'] ) : 0;
        $last_error = isset( $opts['oauth_last_error'] ) ? (string) $opts['oauth_last_error'] : '';
        $now = time();
        $expired = $expires_at > 0 && $expires_at <= $now;
        $has_refresh = '' !== $refresh;
        $connected = '' !== $access && ( ! $expired || ( $has_refresh && '' === $last_error ) );

        return array(
            'connected' => $connected,
            'access_token' => $access,
            'refresh_token' => $refresh,
            'expires_at' => $expires_at,
            'expired' => $expired,
            'expiring_soon' => $expires_at > 0 && $expires_at <= ( $now + self::OAUTH_STATUS_WARN_WINDOW_SECONDS ),
            'has_refresh_token' => $has_refresh,
            'email' => isset( $opts['oauth_email'] ) ? (string) $opts['oauth_email'] : '',
            'account_id' => isset( $opts['oauth_account_id'] ) ? (string) $opts['oauth_account_id'] : '',
            'last_error' => $last_error,
        );
    }

    /**
     * Persist an error message in the auth options so it is visible in the
     * admin OAuth panel after a redirect (not only as a transient admin notice).
     *
     * @param string $message Error message to store.
     */
    public function set_oauth_last_error( $message ) {
        $opts = $this->get_auth_options();
        $opts['oauth_last_error'] = sanitize_text_field( (string) $message );
        $this->save_auth_options( $opts );
    }

    public function clear_oauth_credentials() {
        $opts = $this->get_auth_options();
        $opts['oauth_access_token'] = '';
        $opts['oauth_refresh_token'] = '';
        $opts['oauth_expires_at'] = 0;
        $opts['oauth_account_id'] = '';
        $opts['oauth_email'] = '';
        $opts['oauth_last_error'] = '';
        if ( isset( $opts['mode'] ) && 'oauth' === $opts['mode'] ) {
            $opts['mode'] = 'api_key';
        }
        $this->save_auth_options( $opts );
        if ( function_exists( 'get_current_user_id' ) ) {
            $user_id = intval( get_current_user_id() );
            if ( $user_id > 0 ) {
                $this->delete_device_session( $user_id );
            }
        }
    }

    public function start_codex_device_flow( $user_id ) {
        $user_id = intval( $user_id );
        if ( $user_id <= 0 ) {
            return array(
                'success' => false,
                'error' => 'Invalid user context for OAuth.',
            );
        }

        $response = wp_remote_post(
            self::OAUTH_DEVICE_USERCODE_ENDPOINT,
            array(
                'timeout' => 20,
                'headers' => array(
                    'Accept' => 'application/json',
                    'Content-Type' => 'application/json',
                ),
                'body' => wp_json_encode(
                    array(
                        'client_id' => self::OAUTH_CLIENT_ID,
                        'scope' => self::OAUTH_SCOPE,
                    )
                ),
            )
        );
        if ( is_wp_error( $response ) ) {
            $err = 'Unable to start device authorization: ' . $response->get_error_message();
            error_log( '[Niwa AI] Codex device flow WP error: ' . $err );
            return array(
                'success' => false,
                'error' => $err,
            );
        }

        $status = wp_remote_retrieve_response_code( $response );
        $raw_body = (string) wp_remote_retrieve_body( $response );
        $body = json_decode( $raw_body, true );
        if ( $status < 200 || $status >= 300 || ! is_array( $body ) ) {
            return array(
                'success' => false,
                'error' => $this->build_oauth_error_message( $status, $body, 'Unable to start device authorization.' ),
            );
        }

        // Accept both OpenAI-proprietary field names and RFC 8628 standard names.
        $device_auth_id = '';
        if ( ! empty( $body['device_auth_id'] ) ) {
            $device_auth_id = trim( (string) $body['device_auth_id'] );
        } elseif ( ! empty( $body['device_code'] ) ) {
            $device_auth_id = trim( (string) $body['device_code'] );
        }
        $user_code = isset( $body['user_code'] ) ? trim( (string) $body['user_code'] ) : '';
        if ( '' === $device_auth_id || '' === $user_code ) {
            $missing = array();
            if ( '' === $device_auth_id ) {
                $missing[] = 'device_auth_id/device_code';
            }
            if ( '' === $user_code ) {
                $missing[] = 'user_code';
            }
            error_log( '[Niwa AI] Codex device flow incomplete response — missing: ' . implode( ', ', $missing ) . '. Keys present: ' . implode( ', ', array_keys( $body ) ) );
            return array(
                'success' => false,
                'error' => 'Device authorization response is incomplete (missing ' . implode( ', ', $missing ) . ').',
            );
        }

        $interval = isset( $body['interval'] ) ? intval( $body['interval'] ) : self::OAUTH_DEVICE_DEFAULT_INTERVAL;
        $interval = max( 1, min( self::OAUTH_DEVICE_MAX_INTERVAL, $interval ) );

        // Accept both "expires_at" (datetime string) and "expires_in" (seconds, RFC 8628).
        $expires_at = $this->parse_datetime_or_zero( isset( $body['expires_at'] ) ? $body['expires_at'] : '' );
        if ( $expires_at <= time() && ! empty( $body['expires_in'] ) ) {
            $expires_at = time() + max( 60, intval( $body['expires_in'] ) );
        }
        if ( $expires_at <= time() ) {
            $expires_at = time() + 900;
        }

        $session = array(
            'user_id' => $user_id,
            'device_auth_id' => $device_auth_id,
            'user_code' => $user_code,
            'interval' => $interval,
            'expires_at' => $expires_at,
            'created_at' => time(),
            'last_polled_at' => 0,
        );
        $ttl = max( 120, ( $expires_at - time() ) + 120 );
        $this->set_device_session( $user_id, $session, $ttl );

        $opts = $this->get_auth_options();
        $opts['oauth_last_error'] = '';
        $this->save_auth_options( $opts );

        return array(
            'success' => true,
            'session' => $session,
            'verification_url' => self::OAUTH_DEVICE_VERIFY_URL,
        );
    }

    public function get_codex_device_status( $user_id ) {
        $user_id = intval( $user_id );
        if ( $user_id <= 0 ) {
            return array( 'active' => false );
        }

        $session = $this->get_device_session( $user_id );
        if ( false === $session ) {
            return array( 'active' => false );
        }

        $expires_at = isset( $session['expires_at'] ) ? intval( $session['expires_at'] ) : 0;
        if ( $expires_at > 0 && $expires_at <= time() ) {
            $this->delete_device_session( $user_id );
            return array( 'active' => false );
        }

        return array(
            'active' => true,
            'user_code' => isset( $session['user_code'] ) ? (string) $session['user_code'] : '',
            'interval' => isset( $session['interval'] ) ? max( 1, intval( $session['interval'] ) ) : self::OAUTH_DEVICE_DEFAULT_INTERVAL,
            'expires_at' => $expires_at,
            'verification_url' => self::OAUTH_DEVICE_VERIFY_URL,
        );
    }

    public function poll_codex_device_flow( $user_id ) {
        $user_id = intval( $user_id );
        if ( $user_id <= 0 ) {
            return array(
                'success' => false,
                'pending' => false,
                'error' => 'Invalid user context for OAuth polling.',
            );
        }

        $session = $this->get_device_session( $user_id );
        if ( false === $session ) {
            return array(
                'success' => false,
                'pending' => false,
                'error' => 'No active device authorization session. Start OAuth again.',
            );
        }

        $expires_at = isset( $session['expires_at'] ) ? intval( $session['expires_at'] ) : 0;
        if ( $expires_at > 0 && $expires_at <= time() ) {
            $this->delete_device_session( $user_id );
            return array(
                'success' => false,
                'pending' => false,
                'error' => 'Device authorization code expired. Start OAuth again.',
            );
        }

        $interval = isset( $session['interval'] ) ? max( 1, intval( $session['interval'] ) ) : self::OAUTH_DEVICE_DEFAULT_INTERVAL;
        $now = time();
        $last_polled_at = isset( $session['last_polled_at'] ) ? intval( $session['last_polled_at'] ) : 0;
        if ( $last_polled_at > 0 && ( $now - $last_polled_at ) < $interval ) {
            $retry_after = max( 1, $interval - ( $now - $last_polled_at ) );
            return array(
                'success' => false,
                'pending' => true,
                'retry_after' => $retry_after,
                'message' => 'Waiting for authorization in browser.',
            );
        }

        $session['last_polled_at'] = $now;
        $this->set_device_session( $user_id, $session, max( 120, ( $expires_at - $now ) + 120 ) );

        $device_auth_id = isset( $session['device_auth_id'] ) ? trim( (string) $session['device_auth_id'] ) : '';
        $user_code = isset( $session['user_code'] ) ? trim( (string) $session['user_code'] ) : '';
        if ( '' === $device_auth_id || '' === $user_code ) {
            $this->delete_device_session( $user_id );
            return array(
                'success' => false,
                'pending' => false,
                'error' => 'Device authorization session is invalid. Start OAuth again.',
            );
        }

        $response = wp_remote_post(
            self::OAUTH_DEVICE_TOKEN_ENDPOINT,
            array(
                'timeout' => 20,
                'headers' => array(
                    'Accept' => 'application/json',
                    'Content-Type' => 'application/json',
                ),
                'body' => wp_json_encode(
                    array(
                        'device_auth_id' => $device_auth_id,
                        'user_code' => $user_code,
                    )
                ),
            )
        );
        if ( is_wp_error( $response ) ) {
            return array(
                'success' => false,
                'pending' => true,
                'retry_after' => $interval,
                'message' => 'Waiting for authorization in browser.',
            );
        }

        $status = wp_remote_retrieve_response_code( $response );
        $body = json_decode( (string) wp_remote_retrieve_body( $response ), true );
        if ( $status >= 200 && $status < 300 && is_array( $body ) ) {
            $authorization_code = isset( $body['authorization_code'] ) ? trim( (string) $body['authorization_code'] ) : '';
            $code_verifier = isset( $body['code_verifier'] ) ? trim( (string) $body['code_verifier'] ) : '';
            if ( '' === $authorization_code || '' === $code_verifier ) {
                $this->delete_device_session( $user_id );
                return array(
                    'success' => false,
                    'pending' => false,
                    'error' => 'Device authorization returned incomplete token data. Start OAuth again.',
                );
            }

            $exchange = $this->request_oauth_tokens(
                array(
                    'grant_type' => 'authorization_code',
                    'code' => $authorization_code,
                    'redirect_uri' => self::OAUTH_DEVICE_CALLBACK_URI,
                    'client_id' => self::OAUTH_CLIENT_ID,
                    'code_verifier' => $code_verifier,
                )
            );
            if ( empty( $exchange['success'] ) ) {
                $this->delete_device_session( $user_id );
                return array(
                    'success' => false,
                    'pending' => false,
                    'error' => isset( $exchange['error'] ) ? (string) $exchange['error'] : 'OAuth token exchange failed.',
                );
            }

            $tokens = isset( $exchange['tokens'] ) && is_array( $exchange['tokens'] ) ? $exchange['tokens'] : array();
            $this->persist_oauth_tokens( $tokens, true );
            $this->mark_oauth_mode_connected();
            $this->delete_device_session( $user_id );
            return array(
                'success' => true,
                'pending' => false,
                'tokens' => $tokens,
            );
        }

        $error_code = '';
        if ( is_array( $body ) && ! empty( $body['error'] ) && is_array( $body['error'] ) && ! empty( $body['error']['code'] ) ) {
            $error_code = sanitize_key( (string) $body['error']['code'] );
        }

        $pending_codes = array(
            'deviceauth_authorization_pending',
            'deviceauth_authorization_unknown',
            'authorization_pending',
            'slow_down',
        );
        if ( $expires_at > $now && ( $status === 403 || in_array( $error_code, $pending_codes, true ) ) ) {
            return array(
                'success' => false,
                'pending' => true,
                'retry_after' => $interval,
                'message' => 'Waiting for authorization in browser.',
            );
        }

        $this->delete_device_session( $user_id );
        $message = $this->build_oauth_error_message( $status, $body, 'Device authorization failed. Start OAuth again.' );
        $opts = $this->get_auth_options();
        $opts['oauth_last_error'] = sanitize_text_field( $message );
        $this->save_auth_options( $opts );
        return array(
            'success' => false,
            'pending' => false,
            'error' => $message,
        );
    }

    public function get_codex_oauth_authorize_url( $redirect_uri, $user_id ) {
        $redirect_uri = trim( (string) $redirect_uri );
        if ( '' === $redirect_uri ) {
            return new WP_Error( 'niwa_oauth_redirect_uri_missing', 'OAuth redirect URI is missing.' );
        }

        try {
            $state = bin2hex( random_bytes( 24 ) );
            $code_verifier = $this->base64url_encode( random_bytes( 64 ) );
        } catch ( Exception $e ) {
            return new WP_Error( 'niwa_oauth_entropy_error', 'Unable to initialize OAuth state securely.' );
        }

        $challenge_bytes = hash( 'sha256', $code_verifier, true );
        $code_challenge = $this->base64url_encode( $challenge_bytes );

        $state_key = $this->state_transient_key( $state );
        Niwa_AI_Durable_Cache::set(
            $state_key,
            array(
                'code_verifier' => $code_verifier,
                'redirect_uri' => $redirect_uri,
                'user_id' => intval( $user_id ),
                'issued_at' => time(),
            ),
            self::OAUTH_STATE_TTL_SECONDS
        );

        $params = array(
            'response_type' => 'code',
            'client_id' => self::OAUTH_CLIENT_ID,
            'redirect_uri' => $redirect_uri,
            'scope' => self::OAUTH_SCOPE,
            'code_challenge' => $code_challenge,
            'code_challenge_method' => 'S256',
            'state' => $state,
            'id_token_add_organizations' => 'true',
            'codex_cli_simplified_flow' => 'true',
            'originator' => self::OAUTH_ORIGINATOR,
        );

        return self::OAUTH_ISSUER . '/oauth/authorize?' . http_build_query( $params, '', '&', PHP_QUERY_RFC3986 );
    }

    public function complete_codex_oauth( $args ) {
        $code = isset( $args['code'] ) ? trim( (string) $args['code'] ) : '';
        $state = isset( $args['state'] ) ? trim( (string) $args['state'] ) : '';
        $redirect_uri = isset( $args['redirect_uri'] ) ? trim( (string) $args['redirect_uri'] ) : '';
        $user_id = isset( $args['user_id'] ) ? intval( $args['user_id'] ) : 0;

        if ( '' === $code || '' === $state || '' === $redirect_uri ) {
            return array(
                'success' => false,
                'error' => 'Missing OAuth callback parameters.',
            );
        }

        $state_payload = Niwa_AI_Durable_Cache::get( $this->state_transient_key( $state ) );
        Niwa_AI_Durable_Cache::delete( $this->state_transient_key( $state ) );

        if ( empty( $state_payload ) || ! is_array( $state_payload ) ) {
            return array(
                'success' => false,
                'error' => 'OAuth session expired or invalid state.',
            );
        }

        $expected_user_id = isset( $state_payload['user_id'] ) ? intval( $state_payload['user_id'] ) : 0;
        if ( $expected_user_id > 0 && $user_id > 0 && $expected_user_id !== $user_id ) {
            return array(
                'success' => false,
                'error' => 'OAuth callback user mismatch.',
            );
        }

        $code_verifier = isset( $state_payload['code_verifier'] ) ? trim( (string) $state_payload['code_verifier'] ) : '';
        if ( '' === $code_verifier ) {
            return array(
                'success' => false,
                'error' => 'OAuth verifier missing.',
            );
        }

        $expected_redirect_uri = isset( $state_payload['redirect_uri'] ) ? trim( (string) $state_payload['redirect_uri'] ) : '';
        if ( '' === $expected_redirect_uri || $expected_redirect_uri !== $redirect_uri ) {
            return array(
                'success' => false,
                'error' => 'OAuth redirect URI mismatch.',
            );
        }

        $exchange = $this->request_oauth_tokens(
            array(
                'grant_type' => 'authorization_code',
                'code' => $code,
                'redirect_uri' => $redirect_uri,
                'client_id' => self::OAUTH_CLIENT_ID,
                'code_verifier' => $code_verifier,
            )
        );
        if ( empty( $exchange['success'] ) ) {
            return $exchange;
        }

        $tokens = isset( $exchange['tokens'] ) && is_array( $exchange['tokens'] ) ? $exchange['tokens'] : array();
        $this->persist_oauth_tokens( $tokens, true );
        $this->mark_oauth_mode_connected();

        return array(
            'success' => true,
            'tokens' => $tokens,
        );
    }

    private function maybe_refresh_oauth_tokens( $opts ) {
        if ( defined( 'NIWA_AI_OAUTH_ACCESS_TOKEN' ) ) {
            return $opts;
        }

        $refresh = isset( $opts['oauth_refresh_token'] ) ? trim( (string) $opts['oauth_refresh_token'] ) : '';
        if ( '' === $refresh ) {
            return $opts;
        }

        $access = isset( $opts['oauth_access_token'] ) ? trim( (string) $opts['oauth_access_token'] ) : '';
        $expires_at = isset( $opts['oauth_expires_at'] ) ? intval( $opts['oauth_expires_at'] ) : 0;
        if ( '' !== $access && $expires_at > ( time() + self::OAUTH_REFRESH_SKEW_SECONDS ) ) {
            return $opts;
        }

        $lock_token = $this->acquire_refresh_lock();
        if ( '' === $lock_token ) {
            return $this->get_auth_options();
        }

        $refresh_result = $this->request_oauth_tokens(
            array(
                'grant_type' => 'refresh_token',
                'refresh_token' => $refresh,
                'client_id' => self::OAUTH_CLIENT_ID,
                'scope' => self::OAUTH_SCOPE,
            )
        );
        try {
            if ( empty( $refresh_result['success'] ) ) {
                $latest_opts    = $this->reload_auth_options_from_storage();
                $latest_access  = isset( $latest_opts['oauth_access_token'] ) ? trim( (string) $latest_opts['oauth_access_token'] ) : '';
                $latest_expires = isset( $latest_opts['oauth_expires_at'] ) ? intval( $latest_opts['oauth_expires_at'] ) : 0;
                if ( '' !== $latest_access && $latest_expires > ( time() + 5 ) ) {
                    return $latest_opts;
                }

                $opts['oauth_last_error'] = isset( $refresh_result['error'] ) ? (string) $refresh_result['error'] : 'OAuth refresh failed.';
                $this->save_auth_options( $opts );
                return $opts;
            }

            $tokens = isset( $refresh_result['tokens'] ) && is_array( $refresh_result['tokens'] ) ? $refresh_result['tokens'] : array();

            // Refresh responses may omit identity claims. Preserve working values.
            if ( empty( $tokens['oauth_refresh_token'] ) ) {
                $tokens['oauth_refresh_token'] = $refresh;
            }
            if ( empty( $tokens['oauth_account_id'] ) && ! empty( $opts['oauth_account_id'] ) ) {
                $tokens['oauth_account_id'] = $opts['oauth_account_id'];
            }
            if ( empty( $tokens['oauth_email'] ) && ! empty( $opts['oauth_email'] ) ) {
                $tokens['oauth_email'] = $opts['oauth_email'];
            }

            $this->persist_oauth_tokens( $tokens, false );
        } finally {
            $this->release_refresh_lock( $lock_token );
        }

        return $this->get_auth_options();
    }

    private function request_oauth_tokens( $fields ) {
        $response = wp_remote_post(
            self::OAUTH_ISSUER . '/oauth/token',
            array(
                'timeout' => 30,
                'headers' => array(
                    'Accept' => 'application/json',
                    'Content-Type' => 'application/x-www-form-urlencoded',
                ),
                'body' => $fields,
            )
        );

        if ( is_wp_error( $response ) ) {
            return array(
                'success' => false,
                'error' => 'OAuth token request failed: ' . $response->get_error_message(),
            );
        }

        $status = wp_remote_retrieve_response_code( $response );
        $body_raw = (string) wp_remote_retrieve_body( $response );
        $body = json_decode( $body_raw, true );

        if ( $status < 200 || $status >= 300 ) {
            $message = 'OAuth token endpoint returned HTTP ' . $status . '.';
            if ( is_array( $body ) ) {
                if ( ! empty( $body['error_description'] ) ) {
                    $message .= ' ' . sanitize_text_field( (string) $body['error_description'] );
                } elseif ( ! empty( $body['error'] ) && is_string( $body['error'] ) ) {
                    $message .= ' ' . sanitize_text_field( (string) $body['error'] );
                }
            }
            return array(
                'success' => false,
                'error' => $message,
            );
        }

        if ( ! is_array( $body ) || empty( $body['access_token'] ) ) {
            return array(
                'success' => false,
                'error' => 'OAuth token response is invalid.',
            );
        }

        $access_token = trim( (string) $body['access_token'] );
        $refresh_token = isset( $body['refresh_token'] ) ? trim( (string) $body['refresh_token'] ) : '';
        $expires_in = isset( $body['expires_in'] ) ? intval( $body['expires_in'] ) : 0;
        $expires_at = $expires_in > 0 ? max( 0, time() + $expires_in - 300 ) : 0;

        $id_claims = array();
        if ( ! empty( $body['id_token'] ) && is_string( $body['id_token'] ) ) {
            $id_claims = $this->decode_jwt_claims( $body['id_token'] );
        }
        $access_claims = $this->decode_jwt_claims( $access_token );
        $claims = is_array( $id_claims ) ? $id_claims : array();
        if ( is_array( $access_claims ) ) {
            $claims = array_merge( $access_claims, $claims );
        }

        $account_id = $this->extract_account_id_from_claims( $claims );
        $email = $this->extract_email_from_claims( $claims );

        return array(
            'success' => true,
            'tokens' => array(
                'oauth_access_token' => $access_token,
                'oauth_refresh_token' => $refresh_token,
                'oauth_expires_at' => $expires_at,
                'oauth_account_id' => $account_id,
                'oauth_email' => $email,
                'oauth_last_error' => '',
            ),
        );
    }

    private function persist_oauth_tokens( $tokens, $force_oauth_mode ) {
        $opts = $this->get_auth_options();
        if ( ! empty( $tokens['oauth_access_token'] ) ) {
            $opts['oauth_access_token'] = trim( (string) $tokens['oauth_access_token'] );
        }
        if ( array_key_exists( 'oauth_refresh_token', $tokens ) ) {
            $refresh = trim( (string) $tokens['oauth_refresh_token'] );
            if ( '' !== $refresh ) {
                $opts['oauth_refresh_token'] = $refresh;
            } elseif ( ! isset( $opts['oauth_refresh_token'] ) ) {
                $opts['oauth_refresh_token'] = '';
            }
        }
        if ( isset( $tokens['oauth_expires_at'] ) ) {
            $opts['oauth_expires_at'] = intval( $tokens['oauth_expires_at'] );
        }
        if ( isset( $tokens['oauth_account_id'] ) ) {
            $opts['oauth_account_id'] = sanitize_text_field( (string) $tokens['oauth_account_id'] );
        }
        if ( isset( $tokens['oauth_email'] ) ) {
            $opts['oauth_email'] = sanitize_email( (string) $tokens['oauth_email'] );
        }
        $opts['oauth_last_error'] = isset( $tokens['oauth_last_error'] ) ? sanitize_text_field( (string) $tokens['oauth_last_error'] ) : '';
        if ( $force_oauth_mode ) {
            $opts['mode'] = 'oauth';
        }
        $this->save_auth_options( $opts );
    }

    private function extract_email_from_claims( $claims ) {
        if ( ! is_array( $claims ) ) {
            return '';
        }
        if ( ! empty( $claims['email'] ) ) {
            return sanitize_email( (string) $claims['email'] );
        }
        return '';
    }

    private function extract_account_id_from_claims( $claims ) {
        if ( ! is_array( $claims ) ) {
            return '';
        }

        if ( ! empty( $claims['chatgpt_account_id'] ) ) {
            return sanitize_text_field( (string) $claims['chatgpt_account_id'] );
        }
        if ( ! empty( $claims['https://api.openai.com/auth.chatgpt_account_id'] ) ) {
            return sanitize_text_field( (string) $claims['https://api.openai.com/auth.chatgpt_account_id'] );
        }
        if ( ! empty( $claims['https://api.openai.com/auth'] ) && is_array( $claims['https://api.openai.com/auth'] ) ) {
            $auth_claim = $claims['https://api.openai.com/auth'];
            if ( ! empty( $auth_claim['chatgpt_account_id'] ) ) {
                return sanitize_text_field( (string) $auth_claim['chatgpt_account_id'] );
            }
        }
        if ( ! empty( $claims['organizations'] ) && is_array( $claims['organizations'] ) ) {
            foreach ( $claims['organizations'] as $org ) {
                if ( is_array( $org ) && ! empty( $org['id'] ) ) {
                    return sanitize_text_field( (string) $org['id'] );
                }
            }
        }

        return '';
    }

    private function decode_jwt_claims( $token ) {
        $token = trim( (string) $token );
        if ( '' === $token ) {
            return array();
        }

        $parts = explode( '.', $token );
        if ( count( $parts ) < 2 ) {
            return array();
        }

        $payload = $parts[1];
        $payload = strtr( $payload, '-_', '+/' );
        $pad_len = strlen( $payload ) % 4;
        if ( $pad_len > 0 ) {
            $payload .= str_repeat( '=', 4 - $pad_len );
        }
        $decoded = base64_decode( $payload, true );
        if ( false === $decoded ) {
            return array();
        }

        $claims = json_decode( $decoded, true );
        return is_array( $claims ) ? $claims : array();
    }

    private function state_transient_key( $state ) {
        return self::OAUTH_STATE_TRANSIENT_PREFIX . hash( 'sha256', (string) $state );
    }

    private function base64url_encode( $raw ) {
        return rtrim( strtr( base64_encode( $raw ), '+/', '-_' ), '=' );
    }

    private function device_option_key( $user_id ) {
        return self::OAUTH_DEVICE_TRANSIENT_PREFIX . intval( $user_id );
    }

    /**
     * Save device flow session directly to wp_options (bypasses object cache).
     *
     * Transients rely on the object cache which can be unreliable on hosts
     * with LiteSpeed LSMCD, misconfigured Redis, or similar setups. Device
     * flow sessions are critical and short-lived, so we store them as a
     * regular option with an embedded expiry timestamp.
     */
    private function set_device_session( $user_id, $session, $ttl ) {
        $session['_expires_at_unix'] = time() + max( 60, intval( $ttl ) );
        update_option( $this->device_option_key( $user_id ), $session, false );
    }

    /**
     * Read device flow session from wp_options, respecting embedded expiry.
     */
    private function get_device_session( $user_id ) {
        $session = get_option( $this->device_option_key( $user_id ), false );
        if ( ! is_array( $session ) ) {
            return false;
        }

        // Honour embedded expiry.
        $hard_expiry = isset( $session['_expires_at_unix'] ) ? intval( $session['_expires_at_unix'] ) : 0;
        if ( $hard_expiry > 0 && $hard_expiry <= time() ) {
            $this->delete_device_session( $user_id );
            return false;
        }

        return $session;
    }

    /**
     * Delete device flow session from wp_options.
     */
    private function delete_device_session( $user_id ) {
        delete_option( $this->device_option_key( $user_id ) );
    }

    private function parse_datetime_or_zero( $value ) {
        if ( ! is_string( $value ) ) {
            return 0;
        }
        $value = trim( $value );
        if ( '' === $value ) {
            return 0;
        }
        $ts = strtotime( $value );
        return false === $ts ? 0 : intval( $ts );
    }

    private function build_oauth_error_message( $status, $body, $fallback ) {
        $message = (string) $fallback;
        if ( intval( $status ) > 0 ) {
            $message .= ' HTTP ' . intval( $status ) . '.';
        }

        if ( is_array( $body ) ) {
            if ( ! empty( $body['error_description'] ) && is_string( $body['error_description'] ) ) {
                $message .= ' ' . sanitize_text_field( $body['error_description'] );
            } elseif ( ! empty( $body['error']['message'] ) && is_string( $body['error']['message'] ) ) {
                $message .= ' ' . sanitize_text_field( $body['error']['message'] );
            } elseif ( ! empty( $body['error'] ) && is_string( $body['error'] ) ) {
                $message .= ' ' . sanitize_text_field( $body['error'] );
            } elseif ( ! empty( $body['message'] ) && is_string( $body['message'] ) ) {
                $message .= ' ' . sanitize_text_field( $body['message'] );
            }
        }

        return trim( $message );
    }

    private function get_auth_options() {
        if ( null === $this->cached_auth_opts ) {
            $this->cached_auth_opts = $this->load_auth_options_from_db();
        }
        return $this->cached_auth_opts;
    }

    private function save_auth_options( $opts ) {
        $opts = is_array( $opts ) ? $opts : array();
        update_option( self::OPTION_KEY, $opts, false );
        wp_cache_delete( self::OPTION_KEY, 'options' );
        wp_cache_delete( 'alloptions', 'options' );
        $this->cached_auth_opts = $opts;
        $this->write_auth_cache_for_all_lanes( $opts );
    }

    private function mark_oauth_mode_connected() {
        $opts                     = $this->get_auth_options();
        $opts['mode']             = 'oauth';
        $opts['oauth_last_error'] = '';
        $this->save_auth_options( $opts );
    }

    private function reload_auth_options_from_storage() {
        $opts = $this->load_auth_options_from_db();
        $this->cached_auth_opts = $opts;
        $this->write_auth_cache_for_all_lanes( $opts );
        return $opts;
    }

    private function write_auth_cache_for_all_lanes( $opts ) {
        foreach ( array( 'default', 'frontend', 'founder' ) as $lane ) {
            delete_transient( self::AUTH_CACHE_TRANSIENT_PREFIX . $lane );
        }
    }

    private function load_auth_options_from_db() {
        global $wpdb;

        if ( isset( $wpdb ) && $wpdb instanceof wpdb ) {
            $raw = $wpdb->get_var(
                $wpdb->prepare(
                    "SELECT option_value FROM {$wpdb->options} WHERE option_name = %s LIMIT 1",
                    self::OPTION_KEY
                )
            );

            if ( null !== $raw ) {
                $decoded = maybe_unserialize( $raw );
                return is_array( $decoded ) ? $decoded : array();
            }
        }

        $opts = get_option( self::OPTION_KEY, array() );
        return is_array( $opts ) ? $opts : array();
    }

    private function determine_preferred_mode( $opts ) {
        $opts = is_array( $opts ) ? $opts : array();
        $mode = 'api_key';

        if ( defined( 'NIWA_AI_AUTH_MODE' ) && in_array( NIWA_AI_AUTH_MODE, array( 'api_key', 'oauth' ), true ) ) {
            $mode = NIWA_AI_AUTH_MODE;
        } else {
            $lane_override = $this->get_lane_mode_override();
            if ( '' !== $lane_override ) {
                $mode = $lane_override;
            } else {
                $mode = isset( $opts['mode'] ) ? $opts['mode'] : 'api_key';
            }
        }

        return in_array( $mode, array( 'api_key', 'oauth' ), true ) ? $mode : 'api_key';
    }

    private function resolve_runtime_auth() {
        $opts           = $this->get_auth_options();
        $preferred_mode = $this->determine_preferred_mode( $opts );
        $opts           = is_array( $opts ) ? $opts : array();
        $api_key        = defined( 'NIWA_AI_OPENAI_API_KEY' )
            ? trim( (string) NIWA_AI_OPENAI_API_KEY )
            : ( isset( $opts['api_key'] ) ? trim( (string) $opts['api_key'] ) : '' );

        $oauth_access = defined( 'NIWA_AI_OAUTH_ACCESS_TOKEN' )
            ? trim( (string) NIWA_AI_OAUTH_ACCESS_TOKEN )
            : ( isset( $opts['oauth_access_token'] ) ? trim( (string) $opts['oauth_access_token'] ) : '' );
        $oauth_refresh = defined( 'NIWA_AI_OAUTH_REFRESH_TOKEN' )
            ? trim( (string) NIWA_AI_OAUTH_REFRESH_TOKEN )
            : ( isset( $opts['oauth_refresh_token'] ) ? trim( (string) $opts['oauth_refresh_token'] ) : '' );
        $oauth_expires = isset( $opts['oauth_expires_at'] ) ? intval( $opts['oauth_expires_at'] ) : 0;
        $oauth_expiring = $oauth_expires > 0 && $oauth_expires <= ( time() + self::OAUTH_REFRESH_SKEW_SECONDS );

        $should_try_refresh = ! defined( 'NIWA_AI_OAUTH_ACCESS_TOKEN' )
            && '' !== $oauth_refresh
            && ( '' === $oauth_access || $oauth_expiring )
            && (
                'oauth' === $preferred_mode
                || ( 'api_key' === $preferred_mode && '' === $api_key )
            );

        if ( $should_try_refresh ) {
            $opts         = $this->maybe_refresh_oauth_tokens( $opts );
            $oauth_access = isset( $opts['oauth_access_token'] ) ? trim( (string) $opts['oauth_access_token'] ) : '';
            $oauth_refresh = isset( $opts['oauth_refresh_token'] ) ? trim( (string) $opts['oauth_refresh_token'] ) : '';
        }

        $oauth_creds = array(
            'access_token'  => $oauth_access,
            'refresh_token' => $oauth_refresh,
            'account_id'    => isset( $opts['oauth_account_id'] ) ? trim( (string) $opts['oauth_account_id'] ) : '',
        );

        if ( 'oauth' === $preferred_mode ) {
            if ( '' !== $oauth_access ) {
                return array(
                    'mode'  => 'oauth',
                    'creds' => $oauth_creds,
                );
            }

            if ( '' !== $api_key ) {
                return array(
                    'mode'  => 'api_key',
                    'creds' => array(
                        'api_key' => $api_key,
                    ),
                );
            }

            return array(
                'mode'  => 'oauth',
                'creds' => $oauth_creds,
            );
        }

        if ( '' !== $api_key ) {
            return array(
                'mode'  => 'api_key',
                'creds' => array(
                    'api_key' => $api_key,
                ),
            );
        }

        if ( '' !== $oauth_access ) {
            return array(
                'mode'  => 'oauth',
                'creds' => $oauth_creds,
            );
        }

        return array(
            'mode'  => 'api_key',
            'creds' => array(
                'api_key' => '',
            ),
        );
    }

    private function acquire_refresh_lock() {
        $deadline = time() + self::OAUTH_REFRESH_LOCK_WAIT_SECONDS;
        $lock_key = $this->refresh_lock_option_key();

        do {
            $token = time() . ':' . wp_generate_password( 12, false, false );
            if ( add_option( $lock_key, $token, '', false ) ) {
                return $token;
            }

            $existing = get_option( $lock_key, '' );
            if ( $this->is_refresh_lock_stale( $existing ) ) {
                delete_option( $lock_key );
                continue;
            }

            usleep( 250000 );
        } while ( time() < $deadline );

        return '';
    }

    private function release_refresh_lock( $token ) {
        $lock_key = $this->refresh_lock_option_key();
        $current = get_option( $lock_key, '' );
        if ( is_string( $current ) && hash_equals( $current, (string) $token ) ) {
            delete_option( $lock_key );
        }
    }

    private function is_refresh_lock_stale( $value ) {
        if ( ! is_string( $value ) || '' === $value ) {
            return true;
        }

        $parts = explode( ':', $value, 2 );
        $issued_at = isset( $parts[0] ) ? intval( $parts[0] ) : 0;
        if ( $issued_at <= 0 ) {
            return true;
        }

        return $issued_at <= ( time() - self::OAUTH_REFRESH_LOCK_TTL_SECONDS );
    }

    private function sanitize_lane( $lane ) {
        $lane = sanitize_key( (string) $lane );
        return '' !== $lane ? $lane : 'default';
    }

    private function get_lane_mode_override() {
        if ( 'founder' !== $this->lane ) {
            return '';
        }

        $opts = get_option( Niwa_AI_Constants::OPT_AUTONOMY, array() );
        if ( ! is_array( $opts ) ) {
            return '';
        }

        $mode = isset( $opts['founder_auth_mode'] ) ? sanitize_key( (string) $opts['founder_auth_mode'] ) : '';
        return in_array( $mode, array( 'api_key', 'oauth' ), true ) ? $mode : '';
    }

    private function refresh_lock_option_key() {
        if ( 'default' === $this->lane ) {
            return self::OAUTH_REFRESH_LOCK_OPTION;
        }

        return self::OAUTH_REFRESH_LOCK_OPTION . '_' . $this->lane;
    }

}
