"""Schema生成模块"""

from typing import Dict, List
from lxml import etree as ET


class SchemaGenerator:
    """Schema生成工具
    
    功能：
    1. 从XML结构提取字段定义
    2. 生成用于AI约束的Schema
    3. 支持多种Schema格式
    """
    
    def __init__(self):
        """初始化Schema生成器"""
        pass
    
    def generate_from_xml(self, xml_content: str, format: str = 'simple') -> Dict[str, str]:
        """从XML生成Schema
        
        Args:
            xml_content: XML内容字符串
            format: Schema格式 ('simple', 'json_schema', 'yaml')
            
        Returns:
            Dict[str, str]: 字段定义字典
        """
        # 解析XML
        try:
            # 尝试直接解析
            root = ET.fromstring(xml_content)
        except ValueError:
            # 如果是Unicode字符串带有编码声明，转换为字节串
            if isinstance(xml_content, str):
                xml_content = xml_content.encode('utf-8')
            root = ET.fromstring(xml_content)
        
        # 提取字段定义
        field_definitions = self._extract_fields(root)
        
        # 根据格式生成Schema
        if format == 'json_schema':
            return self._generate_json_schema(field_definitions)
        elif format == 'yaml':
            return self._generate_yaml_schema(field_definitions)
        else:
            # 默认返回简单格式
            return field_definitions
    
    def generate_from_file(self, xml_file: str, format: str = 'simple') -> Dict[str, str]:
        """从XML文件生成Schema
        
        Args:
            xml_file: XML文件路径
            format: Schema格式
            
        Returns:
            Dict[str, str]: 字段定义字典
        """
        # 读取XML文件，尝试不同编码
        encodings = ['utf-8', 'gbk', 'ansi']
        xml_content = None
        
        for encoding in encodings:
            try:
                with open(xml_file, 'r', encoding=encoding) as f:
                    xml_content = f.read()
                break
            except UnicodeDecodeError:
                continue
        
        if xml_content is None:
            raise Exception("无法读取文件：所有编码尝试都失败")
        
        return self.generate_from_xml(xml_content, format)
    
    def _extract_fields(self, element: ET.Element, parent_path: str = '') -> Dict[str, str]:
        """递归提取字段定义
        
        Args:
            element: XML元素
            parent_path: 父元素路径
            
        Returns:
            Dict[str, str]: 字段定义字典
        """
        field_definitions = {}
        
        # 构建当前路径
        if parent_path:
            current_path = f"{parent_path}/{element.tag}"
        else:
            current_path = element.tag
        
        # 检查当前元素是否有文本内容
        if element.text and element.text.strip():
            # 推断字段类型
            field_type = self._infer_field_type(element.text.strip())
            field_definitions[current_path] = field_type
        
        # 递归处理子元素
        for child in element:
            child_definitions = self._extract_fields(child, current_path)
            field_definitions.update(child_definitions)
        
        return field_definitions
    
    def _infer_field_type(self, text: str) -> str:
        """推断字段类型
        
        Args:
            text: 字段文本内容
            
        Returns:
            str: 推断的字段类型
        """
        # 简单的类型推断
        text = text.strip()
        
        # 检查是否为数字
        if text.replace('.', '', 1).replace('-', '', 1).isdigit():
            if '.' in text:
                return 'number'
            else:
                return 'integer'
        
        # 检查是否为日期
        date_formats = ['%Y-%m-%d', '%Y/%m/%d', '%d-%m-%Y', '%d/%m/%Y']
        for fmt in date_formats:
            try:
                import datetime
                datetime.datetime.strptime(text, fmt)
                return 'date'
            except ValueError:
                pass
        
        # 检查是否为布尔值
        if text.lower() in ['true', 'false', 'yes', 'no', '1', '0']:
            return 'boolean'
        
        # 检查是否为URL
        import re
        url_pattern = re.compile(r'https?://\S+')
        if url_pattern.match(text):
            return 'url'
        
        # 检查是否为Markdown
        if any(marker in text for marker in ['# ', '## ', '**', '*', '- ', '```']):
            return 'markdown'
        
        # 默认类型
        return 'string'
    
    def _generate_json_schema(self, field_definitions: Dict[str, str]) -> Dict:
        """生成JSON Schema
        
        Args:
            field_definitions: 字段定义字典
            
        Returns:
            Dict: JSON Schema字典
        """
        json_schema = {
            "type": "object",
            "properties": {}
        }
        
        # 转换字段定义为JSON Schema
        for field_path, field_type in field_definitions.items():
            # 简化字段名（只使用最后一部分）
            field_name = field_path.split('/')[-1]
            
            # 映射类型到JSON Schema类型
            type_mapping = {
                'string': 'string',
                'number': 'number',
                'integer': 'integer',
                'boolean': 'boolean',
                'date': 'string',
                'url': 'string',
                'markdown': 'string'
            }
            
            json_type = type_mapping.get(field_type, 'string')
            
            # 添加字段到properties
            json_schema["properties"][field_name] = {
                "type": json_type
            }
            
            # 添加格式约束
            if field_type == 'date':
                json_schema["properties"][field_name]["format"] = "date"
            elif field_type == 'url':
                json_schema["properties"][field_name]["format"] = "uri"
            elif field_type == 'markdown':
                json_schema["properties"][field_name]["description"] = "Markdown content"
        
        return json_schema
    
    def _generate_yaml_schema(self, field_definitions: Dict[str, str]) -> str:
        """生成YAML格式的Schema
        
        Args:
            field_definitions: 字段定义字典
            
        Returns:
            str: YAML格式的Schema字符串
        """
        yaml_lines = []
        yaml_lines.append("# Schema Definition")
        yaml_lines.append("")
        
        for field_path, field_type in field_definitions.items():
            # 简化字段名
            field_name = field_path.split('/')[-1]
            yaml_lines.append(f"{field_name}: {field_type}")
        
        return '\n'.join(yaml_lines)
