<?php

namespace Antimanual;

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

/**
 * Usage Tracking Class
 *
 * Tracks usage of features (e.g., chatbot messages, auto-posting) to enforce
 * free/pro tier limits, and — since v3.7 — records per-request provider, model,
 * token counts, and estimated cost so the Usage & Cost dashboard and per-role
 * spend limits can report and govern AI consumption.
 *
 * Every usage event is a single row. COUNT(*)-based monthly/lifetime limits
 * keep working exactly as before; the new columns are optional metadata.
 *
 * @package Antimanual
 */
class UsageTracker {
    private static $table_name = 'antimanual_usage';

    /**
     * Cached result of the cost-columns existence check.
     *
     * @var bool|null
     */
    private static $has_cost_columns = null;

    /**
     * Get the full table name with prefix.
     *
     * @return string Table name.
     */
    public static function get_table_name() {
        global $wpdb;
        return $wpdb->prefix . self::$table_name;
    }

    /**
     * Get usage count for a specific feature in the current month.
     *
     * @param string $feature Feature identifier.
     * @return int Usage count.
     */
    public static function get_monthly_count($feature) {
        global $wpdb;
        $table = self::get_table_name();
        $month = gmdate('Y-m');

        return intval($wpdb->get_var($wpdb->prepare(
            "SELECT COUNT(*) FROM $table WHERE feature = %s AND month = %s",
            $feature, $month
        )));
    }

    /**
     * Get total (lifetime) count of usage for a feature.
     * Used for features with lifetime limits (e.g., search_block).
     */
    public static function get_total_count($feature) {
        global $wpdb;
        $table = self::get_table_name();

        return intval($wpdb->get_var($wpdb->prepare(
            "SELECT COUNT(*) FROM $table WHERE feature = %s",
            $feature
        )));
    }

    /**
     * Get monthly counts for all features.
     *
     * @return array Associative array of feature counts.
     */
    public static function get_all_monthly_counts() {
        global $wpdb;
        $table = self::get_table_name();
        $month = gmdate('Y-m');

        $results = $wpdb->get_results($wpdb->prepare(
            "SELECT feature, COUNT(*) as count FROM $table
             WHERE month = %s GROUP BY feature",
            $month
        ), ARRAY_A);

        $counts = [];
        foreach ($results as $row) {
            $counts[$row['feature']] = intval($row['count']);
        }

        return $counts;
    }

    /**
     * Increment usage for a feature.
     *
     * Back-compatible: existing call sites pass only a feature string and a
     * bare counting row is inserted. New call sites may pass a metadata array
     * with provider/model/token/cost fields to populate the analytics columns.
     *
     * @param string $feature Feature identifier.
     * @param array  $meta    Optional. {
     *     @type string $provider      Provider slug (e.g. 'openai').
     *     @type string $model         Model identifier.
     *     @type int    $input_tokens  Input/prompt tokens.
     *     @type int    $output_tokens Output/completion tokens.
     *     @type float  $cost          Pre-computed cost in USD (optional; auto-derived when omitted).
     *     @type int    $user_id       User ID (defaults to current user).
     * }
     * @return int|false Row ID on success, false on error.
     */
    public static function increment($feature, array $meta = []) {
        global $wpdb;
        $table = self::get_table_name();

        $data = [
            'feature'    => $feature,
            'month'      => gmdate('Y-m'),
            'user_id'    => isset($meta['user_id']) ? absint($meta['user_id']) : get_current_user_id(),
            'created_at' => current_time('mysql'),
        ];
        $format = ['%s', '%s', '%d', '%s'];

        // Only write the analytics columns when the schema supports them, so
        // installs that have not yet migrated keep working without warnings.
        if (self::has_cost_columns()) {
            $input  = isset($meta['input_tokens']) ? absint($meta['input_tokens']) : 0;
            $output = isset($meta['output_tokens']) ? absint($meta['output_tokens']) : 0;
            $provider = isset($meta['provider']) ? sanitize_key((string) $meta['provider']) : '';
            $model    = isset($meta['model']) ? sanitize_text_field((string) $meta['model']) : '';

            if (isset($meta['cost'])) {
                $cost = (float) $meta['cost'];
            } else {
                $cost = self::estimate_cost($provider, $model, $input, $output);
            }

            $data['provider']      = $provider;
            $data['model']         = substr($model, 0, 60);
            $data['input_tokens']  = $input;
            $data['output_tokens'] = $output;
            $data['total_tokens']  = $input + $output;
            $data['cost']          = round($cost, 6);

            array_push($format, '%s', '%s', '%d', '%d', '%d', '%f');
        }

        return $wpdb->insert($table, $data, $format);
    }

