AX-Copilot-Codex/src/AxCopilot/Services/TokenEstimator.cs

namespace AxCopilot.Services;

/// <summary>
/// 모델별 토큰 추정 및 비용 계산 서비스.
/// tiktoken 없이 문자 기반 가중치 추정을 사용합니다.
/// </summary>
public static class TokenEstimator
{
    /// <summary>텍스트의 토큰 수를 추정합니다 (CJK/영문 가중치 적용).</summary>
    public static int Estimate(string text)
    {
        if (string.IsNullOrEmpty(text)) return 0;

        int cjkChars = 0;
        int totalChars = text.Length;
        foreach (var c in text)
        {
            if (c >= 0xAC00 && c <= 0xD7A3) cjkChars++; // 한글
            else if (c >= 0x3000 && c <= 0x9FFF) cjkChars++; // CJK
        }

        double ratio = totalChars > 0 ? (double)cjkChars / totalChars : 0;
        // 영문: ~4글자/토큰, 한글: ~2글자/토큰
        double charsPerToken = 4.0 - ratio * 2.0;
        return Math.Max(1, (int)(totalChars / charsPerToken));
    }

    /// <summary>메시지 리스트의 총 토큰 수를 추정합니다.</summary>
    public static int EstimateMessages(IEnumerable<Models.ChatMessage> messages)
    {
        int total = 0;
        foreach (var m in messages)
        {
            var content = m.Content;
            // _tool_use_blocks JSON은 API 전송 시 구조화된 형식으로 변환되므로
            // 원시 JSON 길이 대비 약 40% 할인 적용 (JSON 래퍼 오버헤드 제거)
            if (m.Role == "assistant" && content != null && content.StartsWith("{\"_tool_use_blocks\""))
                total += (int)(Estimate(content) * 0.6) + 4;
            // tool_result JSON도 구조화 변환됨
            else if (m.Role == "user" && content != null && content.StartsWith("{\"type\":\"tool_result\""))
                total += (int)(Estimate(content) * 0.7) + 4;
            else
                total += Estimate(content ?? "") + 4;
        }
        return total;
    }

    /// <summary>시스템 프롬프트 + 도구 정의에 의한 기본 오버헤드 토큰을 추정합니다.</summary>
    /// <param name="systemPromptLength">시스템 프롬프트 문자 수</param>
    /// <param name="toolCount">등록된 도구 수</param>
    /// <returns>추정 토큰 수</returns>
    public static int EstimateBaseOverhead(int systemPromptLength, int toolCount)
    {
        // 시스템 프롬프트 토큰
        var sysTokens = systemPromptLength > 0 ? (int)(systemPromptLength / 3.5) : 0;
        // 도구 정의: 이름 + 설명 + 파라미터 스키마 ≈ 평균 180토큰/도구
        var toolTokens = toolCount * 180;
        return sysTokens + toolTokens;
    }

    /// <summary>비용을 추정합니다 (USD 기준).</summary>
    public static (double InputCost, double OutputCost) EstimateCost(
        int promptTokens, int completionTokens, string service, string model)
    {
        var (inputPrice, outputPrice) = GetPricing(service, model);
        return (promptTokens / 1_000_000.0 * inputPrice,
                completionTokens / 1_000_000.0 * outputPrice);
    }

    /// <summary>모델별 1M 토큰 가격 (USD). (inputPrice, outputPrice)</summary>
    public static (double Input, double Output) GetPricing(string service, string model)
    {
        var key = $"{service}:{model}".ToLowerInvariant();

        // 2026년 기준 가격표
        return key switch
        {
            _ when key.Contains(string.Concat("cl", "aude-opus")) => (15.0, 75.0),
            _ when key.Contains(string.Concat("cl", "aude-sonnet")) => (3.0, 15.0),
            _ when key.Contains(string.Concat("cl", "aude-haiku")) => (0.25, 1.25),
            _ when key.Contains("gemini-2.5-pro") => (1.25, 10.0),
            _ when key.Contains("gemini-2.5-flash") => (0.15, 0.6),
            _ when key.Contains("gemini-2.0") => (0.1, 0.4),
            _ when key.Contains("ollama") => (0, 0), // 로컬
            _ when key.Contains("vllm") => (0, 0),   // 로컬
            _ => (1.0, 3.0), // 기본 추정
        };
    }

    /// <summary>토큰 수를 읽기 쉬운 문자열로 포맷합니다. 이진 단위(1K=1024) 사용.</summary>
    public static string Format(int count) => count switch
    {
        >= 1_048_576 => $"{count / 1_048_576.0:0.#}M",  // 1024*1024
        >= 1_024 => $"{count / 1_024.0:0.#}K",           // 1024
        _ => count.ToString(),
    };

    /// <summary>비용을 읽기 쉬운 문자열로 포맷합니다.</summary>
    public static string FormatCost(double usd) => usd switch
    {
        0 => "무료",
        < 0.01 => $"${usd:F4}",
        < 1.0 => $"${usd:F3}",
        _ => $"${usd:F2}",
    };

    /// <summary>컨텍스트 사용률을 계산합니다 (0.0~1.0).</summary>
    public static double GetContextUsage(int currentTokens, int maxContextTokens)
    {
        if (maxContextTokens <= 0) return 0;
        return Math.Min(1.0, (double)currentTokens / maxContextTokens);
    }
}