<?php

namespace AgoLab\AIChatbot;

defined( 'ABSPATH' ) || exit;

class GeminiAPI {

    private const BASE_URL = 'https://generativelanguage.googleapis.com/v1beta';

    /**
     * Google keeps this alias pointing at its current Flash Lite model. A fixed
     * id rots: Google retires versions, and some stay in the model list while
     * answering 404 to keys created after the retirement date.
     */
    public const DEFAULT_MODEL = 'gemini-flash-lite-latest';

    /**
     * Send a chat message using File Search (if store exists) or file_data fallback.
     */
    public static function chat( array $settings, string $message, array $history = [] ): array {
        $api_key = $settings['api_key'];
        $model   = ! empty( $settings['model'] ) ? $settings['model'] : self::DEFAULT_MODEL;

        // Build contents array with history
        $contents = [];
        foreach ( $history as $msg ) {
            $role = $msg['role'] === 'user' ? 'user' : 'model';
            $contents[] = [ 'role' => $role, 'parts' => [ [ 'text' => $msg['text'] ] ] ];
        }

        // Add current message
        $contents[] = [ 'role' => 'user', 'parts' => [ [ 'text' => $message ] ] ];

        // Build system instruction
        $system_text = self::build_system_prompt( $settings );
        $body = [ 'contents' => $contents ];
        if ( $system_text ) {
            $body['systemInstruction'] = [ 'parts' => [ [ 'text' => $system_text ] ] ];
        }

        $body['generationConfig'] = [
            'temperature'     => 0.7,
            'maxOutputTokens' => (int) ( $settings['max_output_tokens'] ?? 1000 ),
        ];

        // Use File Search Store if available (cheaper: only relevant fragments)
        $store_name = get_option( 'agoaichat_store_name', '' );
        if ( $store_name ) {
            $body['tools'] = [ [ 'fileSearch' => [ 'fileSearchStoreNames' => [ $store_name ] ] ] ];
        } else {
            // Fallback: attach files directly (sends entire file each time = expensive)
            $files = get_option( 'agoaichat_files', [] );
            if ( $files ) {
                $file_parts = [];
                foreach ( $files as $file ) {
                    if ( ! empty( $file['uri'] ) ) {
                        $file_parts[] = [ 'file_data' => [ 'mime_type' => $file['mime_type'], 'file_uri' => $file['uri'] ] ];
                    }
                }
                if ( $file_parts ) {
                    // Prepend file parts to the last user message
                    $last_idx = count( $contents ) - 1;
                    $contents[ $last_idx ]['parts'] = array_merge( $file_parts, $contents[ $last_idx ]['parts'] );
                    $body['contents'] = $contents;
                }
            }
        }

        $result = self::generate( $api_key, $model, $body );

        // The chosen model is gone for this key or overloaded: answer with the
        // alias instead of leaving the visitor with an error, and say so.
        if ( $result['error'] && $result['retry'] && $model !== self::DEFAULT_MODEL ) {
            $fallback = self::generate( $api_key, self::DEFAULT_MODEL, $body );
            if ( ! $fallback['error'] ) {
                $fallback['model_error'] = $model . ': ' . $result['error'];
                $fallback['model']       = self::DEFAULT_MODEL;
                return $fallback;
            }
        }

        $result['model'] = $model;
        return $result;
    }

