"""
Token 使用统计服务

负责记录、查询和分析 LLM 调用的 token 使用情况
数据持久化到 data/token_usage 目录，按月份组织
采用双层存储：
1. JSONL 文件 - 详细记录（按月存储）
2. 索引文件 - 快速聚合统计（按 API URL、日期等维度）
"""
import json
import uuid
from pathlib import Path
from typing import List, Dict, Optional, Any
from datetime import datetime
from collections import defaultdict

try:
    from backend.models.internal import TokenUsageRecord, TokenUsageStatus
    from backend.core.config import settings
except ImportError:
    from models.internal import TokenUsageRecord, TokenUsageStatus
    from core.config import settings


class TokenUsageService:
    """
    Token 使用统计服务
    
    功能：
    - 记录每次 LLM 调用的 token 使用情况
    - 按月份、日期、角色、聊天、API URL 维度统计
    - 支持中断和失败标记
    - 数据持久化到文件系统（JSONL + 索引）
    """
    
    def __init__(self):
        self.token_usage_dir = settings.DATA_PATH / "token_usage"
        self.token_usage_dir.mkdir(parents=True, exist_ok=True)
        
        # ✅ 索引文件目录 - 用于快速聚合查询
        self.index_dir = self.token_usage_dir / "indexes"
        self.index_dir.mkdir(parents=True, exist_ok=True)
    
    def _get_month_file(self, year: int, month: int) -> Path:
        """获取指定月份的统计文件路径"""
        month_dir = self.token_usage_dir / f"{year}"
        month_dir.mkdir(parents=True, exist_ok=True)
        return month_dir / f"{month:02d}.jsonl"
    
    def _load_month_records(self, year: int, month: int) -> List[TokenUsageRecord]:
        """加载指定月份的所有记录"""
        file_path = self._get_month_file(year, month)
        
        if not file_path.exists():
            return []
        
        records = []
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                for line in f:
                    if line.strip():
                        data = json.loads(line)
                        records.append(TokenUsageRecord(**data))
        except Exception as e:
            print(f"[TokenUsage] 加载记录失败: {e}")
        
        return records
    
    def _save_record(self, record: TokenUsageRecord):
        """保存单条记录到对应的月份文件"""
        dt = datetime.fromtimestamp(record.timestamp)
        file_path = self._get_month_file(dt.year, dt.month)
        
        try:
            with open(file_path, 'a', encoding='utf-8') as f:
                f.write(json.dumps(record.model_dump(), ensure_ascii=False) + '\n')
            
            # ✅ 同时更新索引文件（用于快速查询）
            self._update_indexes(record)
        except Exception as e:
            print(f"[TokenUsage] 保存记录失败: {e}")
    
    def _update_indexes(self, record: TokenUsageRecord):
        """
        更新索引文件 - 实现高效的按维度聚合查询
        
        索引结构：
        - indexes/api_urls.json - 按 API URL 聚合
        - indexes/daily/{year}-{month}.json - 按日聚合
        """
        dt = datetime.fromtimestamp(record.timestamp)
        
        # 1. 更新 API URL 索引
        if record.apiUrl:
            api_url_index = self.index_dir / "api_urls.json"
            self._update_api_url_index(api_url_index, record)
        
        # 2. 更新每日索引
        daily_index = self.index_dir / "daily" / f"{dt.year}-{dt.month:02d}.json"
        daily_index.parent.mkdir(parents=True, exist_ok=True)
        self._update_daily_index(daily_index, record)
    
    def _update_api_url_index(self, index_file: Path, record: TokenUsageRecord):
        """更新 API URL 索引文件"""
        index_data = {}
        
        # 加载现有索引
        if index_file.exists():
            try:
                with open(index_file, 'r', encoding='utf-8') as f:
                    index_data = json.load(f)
            except:
                index_data = {}
        
        # 更新统计
        api_url = record.apiUrl
        if api_url not in index_data:
            index_data[api_url] = {
                "totalPromptTokens": 0,
                "totalCompletionTokens": 0,
                "totalTokens": 0,
                "count": 0,
                "firstUsed": record.timestamp,
                "lastUsed": record.timestamp
            }
        
        stats = index_data[api_url]
        stats["totalPromptTokens"] += record.promptTokens
        stats["totalCompletionTokens"] += record.completionTokens
        stats["totalTokens"] += record.totalTokens
        stats["count"] += 1
        stats["lastUsed"] = max(stats["lastUsed"], record.timestamp)
        stats["firstUsed"] = min(stats["firstUsed"], record.timestamp)
        
        # 保存索引
        with open(index_file, 'w', encoding='utf-8') as f:
            json.dump(index_data, f, ensure_ascii=False, indent=2)
    
    def _update_daily_index(self, index_file: Path, record: TokenUsageRecord):
        """更新每日索引文件"""
        dt = datetime.fromtimestamp(record.timestamp)
        day_key = f"{dt.year}-{dt.month:02d}-{dt.day:02d}"
        
        index_data = {}
        
        # 加载现有索引
        if index_file.exists():
            try:
                with open(index_file, 'r', encoding='utf-8') as f:
                    index_data = json.load(f)
            except:
                index_data = {}
        
        # 更新统计
        if day_key not in index_data:
            index_data[day_key] = {
                "promptTokens": 0,
                "completionTokens": 0,
                "totalTokens": 0,
                "count": 0
            }
        
        stats = index_data[day_key]
        stats["promptTokens"] += record.promptTokens
        stats["completionTokens"] += record.completionTokens
        stats["totalTokens"] += record.totalTokens
        stats["count"] += 1
        
        # 保存索引
        with open(index_file, 'w', encoding='utf-8') as f:
            json.dump(index_data, f, ensure_ascii=False, indent=2)
    
    async def record_usage(
        self,
        chat_id: str,
        role_name: str,
        chat_name: str,
        prompt_tokens: int,
        completion_tokens: int,
        total_tokens: int,
        status: TokenUsageStatus = TokenUsageStatus.COMPLETED,
        message_id: Optional[str] = None,
        floor: Optional[int] = None,
        error_message: Optional[str] = None,
        duration: Optional[float] = None,
        model: Optional[str] = None,
        api_provider: Optional[str] = None,
        api_url: Optional[str] = None
    ) -> TokenUsageRecord:
        """
        记录一次 LLM 调用的 token 使用情况
        
        Args:
            chat_id: 聊天ID
            role_name: 角色名称
            chat_name: 聊天名称
            prompt_tokens: 输入 token 数
            completion_tokens: 输出 token 数
            total_tokens: 总 token 数
            status: 请求状态
            message_id: 关联的消息ID
            floor: 楼层号
            error_message: 错误信息
            duration: 请求耗时
            model: 使用的模型
            api_provider: API 提供商
            api_url: API URL地址
            
        Returns:
            TokenUsageRecord: 创建的记录
        """
        record = TokenUsageRecord(
            id=str(uuid.uuid4()),
            chatId=chat_id,
            roleName=role_name,
            chatName=chat_name,
            messageId=message_id,
            floor=floor,
            promptTokens=prompt_tokens,
            completionTokens=completion_tokens,
            totalTokens=total_tokens,
            status=status,
            errorMessage=error_message,
            duration=duration,
            model=model,
            apiProvider=api_provider,
            apiUrl=api_url
        )
        
        self._save_record(record)
        return record
    
    async def get_stats_by_month(
        self,
        year: int,
        month: int,
        role_name: Optional[str] = None,
        chat_name: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        获取指定月份的统计数据
        
        Args:
            year: 年份
            month: 月份
            role_name: 角色名称（可选，用于过滤）
            chat_name: 聊天名称（可选，用于过滤）
            
        Returns:
            统计数据字典
        """
        records = self._load_month_records(year, month)
        
        # 过滤
        if role_name:
            records = [r for r in records if r.roleName == role_name]
        if chat_name:
            records = [r for r in records if r.chatName == chat_name]
        
        # 统计
        total_prompt = sum(r.promptTokens for r in records)
        total_completion = sum(r.completionTokens for r in records)
        total_tokens = sum(r.totalTokens for r in records)
        
        completed_count = sum(1 for r in records if r.status == TokenUsageStatus.COMPLETED)
        interrupted_count = sum(1 for r in records if r.status == TokenUsageStatus.INTERRUPTED)
        failed_count = sum(1 for r in records if r.status == TokenUsageStatus.FAILED)
        
        # 按日期分组
        daily_stats = defaultdict(lambda: {
            "promptTokens": 0,
            "completionTokens": 0,
            "totalTokens": 0,
            "count": 0
        })
        
        for r in records:
            dt = datetime.fromtimestamp(r.timestamp)
            day_key = f"{dt.year}-{dt.month:02d}-{dt.day:02d}"
            daily_stats[day_key]["promptTokens"] += r.promptTokens
            daily_stats[day_key]["completionTokens"] += r.completionTokens
            daily_stats[day_key]["totalTokens"] += r.totalTokens
            daily_stats[day_key]["count"] += 1
        
        # 按角色分组
        role_stats = defaultdict(lambda: {
            "promptTokens": 0,
            "completionTokens": 0,
            "totalTokens": 0,
            "count": 0
        })
        
        for r in records:
            role_stats[r.roleName]["promptTokens"] += r.promptTokens
            role_stats[r.roleName]["completionTokens"] += r.completionTokens
            role_stats[r.roleName]["totalTokens"] += r.totalTokens
            role_stats[r.roleName]["count"] += 1
        
        # 按聊天分组
        chat_stats = defaultdict(lambda: {
            "promptTokens": 0,
            "completionTokens": 0,
            "totalTokens": 0,
            "count": 0
        })
        
        for r in records:
            chat_key = f"{r.roleName}/{r.chatName}"
            chat_stats[chat_key]["promptTokens"] += r.promptTokens
            chat_stats[chat_key]["completionTokens"] += r.completionTokens
            chat_stats[chat_key]["totalTokens"] += r.totalTokens
            chat_stats[chat_key]["count"] += 1
        
        # ✅ 按 API URL 分组
        api_url_stats = defaultdict(lambda: {
            "promptTokens": 0,
            "completionTokens": 0,
            "totalTokens": 0,
            "count": 0
        })
        
        for r in records:
            if r.apiUrl:
                api_url_stats[r.apiUrl]["promptTokens"] += r.promptTokens
                api_url_stats[r.apiUrl]["completionTokens"] += r.completionTokens
                api_url_stats[r.apiUrl]["totalTokens"] += r.totalTokens
                api_url_stats[r.apiUrl]["count"] += 1
        
        return {
            "year": year,
            "month": month,
            "totalRecords": len(records),
            "totalPromptTokens": total_prompt,
            "totalCompletionTokens": total_completion,
            "totalTokens": total_tokens,
            "completedCount": completed_count,
            "interruptedCount": interrupted_count,
            "failedCount": failed_count,
            "dailyStats": dict(daily_stats),
            "roleStats": dict(role_stats),
            "chatStats": dict(chat_stats),
            "apiUrlStats": dict(api_url_stats),  # ✅ 新增
            "records": [r.model_dump() for r in records[:100]]  # 最近100条记录
        }
    
    async def list_months(self) -> List[Dict[str, int]]:
        """列出所有有数据的月份"""
        months = []
        
        if not self.token_usage_dir.exists():
            return months
        
        for year_dir in sorted(self.token_usage_dir.iterdir()):
            if year_dir.is_dir() and year_dir.name.isdigit():
                year = int(year_dir.name)
                for month_file in sorted(year_dir.glob("*.jsonl")):
                    month = int(month_file.stem)
                    months.append({"year": year, "month": month})
        
        return months
    
    async def get_api_url_stats(self) -> Dict[str, Any]:
        """
        ✅ 获取按 API URL 分组的统计数据（从索引文件快速读取）
        
        Returns:
            {api_url: {totalPromptTokens, totalCompletionTokens, totalTokens, count, firstUsed, lastUsed}}
        """
        api_url_index = self.index_dir / "api_urls.json"
        
        if not api_url_index.exists():
            return {}
        
        try:
            with open(api_url_index, 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception as e:
            print(f"[TokenUsage] 读取 API URL 索引失败: {e}")
            return {}
    
    async def get_daily_stats(self, year: int, month: int) -> Dict[str, Any]:
        """
        ✅ 获取指定月份的每日统计数据（从索引文件快速读取）
        
        Args:
            year: 年份
            month: 月份
            
        Returns:
            {day_key: {promptTokens, completionTokens, totalTokens, count}}
        """
        daily_index = self.index_dir / "daily" / f"{year}-{month:02d}.json"
        
        if not daily_index.exists():
            return {}
        
        try:
            with open(daily_index, 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception as e:
            print(f"[TokenUsage] 读取每日索引失败: {e}")
            return {}
    
    async def get_available_roles(self, year: int, month: int) -> List[str]:
        """获取指定月份有数据的角色列表"""
        records = self._load_month_records(year, month)
        roles = set(r.roleName for r in records)
        return sorted(list(roles))
    
    async def get_available_chats(
        self,
        year: int,
        month: int,
        role_name: Optional[str] = None
    ) -> List[str]:
        """获取指定月份有数据的聊天列表"""
        records = self._load_month_records(year, month)
        
        if role_name:
            records = [r for r in records if r.roleName == role_name]
        
        chats = set(f"{r.roleName}/{r.chatName}" for r in records)
        return sorted(list(chats))


# 全局实例
token_usage_service = TokenUsageService()
