from __future__ import annotations
from typing import TYPE_CHECKING, Any, Iterator
if TYPE_CHECKING:
    from lib.core.translator import Translator

import ast
import traceback

from lib.core.Handles.HandlesBase import CTypeInfo
from lib.includes import t
from lib.core.VLogger import get_logger as _vlog
from lib.core.DiagnosticCollector import DiagnosticCollector
from lib.core.LLVMTypeMapper import LLVMTypeMapper
from lib.core.SymbolExtractor import ASTSymbolExtractor
from lib.core.SymbolTypeResolver import TypeResolver
from lib.core.SymbolInserter import SymbolInserter
from lib.core.SymbolReexporter import PackageReexporter
from lib.core.ConstEvaluator import ConstEvaluator
from lib.core.SymbolData import ModuleSymbols

# 从 SymbolUtils 重导出，保持向后兼容
from lib.core.SymbolUtils import CheckAnnotationHasCInline as _CheckAnnotationHasCInline


class SymbolTable:
    def __init__(self, translator: Translator) -> None:
        self.translator: Translator = translator
        self._symbols: dict[str, CTypeInfo] = {}
        self._namespaces: dict[str, dict[str, CTypeInfo]] = {}  # 模块名 -> 符号字典
        self._t_type_symbols: dict[str, type] = {}  # 别名 -> t.CType 子类 (from t import CEnum as en)
        self.diagnostics: DiagnosticCollector = DiagnosticCollector()
        self._type_mapper: LLVMTypeMapper = LLVMTypeMapper(self)
        self.import_aliases: dict[str, str] = {}  # 别名 -> 模块路径

    def __getstate__(self) -> dict:
        """pickle 序列化：排除不可序列化的字段（translator/_type_mapper/diagnostics）"""
        state: dict = self.__dict__.copy()
        state.pop('translator', None)
        state.pop('_type_mapper', None)
        state.pop('diagnostics', None)
        return state

    def __setstate__(self, state: dict) -> None:
        """pickle 反序列化：重建可序列化的字段，translator 由调用方后续设置"""
        self.__dict__.update(state)
        self.translator = None  # 由调用方赋值
        self.diagnostics = DiagnosticCollector()
        self._type_mapper = LLVMTypeMapper(self)

    def clear(self) -> None:
        self._symbols.clear()
        self._namespaces.clear()
        self._t_type_symbols.clear()
        self.import_aliases.clear()
    
    def get(self, name: str, default: CTypeInfo | None = None) -> CTypeInfo | None:
        return self._symbols.get(name, default)
    
    def set(self, name: str, value: CTypeInfo | dict[str, Any]) -> None:
        if isinstance(value, CTypeInfo):
            self._symbols[name] = value
        else:
            # 兼容旧代码传入 dict 的情况
            info: CTypeInfo = self._CTypeInfoFromDict(name, value)
            self._symbols[name] = info
    
    def update(self, symbols: dict[str, CTypeInfo | dict[str, Any]]) -> None:
        for name, value in symbols.items():
            self.set(name, value)
    
    def keys(self) -> Iterator[str]:
        return iter(self._symbols.keys())
    
    def values(self) -> Iterator[CTypeInfo]:
        return iter(self._symbols.values())
    
    def items(self) -> Iterator[tuple[str, CTypeInfo]]:
        return iter(self._symbols.items())
    
    def pop(self, name: str, *args: Any) -> CTypeInfo:
        return self._symbols.pop(name, *args)
    
    def __getitem__(self, name: str) -> CTypeInfo:
        if name in self._symbols:
            return self._symbols[name]
        raise KeyError(name)
    
    def __setitem__(self, name: str, value: CTypeInfo | dict[str, Any]) -> None:
        self.set(name, value)
    
    def __delitem__(self, name: str) -> None:
        del self._symbols[name]
    
    def __contains__(self, name: str) -> bool:
        return name in self._symbols
    
    def __len__(self) -> int:
        return len(self._symbols)
    
    def __iter__(self) -> Iterator[str]:
        return iter(self._symbols)

    # ==================================================================
    #  服务接口 — 语义化的符号查询与操作方法
    # ==================================================================

    def lookup(self, name: str) -> CTypeInfo | None:
        """查找符号，支持别名解析和 dotted name

        解析顺序：
        1. 直接匹配 _symbols
        2. 如果包含 '.'，解析第一个段为别名后查找剩余路径
        3. 查 _t_type_symbols（t/c 类型别名，返回模拟 CTypeInfo）
        """
        if name in self._symbols:
            return self._symbols[name]

        if '.' not in name:
            return self._lookup_t_type(name)

        # dotted name: 尝试解析第一个段为别名
        parts: list[str] = name.split('.', 1)
        first: str = parts[0]
        rest: str = parts[1]

        # 1) 别名解析：import X as Y → Y.Symbol → X.Symbol
        resolved: str | None = self.import_aliases.get(first)
        if resolved:
            alias_name: str = f"{resolved}.{rest}"
            if alias_name in self._symbols:
                return self._symbols[alias_name]

        # 2) 纯路径查找：module.ClassName
        # （已在 line 100 检查过，此处为兜底，保留注释说明路径查找语义）

        return None

    def _lookup_t_type(self, name: str) -> CTypeInfo | None:
        """在 _t_type_symbols 中查找 t/c 类型别名，返回模拟 CTypeInfo"""
        cls: type | None = self._t_type_symbols.get(name)
        if cls is None:
            return None
        info: CTypeInfo = CTypeInfo()
        info.Name = name
        info.BaseType = cls()
        if issubclass(cls, (t.CEnum, t.REnum)):
            info.IsEnum = True
        elif issubclass(cls, t.CUnion):
            info.IsUnion = True
        elif issubclass(cls, t.CStruct):
            info.IsStruct = True
        return info

    def register_namespace(self, module_name: str) -> None:
        """注册一个模块命名空间"""
        if module_name not in self._namespaces:
            self._namespaces[module_name] = {}

    def has(self, name: str) -> bool:
        """检查符号是否存在（替代 `name in SymbolTable`）"""
        return name in self._symbols

    def is_struct(self, name: str) -> bool:
        """检查名称是否为结构体"""
        entry: CTypeInfo | None = self._symbols.get(name)
        return entry is not None and (entry.IsStruct or entry.IsCpythonObject)

    def is_define(self, name: str) -> bool:
        """检查名称是否为常量定义"""
        entry: CTypeInfo | None = self._symbols.get(name)
        return entry is not None and entry.IsDefine

    def insert(self, name: str, info: CTypeInfo) -> None:
        """插入符号（替代 `SymbolTable[name] = info`）"""
        self._symbols[name] = info

    def resolve_alias(self, module_path: str) -> str:
        """解析 import 别名，无别名返回原路径"""
        return self.import_aliases.get(module_path, module_path)
    
    def ToDict(self) -> dict[str, Any]:
        """序列化为字典格式，BaseType 使用可重建的 dict 格式

        返回的 dict 包含：
        - 符号条目：{name: {attrs}}（与旧格式兼容）
        - '__namespaces__': {namespace: [symbol_names]}  命名空间索引
        - '__import_aliases__': {alias: module_path}      导入别名
        """
        result: dict[str, Any] = {}
        for name, info in self._symbols.items():
            entry: dict[str, Any] = {'name': name}
            if info.BaseType:
                entry['BaseType'] = self._SerializeBaseType(info.BaseType)
            if info.PtrCount:
                entry['PtrCount'] = info.PtrCount
            if info.IsTypedef:
                entry['IsTypedef'] = True
            if info.IsStruct:
                entry['IsStruct'] = True
            if info.IsEnum:
                entry['IsEnum'] = True
            if info.IsUnion:
                entry['IsUnion'] = True
            if info.IsFunction:
                entry['IsFunction'] = True
            if info.IsVariable:
                entry['IsVariable'] = True
            if info.IsDefine:
                entry['IsDefine'] = True
            if info.OriginalType:
                entry['OriginalType'] = str(info.OriginalType)
            if info.Lineno:
                entry['lineno'] = info.Lineno
            if info.file:
                entry['file'] = info.file
            if info.Members:
                entry['members'] = info.Members
            extra: dict[str, Any] = info._sm._extra
            if extra:
                entry.update(extra)
            result[name] = entry

        # 序列化命名空间信息
        if self._namespaces:
            result['__namespaces__'] = {
                ns: list(symbols.keys()) for ns, symbols in self._namespaces.items()
            }
        if self.import_aliases:
            result['__import_aliases__'] = dict(self.import_aliases)
        if self._t_type_symbols:
            result['__t_type_symbols__'] = {
                name: cls.__name__ for name, cls in self._t_type_symbols.items()
            }

        return result

    @staticmethod
    def _SerializeBaseType(base_type: t.CType | tuple | list | str | None) -> dict | str:
        """将 BaseType 序列化为可重建的格式"""
        if isinstance(base_type, t.CType):
            return {'class': type(base_type).__name__}
        if isinstance(base_type, (tuple, list)):
            return {'classes': [type(b).__name__ for b in base_type if isinstance(b, t.CType)]}
        return str(base_type)

    @staticmethod
    def _DeserializeBaseType(data: Any) -> Any:
        """从序列化格式恢复 BaseType"""
        if isinstance(data, dict):
            if 'class' in data:
                cls: type | None = getattr(t, data['class'], None)
                if cls and isinstance(cls, type) and issubclass(cls, t.CType):
                    return cls()
            elif 'classes' in data:
                types: list[t.CType] = []
                for cname in data['classes']:
                    cls: type | None = getattr(t, cname, None)
                    if cls and isinstance(cls, type) and issubclass(cls, t.CType):
                        types.append(cls())
                if types:
                    return tuple(types)
        elif isinstance(data, str):
            cls: type | None = getattr(t, data, None)
            if cls and isinstance(cls, type) and issubclass(cls, t.CType):
                return cls()
        return None
    
    def FromDict(self, symbols: dict[str, Any]) -> None:
        """从字典格式反序列化（兼容带/不带命名空间信息的两种格式）"""
        self._symbols.clear()
        self._namespaces.clear()
        self.import_aliases.clear()
        self._t_type_symbols.clear()

        _SPECIAL_KEYS: frozenset[str] = frozenset({'__namespaces__', '__import_aliases__', '__t_type_symbols__'})

        for name, attrs in symbols.items():
            if name in _SPECIAL_KEYS:
                continue
            if isinstance(attrs, dict):
                info: CTypeInfo = self._CTypeInfoFromDict(name, attrs)
                self._symbols[name] = info
            elif isinstance(attrs, CTypeInfo):
                self._symbols[name] = attrs

        # 恢复命名空间信息
        ns_data: dict[str, list[str]] | None = symbols.get('__namespaces__')
        if ns_data:
            for ns, sym_names in ns_data.items():
                self._namespaces[ns] = {}
                for sym_name in sym_names:
                    if sym_name in self._symbols:
                        self._namespaces[ns][sym_name] = self._symbols[sym_name]
        aliases_data: dict[str, str] | None = symbols.get('__import_aliases__')
        if aliases_data:
            self.import_aliases.update(aliases_data)

        # 恢复 t 类型别名
        t_type_data: dict[str, str] | None = symbols.get('__t_type_symbols__')
        if t_type_data:
            for name, cls_name in t_type_data.items():
                cls: type | None = getattr(t, cls_name, None)
                if cls and isinstance(cls, type) and issubclass(cls, t.CType):
                    self._t_type_symbols[name] = cls
    
    def _CTypeInfoFromDict(self, name: str, attrs: dict[str, Any]) -> CTypeInfo:
        """从旧 dict 格式创建 CTypeInfo（兼容旧代码）"""
        info: CTypeInfo = CTypeInfo()
        info.Name = name
        node_type: str = attrs.get('type', '')
        
        if node_type == 'struct' or attrs.get('IsStruct'):
            info.IsStruct = True
        elif node_type == 'enum' or attrs.get('IsEnum'):
            info.IsEnum = True
        elif node_type == 'union' or attrs.get('IsUnion'):
            info.IsUnion = True
        elif node_type == 'typedef' or attrs.get('IsTypedef'):
            info.IsTypedef = True
        elif node_type == 'function' or attrs.get('IsFunction'):
            info.IsFunction = True
        elif node_type == 'variable' or attrs.get('IsVariable'):
            info.IsVariable = True
        elif node_type == 'define' or attrs.get('IsDefine'):
            info.IsDefine = True
        elif node_type == 'enum_member' or attrs.get('IsEnumMember'):
            info.IsEnumMember = True
        
        if 'PtrCount' in attrs:
            info.PtrCount = attrs['PtrCount']
        if 'OriginalType' in attrs:
            info.OriginalType = attrs['OriginalType']
        if 'lineno' in attrs:
            info.Lineno = attrs['lineno']
        if 'file' in attrs:
            info.file = attrs['file']
        if 'members' in attrs:
            info.Members = attrs['members']
        if 'IsPtr' in attrs:
            info.IsPtr = attrs['IsPtr']
        if 'dims' in attrs:
            info.ArrayDims = attrs['dims']
        if 'IsCpythonObject' in attrs:
            info.IsCpythonObject = attrs['IsCpythonObject']
        if 'IsAnonymous' in attrs:
            info.IsAnonymous = attrs['IsAnonymous']
        if 'IsPacked' in attrs:
            info.IsPacked = attrs['IsPacked']
        if 'EnumName' in attrs:
            info.EnumName = attrs['EnumName']
        if 'DefineValue' in attrs:
            info.DefineValue = attrs['DefineValue']
        if 'BaseType' in attrs:
            info.BaseType = self._DeserializeBaseType(attrs['BaseType'])
        if 'name' in attrs and attrs['name'] != name:
            info.Name = attrs['name']
        
        # 其他属性存入 _extra
        skip_keys: frozenset[str] = frozenset({'type', 'name', 'PtrCount', 'OriginalType', 'lineno', 'file',
                      'members', 'IsPtr', 'dims', 'IsCpythonObject', 'IsAnonymous',
                      'IsPacked', 'EnumName', 'DefineValue', 'BaseType',
                      'IsStruct', 'IsEnum', 'IsUnion', 'IsTypedef',
                      'IsFunction', 'IsVariable', 'IsDefine', 'IsEnumMember'})
        for key, value in attrs.items():
            if key not in skip_keys:
                info.set(key, value)
        
        return info
    
    # ==================================================================
    #  模块符号加载：四阶段管线
    # ==================================================================
    
    def LoadModuleSymbols(self, FullModulePath: str, asname: str, lineno: int = 0) -> list[str]:
        return self._LoadSymbolsFromFile(FullModulePath, asname, lineno)
    
    def _LoadSymbolsFromFile(self, FilePath: str, namespace_prefix: str | list[str], lineno: int = 0) -> list[str]:
        try:
            # 提取模块名作为命名空间
            module_name: str = namespace_prefix[0] if isinstance(namespace_prefix, list) else namespace_prefix

            # Pass 1: 从 AST 提取原始符号信息
            extractor: ASTSymbolExtractor = ASTSymbolExtractor(self)
            module_symbols: ModuleSymbols = extractor.extract(FilePath)

            # Pass 2: 解析类型信息（typedef、函数签名、define 值）
            resolver: TypeResolver = TypeResolver(self)
            resolver.resolve(module_symbols)

            # Pass 3: 插入符号到主命名空间（无前缀）
            inserter: SymbolInserter = SymbolInserter(self)
            inserter.set_namespace(module_name)
            loaded: list[str] = inserter.insert(module_symbols)

            # Pass 4: 重新导出到命名空间前缀下
            prefixes: list[str] = [namespace_prefix] if isinstance(namespace_prefix, str) else namespace_prefix
            reexporter: PackageReexporter = PackageReexporter(self)
            reexporter.set_namespace(module_name)
            loaded.extend(reexporter.reexport(module_symbols, prefixes, lineno))

            return loaded

        except Exception as e:
            _vlog().error(traceback.format_exc())
            self.diagnostics.error(FilePath, lineno, f"加载模块失败: {e}")
            return []
    
    # ==================================================================
    #  类型解析辅助方法（委托到 LLVMTypeMapper）
    # ==================================================================

    @staticmethod
    def _CheckAnnotationHasCInline(annotation_node: ast.AST | None) -> bool:
        return _CheckAnnotationHasCInline(annotation_node)

    # ==================================================================
    #  符号插入方法（供 SymbolInserter / PackageReexporter 调用）
    # ==================================================================

    def _InsertClassSymbol(self, FullName: str, TypeKind: str, lineno: int, FilePath: str, members: dict[str, CTypeInfo], IsCpythonObject: bool, IsPacked: bool = False) -> None:
        info: CTypeInfo = CTypeInfo()
        info.Lineno = lineno
        info.file = FilePath
        info.Members = members if members else {}
        
        if TypeKind == 'struct':
            info.IsStruct = True
        elif TypeKind == 'union':
            info.IsUnion = True
        elif TypeKind == 'enum':
            info.IsEnum = True
        elif TypeKind == 'renum':
            info.IsStruct = True
            info.IsRenum = True
            info.IsEnum = True
        elif TypeKind == 'exception':
            info.IsStruct = True
            info.IsExceptionClass = True
        
        if IsCpythonObject:
            info.IsCpythonObject = True
        if IsPacked:
            info.IsPacked = True
        
        self._symbols[FullName] = info
    
    def _InsertEnumMemberSymbol(self, FullName: str, EnumName: str, lineno: int, FilePath: str, value: int | None = None) -> None:
        info: CTypeInfo = CTypeInfo()
        info.IsEnumMember = True
        info.EnumName = EnumName
        info.Lineno = lineno
        info.file = FilePath
        if value is not None:
            info.value = value
        self._symbols[FullName] = info

    def _InsertTypedefSymbol(self, FullName: str, OriginalType_kind: str | None, OriginalClass: CTypeInfo | str | None, lineno: int, FilePath: str, members: dict[str, CTypeInfo] | None = None) -> None:
        info: CTypeInfo = CTypeInfo()
        info.IsTypedef = True
        info.Name = FullName
        info.Lineno = lineno
        info.file = FilePath
        if members:
            info.Members = members
        
        if isinstance(OriginalClass, CTypeInfo):
            info.OriginalType = OriginalClass
        elif OriginalType_kind == 'typedef' and OriginalClass and isinstance(OriginalClass, str) and ('*' in OriginalClass):
            OriginalType: str = OriginalClass
            if OriginalType == 'CVoid *':
                OriginalType = 'void *'
            info.OriginalType = OriginalType
        elif OriginalType_kind == 'typedef' and OriginalClass == 'void':
            info.OriginalType = 'void *'
        elif OriginalType_kind == 'typedef' and OriginalClass:
            info.OriginalType = OriginalClass
        elif OriginalType_kind and OriginalClass and isinstance(OriginalClass, str):
            info.OriginalType = f'{OriginalType_kind} {OriginalClass}'
        else:
            info.OriginalType = OriginalClass
        
        self._symbols[FullName] = info
    
    def _InsertModuleSymbol(self, name: str, lineno: int, FilePath: str) -> None:
        info: CTypeInfo = CTypeInfo()
        info.IsModuleAlias = True
        info.Lineno = lineno
        info.file = FilePath
        self._symbols[name] = info
    
    def _InsertFuncSymbol(self, FullName: str, RetType: str | CTypeInfo | None, ParamTypes: list[str], lineno: int, FilePath: str, IsVariadic: bool = False, IsInline: bool = False) -> None:
        info: CTypeInfo = CTypeInfo()
        info.IsFunction = True
        if isinstance(RetType, str):
            info.FuncPtrReturn = CTypeInfo.FromTypeName(RetType) if RetType else CTypeInfo.VoidTypeInfo()
        elif isinstance(RetType, CTypeInfo):
            info.FuncPtrReturn = RetType
        else:
            info.FuncPtrReturn = CTypeInfo.VoidTypeInfo()
        info.FuncPtrParams = [(f'arg{i}', pt) for i, pt in enumerate(ParamTypes)]
        info.IsVariadic = IsVariadic
        info.IsInline = IsInline
        if IsInline:
            info.Storage = t.CInline()
        info.Lineno = lineno
        info.file = FilePath
        self._symbols[FullName] = info

    def _InsertDefineSymbol(self, FullName: str, DefineValue: int | str | float | bool | None, lineno: int, FilePath: str) -> None:
        info: CTypeInfo = CTypeInfo()
        info.IsDefine = True
        info.DefineValue = DefineValue
        info.Lineno = lineno
        info.file = FilePath
        self._symbols[FullName] = info
    
    def _eval_const_expr(self, node: ast.AST) -> int | float | str | None:
        """计算常量表达式的值（委托到 ConstEvaluator）"""
        return ConstEvaluator.eval_with_symtab(node, self)
    
    def _InsertAnonymousSymbol(self, FullName: str, IsUnion: bool, members: dict[str, CTypeInfo], lineno: int, FilePath: str) -> None:
        info: CTypeInfo = CTypeInfo()
        if IsUnion:
            info.IsUnion = True
        else:
            info.IsStruct = True
        info.IsAnonymous = True
        info.Members = members if members else {}
        info.Lineno = lineno
        info.file = FilePath
        self._symbols[FullName] = info