    /**
     * Record an AI usage event, estimating tokens from the request and reply
     * text when explicit token counts are not available.
     *
     * Convenience wrapper around increment() for the common case where a
     * feature has the input and output text on hand but no provider token
     * accounting. Uses the active AIProvider for provider/model when omitted.
     *
     * @param string      $feature  Feature identifier.
     * @param string      $input    Combined prompt/input text sent to the model.
     * @param string      $output   Model reply text.
     * @param string|null $provider Optional provider slug override.
     * @param string|null $model    Optional model identifier override.
     * @return int|false Row ID on success, false on error.
     */
    public static function record($feature, string $input = '', string $output = '', ?string $provider = null, ?string $model = null) {
        if (null === $provider) {
            $provider = class_exists(AIProvider::class) ? AIProvider::get_name() : '';
        }

        if (null === $model) {
            $model = self::resolve_active_model($provider);
        }

        return self::increment($feature, [
            'provider'      => $provider,
            'model'         => $model,
            'input_tokens'  => self::estimate_tokens($input),
            'output_tokens' => self::estimate_tokens($output),
        ]);
    }

    /**
     * Estimate the token count of a string.
     *
     * Uses the widely-used ~4-characters-per-token heuristic, which is accurate
     * enough for cost reporting and spend governance without a tokenizer
     * dependency. Filterable for callers that have exact counts.
     *
     * @param string $text Text to estimate.
     * @return int Estimated token count.
     */
    public static function estimate_tokens(string $text): int {
        if ('' === $text) {
            return 0;
        }

        $length = function_exists('mb_strlen') ? mb_strlen($text) : strlen($text);
        $tokens = (int) ceil($length / 4);

        return max(0, (int) apply_filters('atml_usage_estimate_tokens', $tokens, $text));
    }

    /**
     * Per-1M-token USD prices keyed by model substring.
     *
     * Prices are approximate and intended for relative cost guidance, not
     * billing. Site owners pay their provider directly. Override via the
     * `atml_usage_model_prices` filter to keep current with provider pricing.
     *
     * @return array<string, array{input:float,output:float}>
     */
    public static function get_model_prices(): array {
        $prices = [
            // OpenAI
            'gpt-5-mini'   => ['input' => 0.25, 'output' => 2.00],
            'gpt-5'        => ['input' => 1.25, 'output' => 10.00],
            'gpt-4.1-mini' => ['input' => 0.40, 'output' => 1.60],
            'gpt-4o-mini'  => ['input' => 0.15, 'output' => 0.60],
            'gpt-4o'       => ['input' => 2.50, 'output' => 10.00],
            // Anthropic
            'claude-haiku'  => ['input' => 1.00, 'output' => 5.00],
            'claude-sonnet' => ['input' => 3.00, 'output' => 15.00],
            'claude-opus'   => ['input' => 15.00, 'output' => 75.00],
            // Google Gemini
            'gemini-3-flash' => ['input' => 0.10, 'output' => 0.40],
            'gemini-3-pro'   => ['input' => 1.25, 'output' => 10.00],
            'gemini-2.5-flash' => ['input' => 0.075, 'output' => 0.30],
            'gemini-2.5-pro'   => ['input' => 1.25, 'output' => 10.00],
            // xAI Grok
            'grok-4'      => ['input' => 3.00, 'output' => 15.00],
            'grok-3-mini' => ['input' => 0.30, 'output' => 0.50],
            'grok-3'      => ['input' => 3.00, 'output' => 15.00],
            // DeepSeek
            'deepseek' => ['input' => 0.27, 'output' => 1.10],
        ];

        return apply_filters('atml_usage_model_prices', $prices);
    }

