<?php

declare(strict_types=1);

namespace WordPress\CloudflareWorkersAiProvider\Usage;

use WordPress\CloudflareWorkersAiProvider\Brand;
use WordPress\CloudflareWorkersAiProvider\Lock;
use WordPress\CloudflareWorkersAiProvider\Routing\TaskContext;

use function WordPress\CloudflareWorkersAiProvider\fetch_available_models;
use function WordPress\CloudflareWorkersAiProvider\get_plugin_settings;
use function WordPress\CloudflareWorkersAiProvider\static_model_catalog;
use function WordPress\CloudflareWorkersAiProvider\task_label;

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

final class UsageTracker
{
    public const OPTION = 'aipcf_usage_daily';

    public const NOTICE_OPTION = 'aipcf_budget_notice';

    public const FREE_TIER = 10000;

    public const USD_PER_1K_NEURONS = 0.011;

    public static function boot(): void
    {
        add_action('admin_bar_menu', [self::class, 'adminBarNode'], 90);
        add_action('admin_enqueue_scripts', [self::class, 'adminBarStyle']);
        add_action('wp_enqueue_scripts', [self::class, 'adminBarStyle']);
    }

    public static function today(): string
    {
        return gmdate('Y-m-d');
    }

    public static function record(string $modelId, float $neurons, int $inTokens, int $outTokens, string $task = ''): void
    {
        // Read-modify-write on one option: serialise concurrent requests and read
        // the stored row, not this request's (or the object cache's) copy.
        $locked = Lock::acquire('usage', 2);
        wp_cache_delete(self::OPTION, 'options');
        $all   = self::all();
        $day   = self::today();
        $entry = $all[$day] ?? self::emptyDay();

        $entry['neurons']  += $neurons;
        $entry['requests'] += 1;
        $entry['tokens_in']  += $inTokens;
        $entry['tokens_out'] += $outTokens;
        $entry['models'][$modelId] = ($entry['models'][$modelId] ?? 0) + $neurons;
        $taskKey = $task !== '' ? $task : 'direct';
        $entry['tasks'][$taskKey] = ($entry['tasks'][$taskKey] ?? 0) + $neurons;
        if (strpos($taskKey, 'backfill') === 0) {
            $entry['backfill'] += $neurons;
        }
        if ($taskKey === 'embed') {
            // Rows written before 1.9.0 have no embed total.
            $entry['embed'] = ($entry['embed'] ?? 0.0) + $neurons;
        }

        $all[$day] = $entry;
        krsort($all);
        $all = array_slice($all, 0, 62, true);
        update_option(self::OPTION, $all, false);
        if ($locked) {
            Lock::release('usage');
        }
        // First write on a fresh install: get_option() may have cached "no such
        // option" earlier in this request; make the new row visible immediately.
        wp_cache_delete('notoptions', 'options');
        wp_cache_delete('alloptions', 'options');

        do_action('aipcf_usage_recorded', $modelId, $neurons, $task, $entry);

        self::notify(max((float) $entry['neurons'], AccountUsage::today() ?? 0.0));
    }

    /**
     * Counts vector-store queries a feature made today. They cost no neurons, but Vectorize
     * bills queried dimensions (30 million a month on the free plan).
     */
    public static function recordQueries(string $task, int $count = 1): void
    {
        if ($count < 1) {
            return;
        }
        // Query counts are informational: if the lock is busy, skip this count rather than
        // rewrite the option unlocked and risk losing a concurrent neuron total.
        if (!Lock::acquire('usage', 2)) {
            return;
        }
        wp_cache_delete(self::OPTION, 'options');
        $all   = self::all();
        $day   = self::today();
        $entry = $all[$day] ?? self::emptyDay();
        $entry['queries'][$task] = (int) ($entry['queries'][$task] ?? 0) + $count;
        $all[$day] = $entry;
        krsort($all);
        update_option(self::OPTION, array_slice($all, 0, 62, true), false);
        Lock::release('usage');
        wp_cache_delete('notoptions', 'options');
        wp_cache_delete('alloptions', 'options');
    }

    /**
     * Store queries this calendar month (UTC), all features together.
     */
    public static function monthQueries(): int
    {
        $month = gmdate('Y-m');
        $n     = 0;
        foreach (self::all() as $day => $entry) {
            if (strpos((string) $day, $month) === 0) {
                $n += array_sum(array_map('intval', (array) ($entry['queries'] ?? [])));
            }
        }
        return $n;
    }

