kards-env/kards_battle/units/unit_loader.py

"""
YAML单位加载器
从YAML文件加载单位定义，将游戏引擎与数据分离
"""
import yaml
import os
from typing import Dict, List, Any, Optional, Set
from pathlib import Path

from .unit import Unit
from ..core.enums import UnitType


class UnitLoader:
    """YAML单位加载器"""
    
    def __init__(self, assets_path: Optional[str] = None):
        """初始化加载器"""
        if assets_path is None:
            # 默认使用项目根目录下的assets文件夹
            project_root = Path(__file__).parent.parent.parent
            assets_path = project_root / "assets" / "units"
        
        self.assets_path = Path(assets_path)
        self._unit_definitions: Dict[str, Dict[str, Any]] = {}
        self._loaded_files: Set[str] = set()
    
    def load_units_from_file(self, filename: str) -> Dict[str, Dict[str, Any]]:
        """从单个YAML文件加载单位定义"""
        file_path = self.assets_path / filename
        
        if not file_path.exists():
            raise FileNotFoundError(f"Unit file not found: {file_path}")
        
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                data = yaml.safe_load(f)
            
            if not data or 'units' not in data:
                raise ValueError(f"Invalid unit file format: {filename}")
            
            units_data = {}
            for unit_data in data['units']:
                if 'id' not in unit_data:
                    raise ValueError(f"Unit missing 'id' field in {filename}")
                
                unit_id = unit_data['id']
                units_data[unit_id] = unit_data
            
            self._unit_definitions.update(units_data)
            self._loaded_files.add(filename)
            
            return units_data
            
        except yaml.YAMLError as e:
            raise ValueError(f"Error parsing YAML file {filename}: {e}")
    
    def load_all_units(self) -> Dict[str, Dict[str, Any]]:
        """加载assets目录下所有YAML单位文件"""
        if not self.assets_path.exists():
            raise FileNotFoundError(f"Assets directory not found: {self.assets_path}")
        
        yaml_files = list(self.assets_path.glob("*.yaml")) + list(self.assets_path.glob("*.yml"))
        
        # 排除规范文件
        yaml_files = [f for f in yaml_files if not f.name.startswith("units_spec")]
        
        for yaml_file in yaml_files:
            if yaml_file.name not in self._loaded_files:
                self.load_units_from_file(yaml_file.name)
        
        return self._unit_definitions
    
    def get_unit_definition(self, unit_id: str) -> Optional[Dict[str, Any]]:
        """获取单位定义"""
        return self._unit_definitions.get(unit_id)
    
    def list_available_units(self) -> List[str]:
        """列出所有可用的单位ID"""
        return list(self._unit_definitions.keys())
    
    def create_unit_from_id(self, unit_id: str) -> Unit:
        """根据单位ID创建Unit实例"""
        unit_data = self.get_unit_definition(unit_id)
        if not unit_data:
            raise ValueError(f"Unit definition not found: {unit_id}")
        
        return self._create_unit_from_data(unit_data)
    
    def create_unit_from_data(self, unit_data: Dict[str, Any]) -> Unit:
        """直接从数据字典创建Unit实例"""
        return self._create_unit_from_data(unit_data)
    
    def _create_unit_from_data(self, unit_data: Dict[str, Any]) -> Unit:
        """内部方法：从数据创建Unit实例"""
        # 验证必需字段
        required_fields = ['name', 'type', 'stats']
        for field in required_fields:
            if field not in unit_data:
                raise ValueError(f"Missing required field: {field}")
        
        # 解析单位类型
        unit_type_str = unit_data['type'].upper()
        try:
            unit_type = UnitType[unit_type_str]
        except KeyError:
            raise ValueError(f"Invalid unit type: {unit_data['type']}")
        
        # 解析属性
        stats = unit_data['stats']
        required_stats = ['attack', 'defense', 'operation_cost']
        for stat in required_stats:
            if stat not in stats:
                raise ValueError(f"Missing required stat: {stat}")
        
        # 解析关键词
        keywords = set(unit_data.get('keywords', []))
        
        # 解析能力（暂时作为原始数据存储，后续可扩展）
        abilities = unit_data.get('abilities', [])
        
        # 创建单位实例
        unit = Unit(
            name=unit_data['name'],
            unit_type=unit_type,
            attack=stats['attack'],
            defense=stats['defense'],
            operation_cost=stats['operation_cost'],
            keywords=keywords,
            abilities=abilities
        )
        
        # 设置额外属性
        if 'nation' in unit_data:
            unit.nation = unit_data['nation']
        
        if 'id' in unit_data:
            unit.definition_id = unit_data['id']
        
        return unit
    
    def get_units_by_nation(self, nation: str) -> List[Dict[str, Any]]:
        """获取指定国家的所有单位定义"""
        return [
            unit_data for unit_data in self._unit_definitions.values()
            if unit_data.get('nation', '').upper() == nation.upper()
        ]
    
    def get_units_by_type(self, unit_type: str) -> List[Dict[str, Any]]:
        """获取指定类型的所有单位定义"""
        return [
            unit_data for unit_data in self._unit_definitions.values()
            if unit_data.get('type', '').upper() == unit_type.upper()
        ]


# 全局单例加载器实例
_global_loader: Optional[UnitLoader] = None


def get_unit_loader() -> UnitLoader:
    """获取全局单位加载器实例"""
    global _global_loader
    if _global_loader is None:
        _global_loader = UnitLoader()
    return _global_loader


def load_unit(unit_id: str) -> Unit:
    """便捷函数：加载单位"""
    loader = get_unit_loader()
    
    # 如果还没有加载单位定义，先加载
    if not loader._unit_definitions:
        loader.load_all_units()
    
    return loader.create_unit_from_id(unit_id)


def list_all_units() -> List[str]:
    """便捷函数：列出所有可用单位"""
    loader = get_unit_loader()
    
    # 如果还没有加载单位定义，先加载
    if not loader._unit_definitions:
        loader.load_all_units()
    
    return loader.list_available_units()