    /**
     * Estimate the USD cost of a request from token counts.
     *
     * @param string $provider Provider slug (used only for the fallback default).
     * @param string $model    Model identifier.
     * @param int    $input    Input tokens.
     * @param int    $output   Output tokens.
     * @return float Estimated cost in USD.
     */
    public static function estimate_cost(string $provider, string $model, int $input, int $output): float {
        $prices  = self::get_model_prices();
        $model_lc = strtolower($model);

        $rate = null;
        foreach ($prices as $needle => $price) {
            if ('' !== $model_lc && false !== strpos($model_lc, strtolower($needle))) {
                $rate = $price;
                break;
            }
        }

        // Neutral default when the model is unknown — keeps reporting non-zero
        // rather than guessing a specific provider's pricing.
        if (null === $rate) {
            $rate = apply_filters('atml_usage_default_price', ['input' => 0.50, 'output' => 1.50], $provider, $model);
        }

        $cost = ($input / 1000000) * (float) $rate['input']
              + ($output / 1000000) * (float) $rate['output'];

        return (float) $cost;
    }

    // ─── Analytics queries (admin dashboard) ───────────────────────────────────

    /**
     * Headline totals for the last N days.
     *
     * @param int $days Look-back window in days.
     * @return array{requests:int,input_tokens:int,output_tokens:int,total_tokens:int,cost:float}
     */
    public static function get_summary(int $days = 30): array {
        return self::summary_between(self::since_datetime($days), null);
    }

    /**
     * Summary for the period immediately preceding the last N days, used to
     * compute period-over-period trend deltas.
     *
     * @param int $days Window length in days.
     * @return array{requests:int,input_tokens:int,output_tokens:int,total_tokens:int,cost:float}
     */
    public static function get_previous_summary(int $days = 30): array {
        $days  = max(1, min(365, $days));
        $start = gmdate('Y-m-d 00:00:00', time() - ((2 * $days - 1) * DAY_IN_SECONDS));
        $end   = self::since_datetime($days);

        return self::summary_between($start, $end);
    }

    /**
     * Aggregate totals for a created_at range [start, end). A null end means
     * "up to now".
     *
     * @param string      $start MySQL datetime (inclusive lower bound).
     * @param string|null $end   MySQL datetime (exclusive upper bound) or null.
     * @return array{requests:int,input_tokens:int,output_tokens:int,total_tokens:int,cost:float}
     */
    private static function summary_between(string $start, ?string $end): array {
        global $wpdb;
        $empty = ['requests' => 0, 'input_tokens' => 0, 'output_tokens' => 0, 'total_tokens' => 0, 'cost' => 0.0];

        if (!self::has_cost_columns()) {
            return $empty;
        }

        $table = self::get_table_name();

        if (null === $end) {
            $row = $wpdb->get_row($wpdb->prepare(
                "SELECT COUNT(*) AS requests,
                        COALESCE(SUM(input_tokens),0) AS input_tokens,
                        COALESCE(SUM(output_tokens),0) AS output_tokens,
                        COALESCE(SUM(total_tokens),0) AS total_tokens,
                        COALESCE(SUM(cost),0) AS cost
                 FROM $table WHERE created_at >= %s",
                $start
            ), ARRAY_A);
        } else {
            $row = $wpdb->get_row($wpdb->prepare(
                "SELECT COUNT(*) AS requests,
                        COALESCE(SUM(input_tokens),0) AS input_tokens,
                        COALESCE(SUM(output_tokens),0) AS output_tokens,
                        COALESCE(SUM(total_tokens),0) AS total_tokens,
                        COALESCE(SUM(cost),0) AS cost
                 FROM $table WHERE created_at >= %s AND created_at < %s",
                $start,
                $end
            ), ARRAY_A);
        }

        if (!$row) {
            return $empty;
        }

        return [
            'requests'      => (int) $row['requests'],
            'input_tokens'  => (int) $row['input_tokens'],
            'output_tokens' => (int) $row['output_tokens'],
            'total_tokens'  => (int) $row['total_tokens'],
            'cost'          => round((float) $row['cost'], 4),
        ];
    }

    /**
     * Usage grouped by feature for the last N days.
     *
     * @param int $days Look-back window in days.
     * @return array<int, array{feature:string,requests:int,total_tokens:int,cost:float}>
     */
    public static function get_by_feature(int $days = 30): array {
        return self::group_by('feature', $days);
    }

    /**
     * Usage grouped by provider for the last N days.
     *
     * @param int $days Look-back window in days.
     * @return array<int, array{provider:string,requests:int,total_tokens:int,cost:float}>
     */
    public static function get_by_provider(int $days = 30): array {
        return self::group_by('provider', $days);
    }