    /**
     * One generateContent call. 'retry' marks the failures another model can
     * solve: the model does not exist for this key, or it is overloaded.
     *
     * @return array{text:string, error:?string, retry:bool}
     */
    private static function generate( string $api_key, string $model, array $body ): array {
        $response = wp_remote_post( self::BASE_URL . '/models/' . rawurlencode( $model ) . ':generateContent', [
            'headers' => [ 'Content-Type' => 'application/json', 'x-goog-api-key' => $api_key ],
            'body'    => wp_json_encode( $body ),
            'timeout' => 30,
        ] );

        if ( is_wp_error( $response ) ) {
            return [ 'text' => '', 'error' => $response->get_error_message(), 'retry' => false ];
        }

        $code = (int) wp_remote_retrieve_response_code( $response );
        $data = json_decode( wp_remote_retrieve_body( $response ), true );

        if ( $code !== 200 ) {
            return [
                'text'  => '',
                'error' => $data['error']['message'] ?? 'API error (HTTP ' . $code . ')',
                'retry' => in_array( $code, [ 404, 500, 503 ], true ),
                'code'  => $code,
            ];
        }

        $candidate = $data['candidates'][0] ?? [];
        $text      = '';
        foreach ( $candidate['content']['parts'] ?? [] as $part ) {
            if ( isset( $part['text'] ) && empty( $part['thought'] ) ) {
                $text .= $part['text'];
            }
        }

        if ( trim( $text ) === '' ) {
            $reason = (string) ( $candidate['finishReason'] ?? ( $data['promptFeedback']['blockReason'] ?? 'unknown' ) );
            return [
                'text'  => '',
                'error' => 'MAX_TOKENS' === $reason
                    ? 'The model used up its output budget before answering. Raise Max Output Tokens.'
                    : 'The model returned an empty answer (' . $reason . ').',
                'retry' => false,
            ];
        }

        return [ 'text' => $text, 'error' => null, 'retry' => false ];
    }

    /**
     * Send a minimal request with the model, so a model the key cannot use is
     * caught when the settings are saved and not by the first visitor.
     *
     * @return array{ok:bool, error?:string}
     */
    public static function check_model( string $api_key, string $model ): array {
        $result = self::generate( $api_key, $model, [
            'contents'         => [ [ 'role' => 'user', 'parts' => [ [ 'text' => 'Hi' ] ] ] ],
            'generationConfig' => [ 'maxOutputTokens' => 256 ],
        ] );
        // Only a model that does not exist for this key blocks the save. An
        // overloaded model or an empty reply still proves the model is there.
        if ( $result['error'] && in_array( $result['code'] ?? 0, [ 400, 404 ], true ) ) {
            return [ 'ok' => false, 'error' => $result['error'] ];
        }
        return [ 'ok' => true ];
    }

    /* ───── File Search Store Management ───── */

    /**
     * Create a File Search Store. Returns store name.
     */
    public static function create_store( string $api_key, string $display_name ): array {
        $response = wp_remote_post( self::BASE_URL . '/fileSearchStores', [
            'headers' => [ 'Content-Type' => 'application/json', 'x-goog-api-key' => $api_key ],
            'body'    => wp_json_encode( [ 'displayName' => $display_name ] ),
            'timeout' => 15,
        ] );

        if ( is_wp_error( $response ) ) {
            return [ 'error' => $response->get_error_message() ];
        }

        $data = json_decode( wp_remote_retrieve_body( $response ), true );
        if ( empty( $data['name'] ) ) {
            return [ 'error' => $data['error']['message'] ?? 'Failed to create store' ];
        }

        return [ 'name' => $data['name'], 'displayName' => $data['displayName'] ?? '' ];
    }