    public static function recordFromResponse(string $modelId, array $body, string $task = ''): void
    {
        $usage = [];
        if (isset($body['usage']) && is_array($body['usage'])) {
            $usage = $body['usage'];
        } elseif (isset($body['result']['usage']) && is_array($body['result']['usage'])) {
            $usage = $body['result']['usage'];
        }

        $neurons = isset($usage['neurons']) ? (float) $usage['neurons'] : 0.0;
        $in      = (int) ($usage['prompt_tokens'] ?? $usage['input_tokens'] ?? 0);
        $out     = (int) ($usage['completion_tokens'] ?? $usage['output_tokens'] ?? 0);

        // Image and some task models report no neurons; estimate is left at 0 but the
        // request is still counted so the dashboard shows activity.
        self::record($modelId, $neurons, $in, $out, $task);
    }

    public static function all(): array
    {
        $all = get_option(self::OPTION, []);
        return is_array($all) ? $all : [];
    }

    public static function todayTotals(): array
    {
        return self::all()[self::today()] ?? self::emptyDay();
    }

    public static function todayNeurons(): float
    {
        return (float) (self::todayTotals()['neurons'] ?? 0);
    }

    public static function todayBackfillNeurons(): float
    {
        return (float) (self::todayTotals()['backfill'] ?? 0);
    }

    /**
     * Today's embedding neurons. Estimated: Workers AI reports no usage for embeddings.
     */
    public static function todayEmbedNeurons(): float
    {
        return (float) (self::todayTotals()['embed'] ?? 0);
    }

    /**
     * Today's neurons for one task, e.g. 'search' or 'links'.
     */
    public static function todayTaskNeurons(string $task): float
    {
        return (float) (self::todayTotals()['tasks'][$task] ?? 0);
    }

    public static function budgetEnabled(): bool
    {
        return !empty(get_plugin_settings()['budget_enabled']);
    }

    public static function budget(): int
    {
        $n = (int) (get_plugin_settings()['budget_neurons'] ?? self::FREE_TIER);
        return $n > 0 ? $n : self::FREE_TIER;
    }

    public static function budgetAction(): string
    {
        $a = (string) (get_plugin_settings()['budget_action'] ?? 'stop');
        return $a === 'cheapest' ? 'cheapest' : 'stop';
    }

    public static function overBudget(): bool
    {
        if (!self::budgetEnabled()) {
            return false;
        }

        $used = self::countedNeurons();
        return (bool) apply_filters('aipcf_over_budget', $used >= self::budget(), $used, self::budget());
    }

    /**
     * What the daily budget counts: this site's neurons today, or the whole Cloudflare account's
     * when that is known and larger (other sites and apps on the same account).
     */
    public static function countedNeurons(): float
    {
        return max(self::todayNeurons(), AccountUsage::today() ?? 0.0);
    }

    public static function usedFraction(): float
    {
        return self::countedNeurons() / max(1, self::budget());
    }

    public static function cheapestTextModelId(): string
    {
        $models = fetch_available_models();
        if ($models === []) {
            $models = static_model_catalog();
        }

        $best   = '';
        $bestPx = PHP_FLOAT_MAX;
        foreach ($models as $m) {
            if (!empty($m['vision'])) {
                continue;
            }
            $px = (float) ($m['price_in'] ?? 0) + (float) ($m['price_out'] ?? 0);
            if ($px <= 0) {
                continue;
            }
            if ($px < $bestPx) {
                $bestPx = $px;
                $best   = (string) $m['id'];
            }
        }

        $best = $best !== '' ? $best : '@cf/meta/llama-3.2-1b-instruct';

        return (string) apply_filters('aipcf_cheapest_model', $best);
    }

    public const NEURONS_PER_SHORT_TASK = 2.0;

    public static function describe(float $neurons): string
    {
        $tasks = (int) floor($neurons / self::NEURONS_PER_SHORT_TASK);
        if ($tasks >= 1) {
            return sprintf(
                /* translators: %s: number */
                _n('about %s short task (meta description, excerpt, title)', 'about %s short tasks (meta descriptions, excerpts, titles)', $tasks, 'aipcf-ai-provider-for-cloudflare'),
                number_format($tasks)
            );
        }
        return __('less than one short task', 'aipcf-ai-provider-for-cloudflare');
    }