    /**
     * Usage grouped by model for the last N days.
     *
     * @param int $days Look-back window in days.
     * @return array<int, array{model:string,requests:int,total_tokens:int,cost:float}>
     */
    public static function get_by_model(int $days = 30): array {
        return self::group_by('model', $days);
    }

    /**
     * Usage grouped by user for the last N days, with display names resolved.
     *
     * Capped at the top 50 spenders and ordered by cost so the breakdown stays
     * useful for spend governance without returning an unbounded list.
     *
     * @param int $days Look-back window in days.
     * @return array<int, array{user:string,user_id:int,requests:int,total_tokens:int,cost:float}>
     */
    public static function get_by_user(int $days = 30): array {
        global $wpdb;

        if (!self::has_cost_columns()) {
            return [];
        }

        $table = self::get_table_name();
        $since = self::since_datetime($days);

        $rows = $wpdb->get_results($wpdb->prepare(
            "SELECT user_id,
                    COUNT(*) AS requests,
                    COALESCE(SUM(total_tokens),0) AS total_tokens,
                    COALESCE(SUM(cost),0) AS cost
             FROM $table
             WHERE created_at >= %s
             GROUP BY user_id
             ORDER BY cost DESC
             LIMIT 50",
            $since
        ), ARRAY_A);

        if (empty($rows)) {
            return [];
        }

        // Batch-fetch display names so we never query inside the loop.
        $user_ids = array_filter(array_map(static fn($r) => (int) $r['user_id'], $rows));
        $names    = [];
        if (!empty($user_ids)) {
            $users = get_users([
                'include' => $user_ids,
                'fields'  => ['ID', 'display_name'],
            ]);
            foreach ($users as $user) {
                $names[(int) $user->ID] = $user->display_name;
            }
        }

        $out = [];
        foreach ($rows as $row) {
            $uid = (int) $row['user_id'];

            if ($uid <= 0) {
                $label = __('System', 'antimanual');
            } elseif (isset($names[$uid])) {
                $label = $names[$uid];
            } else {
                /* translators: %d: user ID of a deleted account. */
                $label = sprintf(__('Deleted user #%d', 'antimanual'), $uid);
            }

            $out[] = [
                'user'         => $label,
                'user_id'      => $uid,
                'requests'     => (int) $row['requests'],
                'total_tokens' => (int) $row['total_tokens'],
                'cost'         => round((float) $row['cost'], 4),
            ];
        }

        return $out;
    }

    /**
     * Generic GROUP BY helper for an allow-listed column.
     *
     * @param string $column One of 'feature' or 'provider'.
     * @param int    $days   Look-back window in days.
     * @return array<int, array<string, mixed>>
     */
    private static function group_by(string $column, int $days): array {
        global $wpdb;

        // Allow-list the column name — it is interpolated into SQL.
        if (!in_array($column, ['feature', 'provider', 'model'], true) || !self::has_cost_columns()) {
            return [];
        }

        $table = self::get_table_name();
        $since = self::since_datetime($days);

        $rows = $wpdb->get_results($wpdb->prepare(
            "SELECT $column AS label,
                    COUNT(*) AS requests,
                    COALESCE(SUM(total_tokens),0) AS total_tokens,
                    COALESCE(SUM(cost),0) AS cost
             FROM $table
             WHERE created_at >= %s
             GROUP BY $column
             ORDER BY cost DESC",
            $since
        ), ARRAY_A);

        $out = [];
        foreach ($rows as $row) {
            $out[] = [
                $column        => (string) ($row['label'] ?: __('Unknown', 'antimanual')),
                'requests'     => (int) $row['requests'],
                'total_tokens' => (int) $row['total_tokens'],
                'cost'         => round((float) $row['cost'], 4),
            ];
        }

        return $out;
    }

