TFT-Strategist/src/data_provider/data_query_api.py

from typing import Dict, List, Optional, Any, Set, Tuple, Union
import logging
from .data_loader import DataLoader

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger("TFT-Strategist-DataQueryAPI")

class DataQueryAPI:
    """
    数据查询API，提供对游戏数据的各种查询功能
    """
    
    def __init__(self, data_loader: Optional[DataLoader] = None):
        """
        初始化数据查询API
        
        Args:
            data_loader: 数据加载器实例，如果为None则创建一个新的实例
        """
        self.data_loader = data_loader if data_loader else DataLoader()
        self._job_cache = {}  # 职业缓存
        self._race_cache = {}  # 特质缓存
        self._chess_cache = {}  # 棋子缓存
        self._job_chess_map = {}  # 职业->棋子映射
        self._race_chess_map = {}  # 特质->棋子映射
        self._chess_job_map = {}  # 棋子->职业映射
        self._chess_race_map = {}  # 棋子->特质映射
        
        # 初始化数据
        self._init_data()
    
    def _init_data(self) -> bool:
        """
        初始化并处理数据
        
        Returns:
            bool: 初始化是否成功
        """
        if not self.data_loader.load_all_data():
            logger.error("数据加载失败，无法初始化数据查询API")
            return False
        
        # 处理职业数据
        job_data = self.data_loader.get_data('job')
        if job_data:
            for job in job_data['data']:
                job_id = job['jobId']
                self._job_cache[job_id] = job
                self._job_chess_map[job_id] = []
        
        # 处理特质数据
        race_data = self.data_loader.get_data('race')
        if race_data:
            for race in race_data['data']:
                race_id = race['raceId']
                self._race_cache[race_id] = race
                self._race_chess_map[race_id] = []
        
        # 处理棋子数据并建立关联
        chess_data = self.data_loader.get_data('chess')
        if chess_data:
            for chess in chess_data['data']:
                chess_id = chess['chessId']
                self._chess_cache[chess_id] = chess
                
                # 建立棋子与职业的关联
                job_ids = chess['jobIds'].split(',') if chess['jobIds'] else []
                self._chess_job_map[chess_id] = job_ids
                for job_id in job_ids:
                    if job_id and job_id in self._job_chess_map:
                        self._job_chess_map[job_id].append(chess_id)
                
                # 建立棋子与特质的关联
                race_ids = chess['raceIds'].split(',') if chess['raceIds'] else []
                self._chess_race_map[chess_id] = race_ids
                for race_id in race_ids:
                    if race_id and race_id in self._race_chess_map:
                        self._race_chess_map[race_id].append(chess_id)
        
        logger.info(f"数据初始化完成: {len(self._chess_cache)}个棋子, {len(self._job_cache)}个职业, {len(self._race_cache)}个特质")
        return True
    
    def reload_data(self, force: bool = False) -> bool:
        """
        重新加载数据
        
        Args:
            force: 是否强制从在线接口获取
            
        Returns:
            bool: 重新加载是否成功
        """
        if self.data_loader.reload_data(force=force):
            # 清空缓存
            self._job_cache.clear()
            self._race_cache.clear()
            self._chess_cache.clear()
            self._job_chess_map.clear()
            self._race_chess_map.clear()
            self._chess_job_map.clear()
            self._chess_race_map.clear()
            
            # 重新初始化数据
            return self._init_data()
        return False
    
    def get_all_jobs(self) -> List[Dict[str, Any]]:
        """
        获取所有职业数据
        
        Returns:
            List[Dict[str, Any]]: 职业数据列表
        """
        return list(self._job_cache.values())
    
    def get_all_races(self) -> List[Dict[str, Any]]:
        """
        获取所有特质数据
        
        Returns:
            List[Dict[str, Any]]: 特质数据列表
        """
        return list(self._race_cache.values())
    
    def get_all_chess(self) -> List[Dict[str, Any]]:
        """
        获取所有棋子数据
        
        Returns:
            List[Dict[str, Any]]: 棋子数据列表
        """
        return list(self._chess_cache.values())
    
    def get_job_by_id(self, job_id: str) -> Optional[Dict[str, Any]]:
        """
        根据ID获取职业数据
        
        Args:
            job_id: 职业ID
            
        Returns:
            Optional[Dict[str, Any]]: 职业数据，如果不存在则返回None
        """
        # 确保job_id是字符串
        job_id_str = str(job_id)
        return self._job_cache.get(job_id_str)
    
    def get_race_by_id(self, race_id: str) -> Optional[Dict[str, Any]]:
        """
        根据ID获取特质数据
        
        Args:
            race_id: 特质ID
            
        Returns:
            Optional[Dict[str, Any]]: 特质数据，如果不存在则返回None
        """
        # 确保race_id是字符串
        race_id_str = str(race_id)
        return self._race_cache.get(race_id_str)
    
    def get_chess_by_id(self, chess_id: str) -> Optional[Dict[str, Any]]:
        """
        根据ID获取棋子数据
        
        Args:
            chess_id: 棋子ID
            
        Returns:
            Optional[Dict[str, Any]]: 棋子数据，如果不存在则返回None
        """
        # 确保chess_id是字符串
        chess_id_str = str(chess_id)
        return self._chess_cache.get(chess_id_str)
    
    def get_job_by_name(self, name: str) -> Optional[Dict[str, Any]]:
        """
        根据名称获取职业数据
        
        Args:
            name: 职业名称
            
        Returns:
            Optional[Dict[str, Any]]: 职业数据，如果不存在则返回None
        """
        for job in self._job_cache.values():
            if job['name'] == name:
                return job
        return None
    
    def get_race_by_name(self, name: str) -> Optional[Dict[str, Any]]:
        """
        根据名称获取特质数据
        
        Args:
            name: 特质名称
            
        Returns:
            Optional[Dict[str, Any]]: 特质数据，如果不存在则返回None
        """
        for race in self._race_cache.values():
            if race['name'] == name:
                return race
        return None
    
    def get_chess_by_name(self, name: str) -> Optional[Dict[str, Any]]:
        """
        根据名称获取棋子数据
        
        Args:
            name: 棋子名称（displayName）
            
        Returns:
            Optional[Dict[str, Any]]: 棋子数据，如果不存在则返回None
        """
        for chess in self._chess_cache.values():
            if chess.get('displayName') == name:
                return chess
        return None
    
    def get_chess_by_job(self, job_id: str) -> List[Dict[str, Any]]:
        """
        获取指定职业的所有棋子
        
        Args:
            job_id: 职业ID
            
        Returns:
            List[Dict[str, Any]]: 棋子数据列表
        """
        if job_id not in self._job_chess_map:
            return []
        
        result = []
        for chess_id in self._job_chess_map[job_id]:
            chess = self._chess_cache.get(chess_id)
            if chess:
                result.append(chess)
        
        return result
    
    def get_chess_by_race(self, race_id: str) -> List[Dict[str, Any]]:
        """
        获取指定特质的所有棋子
        
        Args:
            race_id: 特质ID
            
        Returns:
            List[Dict[str, Any]]: 棋子数据列表
        """
        if race_id not in self._race_chess_map:
            return []
        
        result = []
        for chess_id in self._race_chess_map[race_id]:
            chess = self._chess_cache.get(chess_id)
            if chess:
                result.append(chess)
        
        return result
    
    def get_jobs_of_chess(self, chess_id: str) -> List[Dict[str, Any]]:
        """
        获取指定棋子的所有职业
        
        Args:
            chess_id: 棋子ID
            
        Returns:
            List[Dict[str, Any]]: 职业数据列表
        """
        if chess_id not in self._chess_job_map:
            return []
        
        result = []
        for job_id in self._chess_job_map[chess_id]:
            job = self._job_cache.get(job_id)
            if job:
                result.append(job)
        
        return result
    
    def get_races_of_chess(self, chess_id: str) -> List[Dict[str, Any]]:
        """
        获取指定棋子的所有特质
        
        Args:
            chess_id: 棋子ID
            
        Returns:
            List[Dict[str, Any]]: 特质数据列表
        """
        if chess_id not in self._chess_race_map:
            return []
        
        result = []
        for race_id in self._chess_race_map[chess_id]:
            race = self._race_cache.get(race_id)
            if race:
                result.append(race)
        
        return result
    
    def get_all_synergies(self) -> List[Dict[str, Any]]:
        """
        获取所有羁绊（职业和特质）数据
        
        Returns:
            List[Dict[str, Any]]: 羁绊数据列表
        """
        return self.get_all_jobs() + self.get_all_races()
    
    def get_synergy_by_id(self, synergy_id: str) -> Optional[Dict[str, Any]]:
        """
        根据ID获取羁绊数据
        
        Args:
            synergy_id: 羁绊ID
            
        Returns:
            Optional[Dict[str, Any]]: 羁绊数据，如果不存在则返回None
        """
        # 确保synergy_id是字符串
        synergy_id_str = str(synergy_id)
        
        # 先在职业中查找
        job = self.get_job_by_id(synergy_id_str)
        if job:
            return job
        
        # 再在特质中查找
        return self.get_race_by_id(synergy_id_str)
    
    def get_synergy_by_name(self, name: str) -> Optional[Dict[str, Any]]:
        """
        根据名称获取羁绊数据
        
        Args:
            name: 羁绊名称
            
        Returns:
            Optional[Dict[str, Any]]: 羁绊数据，如果不存在则返回None
        """
        # 先在职业中查找
        job = self.get_job_by_name(name)
        if job:
            return job
        
        # 再在特质中查找
        return self.get_race_by_name(name)
    
    def get_chess_by_synergy(self, synergy_id: str) -> List[Dict[str, Any]]:
        """
        获取指定羁绊的所有棋子
        
        Args:
            synergy_id: 羁绊ID
            
        Returns:
            List[Dict[str, Any]]: 棋子数据列表
        """
        # 确保synergy_id是字符串
        synergy_id_str = str(synergy_id)
        
        # 先在职业中查找
        chess_list = self.get_chess_by_job(synergy_id_str)
        if chess_list:
            return chess_list
        
        # 再在特质中查找
        return self.get_chess_by_race(synergy_id_str)
    
    def get_synergies_of_chess(self, chess_id: str) -> List[Dict[str, Any]]:
        """
        获取指定棋子的所有羁绊
        
        Args:
            chess_id: 棋子ID
            
        Returns:
            List[Dict[str, Any]]: 羁绊数据列表
        """
        return self.get_jobs_of_chess(chess_id) + self.get_races_of_chess(chess_id)
    
    def get_chess_cost_distribution(self) -> Dict[str, int]:
        """
        获取棋子费用分布
        
        Returns:
            Dict[str, int]: 费用分布统计，键为费用，值为数量
        """
        result = {}
        for chess in self._chess_cache.values():
            cost = chess.get('price', '0')
            if cost not in result:
                result[cost] = 0
            result[cost] += 1
        
        return result
    
    def get_chess_by_cost(self, cost: str) -> List[Dict[str, Any]]:
        """
        获取指定费用的所有棋子
        
        Args:
            cost: 棋子费用
            
        Returns:
            List[Dict[str, Any]]: 棋子数据列表
        """
        result = []
        for chess in self._chess_cache.values():
            if chess.get('price') == cost:
                result.append(chess)
        
        return result
    
    def get_synergy_levels(self, synergy_id: str) -> Dict[str, str]:
        """
        获取指定羁绊的等级信息
        
        Args:
            synergy_id: 羁绊ID
            
        Returns:
            Dict[str, str]: 羁绊等级信息，键为等级，值为效果描述
        """
        # 确保synergy_id是字符串
        synergy_id_str = str(synergy_id)
        synergy = self.get_synergy_by_id(synergy_id_str)
        if not synergy:
            return {}
        
        return synergy.get('level', {})


