<?php
/**
 * HuggingFace Provider
 * AI models via HuggingFace Router API
 */

if ( ! defined( 'ABSPATH' ) ) exit; // Thoát nếu truy cập trực tiếp

class AIAutotool_AI_Provider_HuggingFace extends AIAutotool_AI_Provider_Base {
    
    protected $name = 'HuggingFace';
    protected $id = 'huggingface';
    protected $base_url = 'https://router.huggingface.co/v1';
    protected $default_model = 'google/gemma-2-2b-it:nebius';  # Google Gemma 2 2B IT - Fast & Efficient
    protected $capabilities = array('text_generation', 'image_generation');
    
    protected $available_models = array(
        // Fast & Efficient models
        'google/gemma-2-2b-it:nebius' => 'Google Gemma 2 2B IT - Fast & Efficient',
        'meta-llama/Meta-Llama-3.1-8B-Instruct' => 'Meta Llama 3.1 8B Instruct - Fast & Lightweight',
        
        // Reasoning models
        'deepseek-ai/DeepSeek-R1' => 'DeepSeek R1 - Powerful Reasoning Model',
        'Qwen/Qwen3-4B-Thinking-2507' => 'Qwen 3 4B Thinking - Powerful Small Model with Reasoning',
        
        // Advanced models
        'context-labs/google-gemma-3-27b-it' => 'Google Gemma 3 27B IT - Advanced Model',
        'Qwen/Qwen2.5-7B-Instruct-1M' => 'Qwen 2.5 7B Instruct 1M - Strong Conversational Model with Long Context',
        
        // Coding models
        'Qwen/Qwen3-Coder-480B-A35B-Instruct' => 'Qwen 3 Coder 480B A35B - Powerful Coding Model',
        'Qwen/Qwen2.5-Coder-32B-Instruct' => 'Qwen 2.5 Coder 32B - Text Generation Model for Code',
        
        // Other models
        'openai/gpt-oss-120b' => 'OpenAI GPT-OSS 120B - Great Text Generation with Tool Calling',
        'zai-org/GLM-4.5' => 'ZAI GLM 4.5 - Powerful Text Generation Model',
        
        // Aliases for easier access
        'gemma-2-2b' => 'google/gemma-2-2b-it:nebius',
        'gemma-2' => 'google/gemma-2-2b-it:nebius',
        'llama-3.1-8b' => 'meta-llama/Meta-Llama-3.1-8B-Instruct',
        'llama-3.1' => 'meta-llama/Meta-Llama-3.1-8B-Instruct',
        'deepseek-r1' => 'deepseek-ai/DeepSeek-R1',
        'gemma-3-27b' => 'context-labs/google-gemma-3-27b-it',
        'gemma-3' => 'context-labs/google-gemma-3-27b-it',
        'qwen-3-coder' => 'Qwen/Qwen3-Coder-480B-A35B-Instruct',
        'qwen-2.5-coder' => 'Qwen/Qwen2.5-Coder-32B-Instruct',
        'qwen-3-thinking' => 'Qwen/Qwen3-4B-Thinking-2507',
        'qwen-2.5-7b' => 'Qwen/Qwen2.5-7B-Instruct-1M',
        'gpt-oss-120b' => 'openai/gpt-oss-120b',
        'glm-4.5' => 'zai-org/GLM-4.5',
    );

    public function get_image_models() {
        return array(
            'black-forest-labs/FLUX.1-schnell' => 'FLUX.1 Schnell',
            'stabilityai/stable-diffusion-xl-base-1.0' => 'Stable Diffusion XL',
            'runwayml/stable-diffusion-v1-5' => 'Stable Diffusion 1.5',
        );
    }

    public function generate_image($prompt, $options = array()) {
        if (!$this->is_configured()) {
            return array('success' => false, 'error' => 'HuggingFace API key not configured');
        }

        $prompt = trim((string) $prompt);
        if ($prompt === '') {
            return array('success' => false, 'error' => 'Image prompt is required');
        }

        $models = $this->get_image_models();
        $model = isset($options['model']) ? sanitize_text_field($options['model']) : 'black-forest-labs/FLUX.1-schnell';
        if (!isset($models[$model])) {
            $model = 'black-forest-labs/FLUX.1-schnell';
        }

        $api_key_to_use = !empty($this->api_key) ? $this->api_key : $this->get_rotated_key();
        if (empty($api_key_to_use)) {
            return array('success' => false, 'error' => 'HuggingFace API key not configured');
        }

        $encoded_model = implode('/', array_map('rawurlencode', explode('/', $model)));

        $response = $this->make_request(
            'https://api-inference.huggingface.co/models/' . $encoded_model,
            array(
                'inputs' => $prompt,
                'parameters' => array(
                    'negative_prompt' => !empty($options['negative_prompt']) ? sanitize_text_field($options['negative_prompt']) : 'worst quality, blurry',
                ),
            ),
            'POST',
            array(
                'Authorization' => 'Bearer ' . $api_key_to_use,
                'Content-Type' => 'application/json',
            ),
            array('timeout' => isset($options['timeout']) ? (int) $options['timeout'] : 90)
        );

        if (!$response['success']) {
            return $response;
        }

        $data = $response['data'];
        if (is_string($data) && $data !== '') {
            return array(
                'success' => true,
                'data_url' => 'data:image/png;base64,' . base64_encode($data),
                'model' => $model,
                'provider' => $this->id,
            );
        }

        return array('success' => false, 'error' => 'HuggingFace image response missing image data');
    }
    