    /**
     * Upload a file and add it to the File Search Store.
     */
    public static function upload_to_store( string $api_key, string $store_name, string $file_path, string $display_name, string $mime_type ): array {
        // Step 1: Upload file via Files API
        $upload_result = self::upload_file( $api_key, $file_path, $display_name, $mime_type );
        if ( ! empty( $upload_result['error'] ) ) {
            return $upload_result;
        }

        $file_name = $upload_result['name'];

        // Step 2: import the file into the store. Indexing runs as a long
        // operation; until it finishes the document is not searchable.
        $import = wp_remote_post( self::BASE_URL . '/' . $store_name . ':importFile', [
            'headers' => [ 'Content-Type' => 'application/json', 'x-goog-api-key' => $api_key ],
            'body'    => wp_json_encode( [ 'fileName' => $file_name ] ),
            'timeout' => 30,
        ] );

        if ( is_wp_error( $import ) ) {
            return [ 'error' => $import->get_error_message() ];
        }

        $code      = (int) wp_remote_retrieve_response_code( $import );
        $operation = json_decode( wp_remote_retrieve_body( $import ), true );
        if ( 200 !== $code || empty( $operation['name'] ) ) {
            return [ 'error' => 'Indexing failed: ' . ( $operation['error']['message'] ?? 'HTTP ' . $code ) ];
        }

        for ( $i = 0; empty( $operation['done'] ) && $i < 10; $i++ ) {
            sleep( 2 );
            $poll = wp_remote_get( self::BASE_URL . '/' . $operation['name'], [
                'headers' => [ 'x-goog-api-key' => $api_key ],
                'timeout' => 15,
            ] );
            if ( is_wp_error( $poll ) ) break;
            $operation = json_decode( wp_remote_retrieve_body( $poll ), true ) ?: $operation;
        }

        if ( ! empty( $operation['error'] ) ) {
            return [ 'error' => 'Indexing failed: ' . ( $operation['error']['message'] ?? '' ) ];
        }

        // The operation and the document it creates share the same id.
        $document = $store_name . '/documents/' . basename( (string) ( $operation['response']['documentName'] ?? $operation['name'] ) );

        return [
            'name'     => $file_name,
            'uri'      => $upload_result['uri'],
            'document' => $document,
            'in_store' => true,
        ];
    }

    /**
     * Documents in the store, as document name => state (STATE_ACTIVE,
     * STATE_PENDING or STATE_FAILED). Null when the store cannot be read.
     */
    public static function store_documents( string $api_key, string $store_name ): ?array {
        $docs  = [];
        $token = '';
        do {
            $url      = self::BASE_URL . '/' . $store_name . '/documents?pageSize=20' . ( $token ? '&pageToken=' . rawurlencode( $token ) : '' );
            $response = wp_remote_get( $url, [ 'headers' => [ 'x-goog-api-key' => $api_key ], 'timeout' => 15 ] );
            if ( is_wp_error( $response ) || 200 !== (int) wp_remote_retrieve_response_code( $response ) ) {
                return null;
            }
            $data = json_decode( wp_remote_retrieve_body( $response ), true );
            foreach ( $data['documents'] ?? [] as $doc ) {
                $docs[ $doc['name'] ] = $doc['state'] ?? '';
            }
            $token = (string) ( $data['nextPageToken'] ?? '' );
        } while ( $token );

        return $docs;
    }

    /**
     * Remove one document from the store.
     */
    public static function delete_document( string $api_key, string $document ): void {
        wp_remote_request( self::BASE_URL . '/' . $document . '?force=true', [
            'method'  => 'DELETE',
            'headers' => [ 'x-goog-api-key' => $api_key ],
            'timeout' => 15,
        ] );
    }

    /**
     * Delete a File Search Store.
     */
    public static function delete_store( string $api_key, string $store_name ): void {
        wp_remote_request( self::BASE_URL . '/' . $store_name . '?force=true', [
            'method'  => 'DELETE',
            'headers' => [ 'x-goog-api-key' => $api_key ],
            'timeout' => 15,
        ] );
    }

    /* ───── Files API (upload/delete) ───── */