    public static function neuronsPer1k(float $usdPerMillion): float
    {
        return $usdPerMillion / self::USD_PER_1K_NEURONS;
    }

    /**
     * The 80% and 100% emails, once each a day.
     *
     * @param bool $account The figure is the whole account's, not only this site's.
     */
    public static function notify(float $used, bool $account = false): void
    {
        if (!self::budgetEnabled() || empty(get_plugin_settings()['budget_email'])) {
            return;
        }
        $budget = self::budget();
        if ($used < $budget * 0.8) {
            return;
        }
        $stage = $used >= $budget ? 'full' : 'warn';
        $flag  = get_option(self::NOTICE_OPTION, []);
        $sent  = is_array($flag) && ($flag['day'] ?? '') === self::today() ? (string) ($flag['stage'] ?? '') : '';
        if ($sent === 'full' || $sent === $stage) {
            return;
        }
        update_option(self::NOTICE_OPTION, ['day' => self::today(), 'stage' => $stage], false);

        $site    = (string) wp_parse_url(home_url(), PHP_URL_HOST);
        $subject = $stage === 'full'
            /* translators: 1: site host name, 2: the plugin's name (AIPCF) */
            ? sprintf(__('[%1$s] %2$s: daily budget reached', 'aipcf-ai-provider-for-cloudflare'), $site, Brand::NAME)
            /* translators: 1: site host name, 2: the plugin's name (AIPCF) */
            : sprintf(__('[%1$s] %2$s: 80%% of the daily budget used', 'aipcf-ai-provider-for-cloudflare'), $site, Brand::NAME);
        $lines   = [
            sprintf(
                /* translators: 1: neurons used, 2: daily budget, 3: date (UTC) */
                __('Usage today (UTC, %3$s): %1$s of %2$s neurons.', 'aipcf-ai-provider-for-cloudflare'),
                number_format($used, 0),
                number_format($budget),
                self::today()
            ),
        ];
        if ($account) {
            $lines[] = __('This is the whole Cloudflare account\'s use, from Cloudflare\'s analytics: other sites and apps on the account count too.', 'aipcf-ai-provider-for-cloudflare');
        }
        if ($stage === 'full') {
            $lines[] = self::budgetAction() === 'cheapest'
                ? __('The daily budget is used up. Text requests switch to the cheapest model; Backfill, the content index, Link assistant chat, the Site assistant, Search, Search insights, image generation, image reading and translation pause until midnight UTC.', 'aipcf-ai-provider-for-cloudflare')
                : __('The daily budget is used up. AI requests are paused until midnight UTC.', 'aipcf-ai-provider-for-cloudflare');
        } else {
            $lines[] = self::budgetAction() === 'cheapest'
                ? __('When the daily budget is used up, text requests switch to the cheapest model. Backfill, the content index, Link assistant chat, the Site assistant, Search, Search insights, image generation, image reading and translation pause until midnight UTC.', 'aipcf-ai-provider-for-cloudflare')
                : __('When the daily budget is used up, AI requests pause until midnight UTC.', 'aipcf-ai-provider-for-cloudflare');
        }
        $lines[] = sprintf(
            /* translators: %s: settings page URL */
            __('Change the daily budget under AIPCF > Settings > Budget: %s', 'aipcf-ai-provider-for-cloudflare'),
            \WordPress\CloudflareWorkersAiProvider\Admin\SettingsPage::tabUrl('budget')
        );
        /* translators: %s: the plugin's full name */
        $lines[] = sprintf(__('Sent by %s.', 'aipcf-ai-provider-for-cloudflare'), Brand::FULL);
        wp_mail((string) get_option('admin_email'), $subject, implode("\n\n", $lines) . "\n");
    }