    /**
     * Model mapping for invalid/old model names to valid ones
     */
    protected $model_mapping = array(
        // Map common model names to HuggingFace models
        'gpt-4' => 'openai/gpt-oss-120b',
        'gpt-4o' => 'openai/gpt-oss-120b',
        'gpt-3.5-turbo' => 'google/gemma-2-2b-it:nebius',
        'gpt-4o-mini' => 'google/gemma-2-2b-it:nebius',
        'claude' => 'deepseek-ai/DeepSeek-R1',
        'gemini' => 'google/gemma-2-2b-it:nebius',
        'gemini-2' => 'google/gemma-2-2b-it:nebius',
        'gemini-3' => 'context-labs/google-gemma-3-27b-it',
    );
    
    /**
     * Models that don't support max_tokens parameter
     */
    protected $models_without_max_tokens = array(
        // Some models may not support max_tokens, add them here if needed
    );
    
    public function __construct($api_key = '') {
        parent::__construct($api_key);
        
        if (empty($this->api_key)) {
            $this->api_key = $this->get_rotated_key();
        }
        
        // HuggingFace default API key
        if (empty($this->api_key)) {
            $this->api_key = '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';
        }
    }
    
    /**
     * Map model name to valid HuggingFace model
     */
    protected function map_model($model) {
        // Check if model needs mapping
        if (isset($this->model_mapping[$model])) {
            $mapped_model = $this->model_mapping[$model];
            return $mapped_model;
        }
        
        // Check if model is an alias
        if (isset($this->available_models[$model]) && strpos($this->available_models[$model], '/') !== false) {
            // This is already a full model path
            return $model;
        }
        
        // Check if model is in available models (could be alias)
        if (isset($this->available_models[$model])) {
            $mapped = $this->available_models[$model];
            // If the value is a full model path, use it; otherwise use the key
            if (strpos($mapped, '/') !== false) {
                return $mapped;
            }
            return $model;
        }
        
        // If model contains /, assume it's already a valid HuggingFace model path
        if (strpos($model, '/') !== false) {
            return $model;
        }
        
        // If model not found, use default
        return $this->default_model;
    }
    