if __name__ == "__main__":
    # 测试代码
    api = DataQueryAPI()
    
    # 获取所有职业和特质
    print(f"职业数量: {len(api.get_all_jobs())}")
    print(f"特质数量: {len(api.get_all_races())}")
    print(f"棋子数量: {len(api.get_all_chess())}")
    
    # 获取特定羁绊的棋子
    heavy_warriors = api.get_job_by_name("重装战士")
    if heavy_warriors:
        heavy_warriors_id = heavy_warriors['jobId']
        chess_list = api.get_chess_by_job(heavy_warriors_id)
        print(f"\n重装战士棋子 ({len(chess_list)}个):")
        for chess in chess_list:
            print(f"  - {chess['displayName']} (费用: {chess['price']})")
    
    # 获取某个棋子的所有羁绊
    some_chess = api.get_all_chess()[0]
    if some_chess:
        chess_id = some_chess['chessId']
        synergies = api.get_synergies_of_chess(chess_id)
        print(f"\n{some_chess['displayName']}的羁绊:")
        for synergy in synergies:
            print(f"  - {synergy['name']}")
    
    # 获取费用分布
    cost_distribution = api.get_chess_cost_distribution()
    print("\n棋子费用分布:")
    for cost, count in sorted(cost_distribution.items()):
        if cost != '0':  # 排除费用为0的棋子（通常是召唤物）
            print(f"  {cost}费棋子: {count}个")