    /**
     * Upload a file to Gemini Files API.
     */
    public static function upload_file( string $api_key, string $file_path, string $display_name, string $mime_type ): array {
        $file_size = filesize( $file_path );

        // Start resumable upload
        $start_url = 'https://generativelanguage.googleapis.com/upload/v1beta/files';

        $start_response = wp_remote_post( $start_url, [
            'headers' => [
                'x-goog-api-key'                     => $api_key,
                'X-Goog-Upload-Protocol'             => 'resumable',
                'X-Goog-Upload-Command'              => 'start',
                'X-Goog-Upload-Header-Content-Length' => $file_size,
                'X-Goog-Upload-Header-Content-Type'  => $mime_type,
                'Content-Type'                       => 'application/json',
            ],
            'body'    => wp_json_encode( [ 'file' => [ 'display_name' => $display_name ] ] ),
            'timeout' => 30,
        ] );

        if ( is_wp_error( $start_response ) ) {
            return [ 'error' => $start_response->get_error_message() ];
        }

        $headers    = wp_remote_retrieve_headers( $start_response );
        $upload_url = $headers['x-goog-upload-url'] ?? '';
        if ( ! $upload_url ) {
            return [ 'error' => 'Failed to get upload URL' ];
        }

        // Upload bytes
        $file_contents = file_get_contents( $file_path );
        $upload_response = wp_remote_request( $upload_url, [
            'method'  => 'PUT',
            'headers' => [
                'Content-Length'         => $file_size,
                'X-Goog-Upload-Offset'  => '0',
                'X-Goog-Upload-Command' => 'upload, finalize',
            ],
            'body'    => $file_contents,
            'timeout' => 60,
        ] );

        if ( is_wp_error( $upload_response ) ) {
            return [ 'error' => $upload_response->get_error_message() ];
        }

        $data = json_decode( wp_remote_retrieve_body( $upload_response ), true );
        if ( empty( $data['file']['uri'] ) ) {
            return [ 'error' => 'Upload failed: ' . wp_json_encode( $data ) ];
        }

        return [ 'name' => $data['file']['name'] ?? '', 'uri' => $data['file']['uri'] ];
    }

    /**
     * Delete a file from Gemini.
     */
    public static function delete_file( string $api_key, string $name ): void {
        wp_remote_request( self::BASE_URL . '/' . $name, [
            'method'  => 'DELETE',
            'headers' => [ 'x-goog-api-key' => $api_key ],
            'timeout' => 15,
        ] );
    }

    /* ───── Models ───── */

    /**
     * Cheap key check against the models endpoint before persisting it.
     */
    public static function validate_key( string $api_key ): array {
        $response = wp_remote_get( self::BASE_URL . '/models?pageSize=1', [
            'headers' => [ 'x-goog-api-key' => $api_key ],
            'timeout' => 15,
        ] );

        if ( is_wp_error( $response ) ) {
            return [ 'ok' => false, 'error' => $response->get_error_message() ];
        }

        $code = wp_remote_retrieve_response_code( $response );
        if ( $code !== 200 ) {
            $data = json_decode( wp_remote_retrieve_body( $response ), true );
            return [ 'ok' => false, 'error' => $data['error']['message'] ?? 'API error (HTTP ' . $code . ')' ];
        }

        return [ 'ok' => true ];
    }

