<?php

namespace AtlasAiDev\Awfah\Service;

use WordPress\AI_Client\AI_Client;
use WordPress\AiClient\Files\DTO\File;
use WordPress\AiClient\Providers\Http\DTO\RequestOptions;

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

class AI_Provider_Manager {

    /**
     * Simple helper: generate text from prompt with optional JSON schema, temperature, etc.
     *
     * @param string $prompt
     * @param array  $args
     *
     * @return string
     */
    public static function generate_text( string $prompt, array $args = [] ): string {
        $result = self::generate_text_with_usage( $prompt, $args );
        return $result['text'];
    }

    /**
     * Generate text and return both text and token usage info.
     *
     * @param string $prompt
     * @param array  $args   Optional: temperature, json_schema, tool_slug, post_id.
     *
     * @return array{text: string, usage: array{prompt_tokens: int, completion_tokens: int, total_tokens: int}}
     */
    public static function generate_text_with_usage( string $prompt, array $args = [] ): array {
        $empty = [
            'text'  => '',
            'usage' => [ 'prompt_tokens' => 0, 'completion_tokens' => 0, 'total_tokens' => 0 ],
        ];

        if ( ! class_exists( AI_Client::class ) ) {
            return $empty;
        }

        $builder = AI_Client::prompt( $prompt );

        if ( isset( $args['temperature'] ) ) {
            $builder = $builder->using_temperature( (float) $args['temperature'] );
        }

        if ( isset( $args['json_schema'] ) && is_array( $args['json_schema'] ) ) {
            if ( method_exists( $builder, 'as_json_response' ) ) {
                $builder = $builder->as_json_response( $args['json_schema'] );
            }
        }

        if ( method_exists( $builder, 'is_supported_for_text_generation' ) ) {
            if ( ! $builder->is_supported_for_text_generation() ) {
                Logger::log( 'Prompt not supported for text generation. No provider/model available.' );
                return $empty;
            }
        }

        try {
            // Use generateTextResult() to get token usage alongside text.
            if ( method_exists( $builder, 'generateTextResult' ) ) {
                $result    = $builder->generateTextResult();
                $text      = method_exists( $result, 'toText' ) ? (string) $result->toText() : '';
                $usage     = self::extract_token_usage( $result );

                // Persist usage to DB if context provided.
                if ( ! empty( $args['tool_slug'] ) ) {
                    self::record_usage( $args['tool_slug'], $usage, $args['post_id'] ?? 0 );
                }

                return [ 'text' => $text, 'usage' => $usage ];
            }

            // Fallback: no usage info available.
            $text = (string) $builder->generate_text();
            return [ 'text' => $text, 'usage' => $empty['usage'] ];
        } catch ( \Throwable $e ) {
            Logger::log(
                'AI text generation failed',
                [
                    'error' => $e->getMessage(),
                    'type'  => get_class( $e ),
                ]
            );
            return $empty;
        }
    }

    /**
     * Extract token usage from a GenerativeAiResult.
     *
     * @param object $result GenerativeAiResult instance.
     * @return array{prompt_tokens: int, completion_tokens: int, total_tokens: int}
     */
    private static function extract_token_usage( $result ): array {
        $usage = [ 'prompt_tokens' => 0, 'completion_tokens' => 0, 'total_tokens' => 0 ];

        if ( ! method_exists( $result, 'getTokenUsage' ) ) {
            return $usage;
        }

        $token_usage = $result->getTokenUsage();
        if ( ! $token_usage ) {
            return $usage;
        }

        if ( method_exists( $token_usage, 'getPromptTokens' ) ) {
            $usage['prompt_tokens'] = (int) $token_usage->getPromptTokens();
        }
        if ( method_exists( $token_usage, 'getCompletionTokens' ) ) {
            $usage['completion_tokens'] = (int) $token_usage->getCompletionTokens();
        }
        if ( method_exists( $token_usage, 'getTotalTokens' ) ) {
            $usage['total_tokens'] = (int) $token_usage->getTotalTokens();
        }

        return $usage;
    }

    /**
     * Record AI usage to the database for dashboard reporting.
     *
     * @param string $tool_slug Tool identifier (e.g. 'generate-excerpt').
     * @param array  $usage     Token usage data.
     * @param int    $post_id   Related post ID (0 if none).
     */
    public static function record_usage( string $tool_slug, array $usage, int $post_id = 0 ): void {
        global $wpdb;
        $table = $wpdb->prefix . 'awfah_ai_usage';

        // Ensure table exists.
        if ( $wpdb->get_var( $wpdb->prepare( 'SHOW TABLES LIKE %s', $table ) ) !== $table ) {
            self::create_usage_table();
        }

        $wpdb->insert(
            $table,
            [
                'tool_slug'         => $tool_slug,
                'post_id'           => $post_id,
                'user_id'           => get_current_user_id(),
                'prompt_tokens'     => $usage['prompt_tokens'] ?? 0,
                'completion_tokens' => $usage['completion_tokens'] ?? 0,
                'total_tokens'      => $usage['total_tokens'] ?? 0,
                'created_at'        => current_time( 'mysql' ),
            ],
            [ '%s', '%d', '%d', '%d', '%d', '%d', '%s' ]
        );
    }