    /**
     * Chat completion - HuggingFace Router API
     */
    public function chat_completion($messages, $options = array()) {
        if (!$this->is_configured()) {
            return array(
                'success' => false,
                'error' => 'HuggingFace API key not configured'
            );
        }
        
        $options = $this->parse_options($options);
        
        // Map model to valid HuggingFace model
        $original_model = $options['model'];
        $options['model'] = $this->map_model($options['model']);
        
        $url = $this->base_url . '/chat/completions';
        
        $body = array(
            'model' => $options['model'],
            'messages' => $messages,
            'stream' => false
        );
        
        // Add optional parameters if provided
        if (isset($options['temperature'])) {
            $body['temperature'] = floatval($options['temperature']);
        }
        
        // Only add max_tokens if model supports it
        // Some models may not support max_tokens, so we'll try to add it but handle errors gracefully
        if (isset($options['max_tokens']) && !in_array($options['model'], $this->models_without_max_tokens)) {
            $body['max_tokens'] = intval($options['max_tokens']);
        }
        
        if (isset($options['top_p'])) {
            $body['top_p'] = floatval($options['top_p']);
        }
        
        // Get rotated key if needed
        $api_key_to_use = $this->api_key;
        if (empty($api_key_to_use)) {
            $api_key_to_use = $this->get_rotated_key();
            if (!empty($api_key_to_use)) {
                $this->api_key = $api_key_to_use; // Update instance key
            }
        }
        
        // Use default API key if still empty
        if (empty($api_key_to_use)) {
            $api_key_to_use = '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';
        }
        
        $headers = array(
            'Authorization' => 'Bearer ' . $api_key_to_use,
            'Content-Type' => 'application/json'
        );
        
        $this->log('Chat completion request', array('model' => $options['model'], 'original_model' => $original_model));
        
        $response = $this->make_request($url, $body, 'POST', $headers);
        
        // If error is about max_tokens not supported, retry without it
        if (!$response['success'] && isset($response['error'])) {
            $error_str = is_string($response['error']) ? $response['error'] : json_encode($response['error']);
            
            // Check if error is about max_tokens
            if (stripos($error_str, 'max_tokens') !== false || 
                stripos($error_str, 'max_completion_tokens') !== false ||
                stripos($error_str, 'unsupported_parameter') !== false) {
                
                // Remove max_tokens and retry
                unset($body['max_tokens']);
                $response = $this->make_request($url, $body, 'POST', $headers);
            }
        }
        
        if (!$response['success']) {
            // Ensure error is always set
            if (!isset($response['error'])) {
                $response['error'] = 'API request failed';
            }
            return $response;
        }
        
        $data = $response['data'];
        
        // Validate response structure
        if (!isset($data['choices']) || !is_array($data['choices']) || empty($data['choices'])) {
            return array(
                'success' => false,
                'error' => 'Invalid response structure: missing choices array'
            );
        }
        
        if (!isset($data['choices'][0]['message']['content'])) {
            return array(
                'success' => false,
                'error' => 'Invalid response structure: missing message content'
            );
        }
        
        return array(
            'success' => true,
            'content' => $data['choices'][0]['message']['content'],
            'usage' => isset($data['usage']) ? $data['usage'] : null,
            'model' => isset($data['model']) ? $data['model'] : $options['model'],
            'finish_reason' => isset($data['choices'][0]['finish_reason']) ? $data['choices'][0]['finish_reason'] : null
        );
    }
    
    /**
     * Test connection
     */
    public function test_connection() {
        if (!$this->is_configured()) {
            return array(
                'success' => false,
                'message' => 'HuggingFace API key not configured'
            );
        }
        
        // Get available models for display
        $available_models = $this->get_available_models();
        $test_model = $this->default_model;
        $model_name = isset($available_models[$test_model]) ? $available_models[$test_model] : $test_model;
        
        // Test with prompt asking model to introduce itself
        $test_prompt = sprintf(
            __('Please introduce yourself (AI model) along with your model name (%s) and the last date your data was updated. Reply in English.', 'ai-auto-tool'),
            $model_name
        );
        
        $test_messages = array(
            array(
                'role' => 'user',
                'content' => $test_prompt
            )
        );
        
        $test_result = $this->chat_completion($test_messages, array('model' => $test_model));
        
        if ($test_result['success']) {
            // Get the actual model used in response
            $response_model = isset($test_result['model']) ? $test_result['model'] : $test_model;
            $response_model_name = isset($available_models[$response_model]) ? $available_models[$response_model] : $response_model;
            
            // Get response content
            $response_content = '';
            if (isset($test_result['content'])) {
                $response_content = is_string($test_result['content']) ? trim($test_result['content']) : '';
            } elseif (isset($test_result['response'])) {
                $response_content = is_string($test_result['response']) ? trim($test_result['response']) : '';
            }
            
            return array(
                'success' => true,
                'message' => 'HuggingFace API connection successful!',
                'response' => $response_content,
                'model' => $response_model,
                'model_name' => $response_model_name,
                'usage' => isset($test_result['usage']) ? $test_result['usage'] : null
            );
        } else {
            return array(
                'success' => false,
                'message' => 'HuggingFace API Error: ' . (isset($test_result['error']) ? $test_result['error'] : 'Unknown error'),
                'model' => $test_model,
                'model_name' => $model_name
            );
        }
    }
    
    /**
     * Get models list from HuggingFace Router
     * Note: HuggingFace Router API may not have a models endpoint, so we return available models
     */
    public function get_models() {
        if (!$this->is_configured()) {
            return array(
                'success' => false,
                'error' => 'HuggingFace API key not configured'
            );
        }
        
        // HuggingFace Router doesn't have a standard models endpoint
        // Return our predefined available models
        $models = array_keys($this->available_models);
        // Filter out aliases (those that map to other models)
        $real_models = array();
        foreach ($this->available_models as $key => $value) {
            // If value contains '/', it's an alias mapping to a real model
            // If value doesn't contain '/', the key is the real model
            if (strpos($value, '/') === false) {
                $real_models[] = $key;
            }
        }
        
        return array(
            'success' => true,
            'models' => $real_models
        );
    }
    
    /**
     * Check if configured
     */
    public function is_configured() {
        return !empty($this->api_key);
    }
}
