using System;

namespace AnotherReplayReader.Utils
{
    /// <summary>
    /// 上下文预算策略：每模型软上限、输出余量、估算安全系数与请求护栏。
    /// </summary>
    internal static class AiContextBudget
    {
        public const int Tier1MBudget = 160_000;
        public const int Tier256KBudget = 100_000;
        public const double EstimatorSafetyFactor = 1.2;
        public const double HardUsageRatio = 0.9;

        /// <summary>
        /// 获取模型的一次请求总 token 软上限。0 表示该模型不支持长录像（只能短录像单 slice）。
        /// </summary>
        public static int GetContextBudget(AiModel model)
        {
            if (model.ContextBudget is { } explicitBudget && explicitBudget > 0)
            {
                return explicitBudget;
            }
            if (model.ContextLength >= 1_000_000)
            {
                return Tier1MBudget;
            }
            if (model.ContextLength >= 200_000)
            {
                return Tier256KBudget;
            }
            return 0;
        }

        /// <summary>为输出/推理 tokens 预留的余量。</summary>
        public static int GetOutputHeadroom(AiProvider provider, AiModel model)
        {
            var maxTokens = provider.DefaultMaxTokens;
            return Math.Max(2 * maxTokens, 32_000);
        }

        /// <summary>带安全系数的 token 估算（对中文偏乐观的 bytes/2.2 估算 × 1.2）。</summary>
        public static int EstimateTokens(string text) =>
            (int)Math.Ceiling(AIAnalyze.EstimateTokenCount(text).EstimatedTokenCount * EstimatorSafetyFactor);

        /// <summary>请求护栏：只按 prompt 估算检查（保留给旧调用/测试使用）。</summary>
        public static ContextCheckResult CheckPromptUsage(
            int estimatedPromptTokens,
            AiProvider provider,
            AiModel model)
        {
            return CheckUsage(estimatedPromptTokens, provider, model, includeOutputHeadroom: false);
        }

        /// <summary>请求护栏：把输出/推理余量也算入总用量。</summary>
        public static ContextCheckResult CheckRequestUsage(
            int estimatedPromptTokens,
            AiProvider provider,
            AiModel model)
        {
            return CheckUsage(estimatedPromptTokens, provider, model, includeOutputHeadroom: true);
        }

        private static ContextCheckResult CheckUsage(
            int estimatedPromptTokens,
            AiProvider provider,
            AiModel model,
            bool includeOutputHeadroom)
        {
            var headroom = includeOutputHeadroom ? GetOutputHeadroom(provider, model) : 0;
            var estimatedTotal = estimatedPromptTokens + headroom;
            if (model.ContextLength > 0)
            {
                var hardLimit = (int)(model.ContextLength * HardUsageRatio);
                if (estimatedTotal > hardLimit)
                {
                    return new ContextCheckResult(
                        true,
                        includeOutputHeadroom
                            ? $"估算用量（输入 {estimatedPromptTokens:N0} + 输出余量 {headroom:N0} = {estimatedTotal:N0}）超过模型上下文 {model.ContextLength:N0} 的 90%，已拒绝发起请求。请改用更长上下文的模型，或缩短操作记录。"
                            : $"估算输入 {estimatedPromptTokens:N0} token 超过模型上下文 {model.ContextLength:N0} 的 90%，已拒绝发起请求。请改用更长上下文的模型，或缩短操作记录。");
                }
            }

            var budget = GetContextBudget(model);
            if (budget > 0 && estimatedTotal > budget)
            {
                return new ContextCheckResult(
                    false,
                    includeOutputHeadroom
                        ? $"估算用量（输入 {estimatedPromptTokens:N0} + 输出余量 {headroom:N0} = {estimatedTotal:N0}）超过上下文预算 {budget:N0}（可在模型设置中调整）。"
                        : $"估算输入 {estimatedPromptTokens:N0} token 超过上下文预算 {budget:N0}（可在模型设置中调整），长录像将自动分段，超出部分会被压缩。");
            }
            return ContextCheckResult.Ok;
        }
    }

    internal sealed record ContextCheckResult(bool Block, string Message)
    {
        public static ContextCheckResult Ok { get; } = new(false, string.Empty);
        public bool IsOk => !Block && string.IsNullOrEmpty(Message);
    }
}