    public static function adminBarNode(\WP_Admin_Bar $bar): void
    {
        if (!current_user_can('manage_options')) {
            return;
        }
        $site    = self::todayNeurons();
        $account = AccountUsage::today();
        $used    = max($site, $account ?? 0.0);
        $guarded = self::budgetEnabled();
        $budget  = $guarded ? self::budget() : self::FREE_TIER;
        $frac    = $used / max(1, $budget);
        if ($guarded) {
            $class = $frac >= 1 ? 'aipcf-usage-over' : ($frac >= 0.8 ? 'aipcf-usage-warn' : 'aipcf-usage-ok');
        } else {
            // No limit to break: past the free allowance usage is billed, which is worth a
            // warning colour but not the red of a budget that stopped requests.
            $class = $frac >= 0.8 ? 'aipcf-usage-warn' : 'aipcf-usage-ok';
        }

        // When the whole account's use is known and larger, that is the number the budget counts.
        $name    = $account !== null && $account > $site
            /* translators: Admin bar label when the figure is the whole Cloudflare account's use; keep it short. */
            ? __('AIPCF account', 'aipcf-ai-provider-for-cloudflare')
            : 'AIPCF';
        $whole   = $account === null ? '' : ' ' . sprintf(
            /* translators: 1: the whole account's neurons today, 2: this site's neurons today */
            __('Whole Cloudflare account today: %1$s neurons; this site: %2$s. The budget counts the larger.', 'aipcf-ai-provider-for-cloudflare'),
            number_format($account),
            number_format($site)
        );

        $bar->add_node([
            'id'    => 'aipcf-usage',
            'title' => sprintf(
                '<span class="%s">%s</span>',
                esc_attr($class),
                esc_html($guarded
                    ? sprintf(
                        /* translators: Admin bar counter; keep it short. 1: "AIPCF" or "AIPCF account", 2: neurons used today, 3: daily budget */
                        __('%1$s %2$s / %3$s', 'aipcf-ai-provider-for-cloudflare'),
                        $name,
                        number_format($used, 0),
                        number_format($budget)
                    )
                    : sprintf(
                        /* translators: Admin bar counter when no daily budget is set; keep it short. 1: "AIPCF" or "AIPCF account", 2: neurons used today, 3: free daily allowance */
                        __('%1$s %2$s / %3$s free', 'aipcf-ai-provider-for-cloudflare'),
                        $name,
                        number_format($used, 0),
                        number_format($budget)
                    ))
            ),
            'href'  => \WordPress\CloudflareWorkersAiProvider\Admin\SettingsPage::pageUrl('aipcf-usage'),
            'meta'  => [
                'title' => $guarded
                    ? sprintf(
                        /* translators: 1: neurons used today, 2: daily budget */
                        __('Cloudflare Workers AI: %1$s of %2$s neurons used today (UTC).', 'aipcf-ai-provider-for-cloudflare'),
                        number_format($used, 0),
                        number_format($budget)
                    ) . $whole . self::featureBreakdown()
                    : sprintf(
                        /* translators: 1: neurons used today, 2: free daily allowance in neurons */
                        __('Cloudflare Workers AI: %1$s neurons used today (UTC); the free daily allowance is %2$s. No daily budget is set, so usage past that is billed.', 'aipcf-ai-provider-for-cloudflare'),
                        number_format($used, 0),
                        number_format($budget)
                    ) . $whole . self::featureBreakdown(),
            ],
        ]);
    }

    /**
     * " By feature: Embeddings 60, Member search 1, ..." for the admin-bar tooltip, largest first.
     */
    private static function featureBreakdown(): string
    {
        $tasks = (array) (self::todayTotals()['tasks'] ?? []);
        arsort($tasks);
        $parts = [];
        foreach (array_slice($tasks, 0, 5, true) as $task => $n) {
            if ((float) $n >= 0.5) {
                $parts[] = task_label((string) $task) . ' ' . number_format((float) $n);
            }
        }
        /* translators: %s: list like "Content index 60, Search 1" */
        return $parts === [] ? '' : ' ' . sprintf(__('By feature: %s.', 'aipcf-ai-provider-for-cloudflare'), implode(', ', $parts));
    }

    public static function adminBarStyle(): void
    {
        if (!is_admin_bar_showing() || !current_user_can('manage_options')) {
            return;
        }
        $css = '#wpadminbar #wp-admin-bar-aipcf-usage .aipcf-usage-warn{color:#f0b849}'
            . '#wpadminbar #wp-admin-bar-aipcf-usage .aipcf-usage-over{color:#f86368;font-weight:600}';
        wp_register_style('aipcf-adminbar', false, [], AIPCF_VERSION);
        wp_enqueue_style('aipcf-adminbar');
        wp_add_inline_style('aipcf-adminbar', $css);
    }

    private static function emptyDay(): array
    {
        return [
            'neurons'    => 0.0,
            'requests'   => 0,
            'tokens_in'  => 0,
            'tokens_out' => 0,
            'backfill'   => 0.0,
            'embed'      => 0.0,
            'models'     => [],
            'tasks'      => [],
            'queries'    => [],
        ];
    }
}