    /**
     * Daily request / token / cost series for the last N days, zero-filled.
     *
     * @param int $days Number of days in the series.
     * @return array<int, array{date:string,requests:int,total_tokens:int,cost:float}>
     */
    public static function get_daily_series(int $days = 30): array {
        global $wpdb;

        $days = max(1, min(365, $days));
        $series = [];
        for ($i = $days - 1; $i >= 0; $i--) {
            $date = gmdate('Y-m-d', time() - ($i * DAY_IN_SECONDS));
            $series[$date] = ['date' => $date, 'requests' => 0, 'total_tokens' => 0, 'cost' => 0.0];
        }

        if (!self::has_cost_columns()) {
            return array_values($series);
        }

        $table = self::get_table_name();
        $since = gmdate('Y-m-d 00:00:00', time() - (($days - 1) * DAY_IN_SECONDS));

        $rows = $wpdb->get_results($wpdb->prepare(
            "SELECT DATE(created_at) AS date,
                    COUNT(*) AS requests,
                    COALESCE(SUM(total_tokens),0) AS total_tokens,
                    COALESCE(SUM(cost),0) AS cost
             FROM $table
             WHERE created_at >= %s
             GROUP BY DATE(created_at)
             ORDER BY date ASC",
            $since
        ), ARRAY_A);

        foreach ($rows as $row) {
            $date = (string) $row['date'];
            if (isset($series[$date])) {
                $series[$date] = [
                    'date'         => $date,
                    'requests'     => (int) $row['requests'],
                    'total_tokens' => (int) $row['total_tokens'],
                    'cost'         => round((float) $row['cost'], 4),
                ];
            }
        }

        return array_values($series);
    }

    /**
     * Sum of tokens consumed by a given user in the current month.
     *
     * Powers per-role monthly spend limits.
     *
     * @param int $user_id User ID.
     * @return int Total tokens this month.
     */
    public static function get_user_monthly_tokens(int $user_id): int {
        global $wpdb;

        if ($user_id <= 0 || !self::has_cost_columns()) {
            return 0;
        }

        $table = self::get_table_name();
        $month = gmdate('Y-m');

        return (int) $wpdb->get_var($wpdb->prepare(
            "SELECT COALESCE(SUM(total_tokens),0) FROM $table WHERE user_id = %d AND month = %s",
            $user_id,
            $month
        ));
    }

    /**
     * Build a UTC datetime string N days before now for range queries.
     *
     * @param int $days Days to subtract.
     * @return string MySQL datetime.
     */
    private static function since_datetime(int $days): string {
        $days = max(1, min(365, $days));
        return gmdate('Y-m-d 00:00:00', time() - (($days - 1) * DAY_IN_SECONDS));
    }

    /**
     * Resolve the active model identifier for a provider slug.
     *
     * @param string $provider Provider slug.
     * @return string Model identifier (empty when unknown).
     */
    private static function resolve_active_model(string $provider): string {
        if ('' === $provider || !class_exists(AIProvider::class)) {
            return '';
        }

        $providers = AIProvider::get_providers();
        $model_key = $providers[$provider]['model'] ?? '';

        return $model_key ? (string) atml_option($model_key) : '';
    }

    /**
     * Whether the analytics columns exist on the usage table.
     *
     * Cached per-request. Used to keep every write/read safe on installs that
     * have not yet run the schema migration.
     *
     * @return bool
     */
    public static function has_cost_columns(): bool {
        if (null !== self::$has_cost_columns) {
            return self::$has_cost_columns;
        }

        global $wpdb;
        $table  = self::get_table_name();
        $column = $wpdb->get_var($wpdb->prepare(
            "SHOW COLUMNS FROM $table LIKE %s",
            'cost'
        ));

        return self::$has_cost_columns = ('cost' === $column);
    }

    /**
     * Create (or migrate) the usage tracking table.
     *
     * dbDelta adds the analytics columns to pre-existing tables on upgrade.
     */
    public static function create_table() {
        global $wpdb;
        $table = self::get_table_name();
        $charset = $wpdb->get_charset_collate();

        $sql = "CREATE TABLE $table (
            id bigint(20) NOT NULL AUTO_INCREMENT,
            feature varchar(50) NOT NULL,
            month varchar(7) NOT NULL,
            user_id bigint(20) DEFAULT 0,
            provider varchar(20) NOT NULL DEFAULT '',
            model varchar(60) NOT NULL DEFAULT '',
            input_tokens int(11) NOT NULL DEFAULT 0,
            output_tokens int(11) NOT NULL DEFAULT 0,
            total_tokens int(11) NOT NULL DEFAULT 0,
            cost decimal(14,6) NOT NULL DEFAULT 0,
            created_at datetime NOT NULL,
            PRIMARY KEY (id),
            KEY feature_month (feature, month),
            KEY user_month (user_id, month),
            KEY created_at (created_at)
        ) $charset;";

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

        // Invalidate the cached column check so the new schema is seen at once.
        self::$has_cost_columns = null;
    }
}