    /**
     * Create the AI usage tracking table.
     */
    public static function create_usage_table(): void {
        global $wpdb;
        $table   = $wpdb->prefix . 'awfah_ai_usage';
        $charset = $wpdb->get_charset_collate();

        $sql = "CREATE TABLE IF NOT EXISTS {$table} (
            id bigint(20) unsigned NOT NULL AUTO_INCREMENT,
            tool_slug varchar(100) NOT NULL,
            post_id bigint(20) unsigned NOT NULL DEFAULT 0,
            user_id bigint(20) unsigned NOT NULL DEFAULT 0,
            prompt_tokens int(11) NOT NULL DEFAULT 0,
            completion_tokens int(11) NOT NULL DEFAULT 0,
            total_tokens int(11) NOT NULL DEFAULT 0,
            created_at datetime NOT NULL,
            PRIMARY KEY (id),
            KEY tool_slug (tool_slug),
            KEY created_at (created_at),
            KEY user_id (user_id)
        ) {$charset};";

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

    /**
     * Get aggregated AI usage for a date range.
     *
     * @param string $from Start date (Y-m-d).
     * @param string $to   End date (Y-m-d).
     * @return array Usage data grouped by tool_slug.
     */
    public static function get_usage_stats( string $from, string $to ): array {
        global $wpdb;
        $table = $wpdb->prefix . 'awfah_ai_usage';

        if ( $wpdb->get_var( $wpdb->prepare( 'SHOW TABLES LIKE %s', $table ) ) !== $table ) {
            return [];
        }

        $results = $wpdb->get_results(
            $wpdb->prepare(
                "SELECT
                    tool_slug,
                    COUNT(*) as request_count,
                    SUM(prompt_tokens) as total_prompt_tokens,
                    SUM(completion_tokens) as total_completion_tokens,
                    SUM(total_tokens) as total_total_tokens
                FROM {$table}
                WHERE created_at >= %s AND created_at < %s
                GROUP BY tool_slug
                ORDER BY total_total_tokens DESC",
                $from . ' 00:00:00',
                $to . ' 23:59:59'
            ),
            ARRAY_A
        );

        return $results ?: [];
    }

    /**
     * Generate an image from a prompt using the configured AI provider.
     *
     * Uses a 120-second request timeout to accommodate longer image generation times.
     * Returns a File object with getBase64Data() and getMimeType() on success.
     *
     * @param string $prompt The image generation prompt.
     * @param array  $args   Optional arguments (temperature, model_preferences).
     *
     * @return File|null File object on success, null on failure.
     */
    public static function generate_image( string $prompt, array $args = [] ): ?File {
        if ( ! class_exists( AI_Client::class ) ) {
            return null;
        }

        $builder = AI_Client::prompt( $prompt );

        // Set 120-second timeout for image generation (default is 30s).
        if ( method_exists( $builder, 'using_request_options' ) ) {
            $builder = $builder->using_request_options(
                RequestOptions::fromArray( [ RequestOptions::KEY_TIMEOUT => 120 ] )
            );
        }

        if ( isset( $args['temperature'] ) ) {
            $builder = $builder->using_temperature( (float) $args['temperature'] );
        }

        // Check if image generation is supported by the configured provider.
        if ( method_exists( $builder, 'is_supported_for_image_generation' ) ) {
            if ( ! $builder->is_supported_for_image_generation() ) {
                Logger::log( 'No provider/model available for image generation.' );
                return null;
            }
        }

        // Hook into WordPress HTTP API to force 120s timeout on AI provider requests.
        $timeout_filter = static function ( $parsed_args, $url ) {
            $ai_endpoints = [
                'api.openai.com/v1/images',
                'generativelanguage.googleapis.com',
                'api.anthropic.com',
            ];

            foreach ( $ai_endpoints as $endpoint ) {
                if ( false !== strpos( $url, $endpoint ) ) {
                    $parsed_args['timeout'] = 120;
                    break;
                }
            }

            return $parsed_args;
        };

        add_filter( 'http_request_args', $timeout_filter, 9999, 2 );

        try {
            $file = $builder->generate_image();

            if ( $file instanceof File && $file->isImage() ) {
                return $file;
            }

            return null;
        } catch ( \Throwable $e ) {
            Logger::log( 'AI image generation failed', [
                'error' => $e->getMessage(),
                'type'  => get_class( $e ),
            ] );
            return null;
        } finally {
            remove_filter( 'http_request_args', $timeout_filter, 9999 );
        }
    }
}