    /**
     * Get available text chat models (Lite, Flash, Pro). Sorted cheapest tier
     * first, newest version first within each tier. Dynamic: any new Gemini
     * model the API exposes appears without code changes.
     */
    public static function list_models( string $api_key ): array {
        $response = wp_remote_get( self::BASE_URL . '/models?pageSize=200', [ 'headers' => [ 'x-goog-api-key' => $api_key ], 'timeout' => 15 ] );
        if ( is_wp_error( $response ) ) return [];

        $data = json_decode( wp_remote_retrieve_body( $response ), true );

        $tiers = [
            'lite'  => [ 'order' => 1, 'tier' => __( 'Lowest cost', 'ago-ai-chatbot' ) ],
            'flash' => [ 'order' => 2, 'tier' => __( 'Mid cost', 'ago-ai-chatbot' ) ],
            'pro'   => [ 'order' => 3, 'tier' => __( 'Highest quality', 'ago-ai-chatbot' ) ],
        ];

        $available = [];
        foreach ( $data['models'] ?? [] as $m ) {
            $name = str_replace( 'models/', '', $m['name'] );

            // Only Gemini text chat models. Non-conversational variants out.
            if ( ! str_starts_with( $name, 'gemini-' ) ) continue;
            if ( preg_match( '/(tts|image|audio|live|robotics|computer-use|embedding|banana|transcribe|customtools)/', $name ) ) continue;
            if ( ! in_array( 'generateContent', $m['supportedGenerationMethods'] ?? [], true ) ) continue;

            if ( str_contains( $name, 'lite' ) ) {
                $group = 'lite';
            } elseif ( str_contains( $name, 'flash' ) ) {
                $group = 'flash';
            } elseif ( str_contains( $name, 'pro' ) ) {
                $group = 'pro';
            } else {
                continue;
            }

            // The alias goes first in its tier: it never gets retired.
            $alias = str_ends_with( $name, '-latest' );
            $label = ucwords( str_replace( '-', ' ', preg_replace( '/^gemini-/', 'Gemini ', $name ) ) );
            $available[] = [
                'id'          => $name,
                'label'       => $label,
                'tier'        => $tiers[ $group ]['tier'],
                'order'       => $tiers[ $group ]['order'] * 2 - ( $alias ? 1 : 0 ),
                'recommended' => self::DEFAULT_MODEL === $name,
            ];
        }

        usort( $available, fn( $a, $b ) => [ $a['order'], $b['id'] ] <=> [ $b['order'], $a['id'] ] );
        return $available;
    }

    /* ───── System Prompt ───── */

    private static function build_system_prompt( array $settings ): string {
        $site_name = get_bloginfo( 'name' );
        $parts = [];

        $bot_name = $settings['bot_name'] ?: 'Asistente';
        $parts[] = "You are {$bot_name}, the friendly assistant for \"{$site_name}\".";

        $style = $settings['response_style'] ?? 'friendly_emoji';
        $styles = [
            'friendly_emoji'  => "Respond in a warm, conversational, and human tone. Use emojis naturally (but not excessively) to make responses feel approachable 😊. Structure your answers with short paragraphs. Use bullet points only when listing more than 2 items. Avoid sounding robotic or overly formal.",
            'friendly_plain'  => "Respond in a warm, conversational, and human tone. Structure your answers with short paragraphs. Be concise and clear. Avoid sounding robotic.",
            'professional'    => "Respond in a professional but approachable tone. Be clear, structured, and concise. Use proper formatting with paragraphs and bullet points when appropriate.",
            'formal'          => "Respond in a formal, respectful tone. Use proper language and structured responses. Be thorough but concise.",
        ];
        $parts[] = $styles[ $style ] ?? $styles['friendly_emoji'];

        $parts[] = "IMPORTANT: You must ONLY answer questions related to the information in the knowledge files provided. If someone asks about unrelated topics (coding, math, general knowledge, etc.), politely redirect them: \"I'm here specifically to help you with information about {$site_name}. Is there something about our services I can help with? 😊\"";

        $tone = $settings['tone'] ?? 'friendly';
        if ( $tone === 'casual' ) {
            $parts[] = 'Be casual and relaxed, like talking to a friend. Use informal language.';
        }

        if ( ! empty( $settings['system_prompt'] ) ) {
            $parts[] = $settings['system_prompt'];
        }

        $parts[] = 'If a user asks to contact a human, suggest they use the contact options visible on this site (phone, email, contact form) and stay friendly.';

        $parts[] = 'If you cannot find the answer in your knowledge files, say so honestly and suggest they reach out through the site contact options. Never invent information.';

        $parts[] = 'Always respond in the same language the user writes in.';
        $parts[] = 'Keep responses concise, aim for 2-4 short paragraphs maximum unless the user asks for detail.';
        $parts[] = 'NEVER discuss your own errors, limitations, or technical issues. If the user mentions something "went wrong", apologize briefly and ask how you can help with ' . $site_name . ' services.';

        return implode( "\n\n", $parts );
    }
}