"""Pass 1: ASTSymbolExtractor — 从 AST 提取原始符号信息"""
from __future__ import annotations
import ast
from typing import TYPE_CHECKING
if TYPE_CHECKING:
    from lib.core.SymbolTable import SymbolTable

from lib.core.SymbolData import (
    ModuleSymbols, ClassSymbolData, TypedefSymbolData,
    FuncSymbolData, DefineSymbolData, AnonymousTypeData
)
from lib.core.Handles.HandlesBase import CTypeInfo
from lib.core.SymbolUtils import AnnotationContainsTType, CheckAnnotationHasCInline, IsTModule, AnnotationContainsName
from lib.includes import t


class ASTSymbolExtractor:
    """从 Python 源文件的 AST 中提取所有符号信息"""

    def __init__(self, symbol_table: SymbolTable) -> None:
        self._symtab: SymbolTable = symbol_table

    def extract(self, file_path: str) -> ModuleSymbols:
        """提取文件中的所有符号信息"""
        with open(file_path, 'r', encoding='utf-8') as f:
            code: str = f.read()
        tree: ast.Module = ast.parse(code)

        result: ModuleSymbols = ModuleSymbols(file_path=file_path)

        # 第一步：收集所有类的成员和匿名类型
        class_members: dict[str, dict[str, CTypeInfo]] = {}
        for node in tree.body:
            if isinstance(node, ast.ClassDef):
                members: dict[str, CTypeInfo] = {}
                self._collect_members(node, members, result.anonymous_types)
                class_members[node.name] = members

        # 第二步：从每个顶层节点提取符号信息
        for node in tree.body:
            if isinstance(node, ast.ClassDef):
                members_info: dict[str, CTypeInfo] = class_members.get(node.name, {})
                result.classes.append(self._extract_class(node, members_info))
            elif isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name):
                self._extract_ann_assign(node, result)
            elif isinstance(node, ast.FunctionDef):
                result.functions.append(self._extract_function(node))

        return result

    # ------------------------------------------------------------------
    #  成员收集（递归处理嵌套匿名类）
    # ------------------------------------------------------------------

    def _collect_members(self, class_node: ast.ClassDef, members: dict[str, CTypeInfo], anonymous_types: dict[str, AnonymousTypeData], prefix: str = '') -> None:
        """递归收集类成员，包括嵌套匿名类型"""
        for item in class_node.body:
            if isinstance(item, ast.AnnAssign) and isinstance(item.target, ast.Name):
                member_name: str = item.target.id
                full_name: str = f'{prefix}{member_name}' if prefix else member_name
                member_type: CTypeInfo | None = CTypeInfo.FromNode(item.annotation, self._symtab)
                if member_type is None:
                    member_type = CTypeInfo()
                    member_type.BaseType = t.CInt()
                is_ptr: bool = member_type.IsPtr
                if isinstance(item.annotation, ast.BinOp):
                    if AnnotationContainsName(item.annotation, 'CPtr'):
                        is_ptr = True
                members[full_name] = member_type.Copy()

            elif isinstance(item, ast.ClassDef):
                is_anonymous: bool = any(
                    any(AnnotationContainsTType(base, 'Anonymous') for base in item.bases)
                ) if item.bases else False

                if is_anonymous:
                    is_union: bool = any(
                        any(AnnotationContainsTType(base, 'CUnion') for base in item.bases)
                    )
                    # 也检查装饰器形式 @t.CUnion
                    if not is_union and hasattr(item, 'decorator_list') and item.decorator_list:
                        for deco in item.decorator_list:
                            if isinstance(deco, ast.Attribute) and hasattr(deco.value, 'id') and IsTModule(deco.value.id, self._symtab) and deco.attr == 'CUnion':
                                is_union = True
                                break
                            elif isinstance(deco, ast.Name):
                                cls2: type | None = getattr(t, deco.id, None)
                                if cls2 is None and hasattr(self._symtab, '_t_type_symbols'):
                                    cls2 = self._symtab._t_type_symbols.get(deco.id)
                                if cls2 is not None and isinstance(cls2, type) and issubclass(cls2, t.CUnion):
                                    is_union = True
                                    break

                    nested_name: str = f'{prefix}{item.name}' if prefix else item.name
                    members[nested_name] = CTypeInfo()
                    members[nested_name].BaseType = f'union {item.name}' if is_union else f'struct {item.name}'
                    members[nested_name].PtrCount = 0
                    members[nested_name].ArrayDims = []

                    new_prefix: str = f'{prefix}{item.name}.' if prefix else f'{item.name}.'
                    self._collect_members(item, members, anonymous_types, new_prefix)

                    anonymous_members: dict[str, CTypeInfo] = {}
                    for full_name, member_info in members.items():
                        if full_name.startswith(f"{nested_name}."):
                            member_name: str = full_name[len(nested_name) + 1:]
                            anonymous_members[member_name] = member_info

                    if item.name not in anonymous_types:
                        anonymous_types[item.name] = AnonymousTypeData(
                            name=item.name,
                            is_union=is_union,
                            members=anonymous_members,
                            lineno=item.lineno
                        )

    # ------------------------------------------------------------------
    #  类符号提取
    # ------------------------------------------------------------------

    def _extract_class(self, node: ast.ClassDef, members_info: dict[str, CTypeInfo]) -> ClassSymbolData:
        """从 ClassDef 节点提取类符号数据"""
        type_kind: str = 'struct'
        is_cpython_object: bool = False
        is_packed: bool = False

        # 检查装饰器
        if hasattr(node, 'decorator_list') and node.decorator_list:
            for decorator in node.decorator_list:
                if isinstance(decorator, ast.Attribute):
                    if (hasattr(decorator.value, 'id') and
                        IsTModule(decorator.value.id, self._symtab)):
                        if decorator.attr == 'Object':
                            is_cpython_object = True
                        elif decorator.attr == 'CUnion':
                            type_kind = 'union'
                        elif decorator.attr == 'CEnum':
                            type_kind = 'enum'
                        elif decorator.attr == 'REnum':
                            type_kind = 'renum'
                        elif decorator.attr == 'CStruct':
                            type_kind = 'struct'
                elif isinstance(decorator, ast.Name):
                    # from t import CStruct / CUnion / CEnum / REnum
                    cls: type | None = getattr(t, decorator.id, None)
                    if cls is None and hasattr(self._symtab, '_t_type_symbols'):
                        cls = self._symtab._t_type_symbols.get(decorator.id)
                    if cls is not None and isinstance(cls, type) and issubclass(cls, t.CType):
                        if issubclass(cls, t.CUnion):
                            type_kind = 'union'
                        elif issubclass(cls, t.REnum):
                            type_kind = 'renum'
                        elif issubclass(cls, t.CEnum):
                            type_kind = 'enum'
                        elif issubclass(cls, t.CStruct):
                            type_kind = 'struct'
                        elif decorator.id == 'Object':
                            is_cpython_object = True
                elif isinstance(decorator, ast.Call):
                    if isinstance(decorator.func, ast.Attribute):
                        if hasattr(decorator.func.value, 'id') and decorator.func.value.id == 'c' and decorator.func.attr == 'Attribute':
                            for arg in decorator.args:
                                if isinstance(arg, ast.Attribute):
                                    if isinstance(arg.value, ast.Attribute):
                                        if hasattr(arg.value.value, 'id') and IsTModule(arg.value.value.id, self._symtab) and arg.value.attr == 'attr' and arg.attr == 'packed':
                                            is_packed = True

        # 检查基类
        for base in node.bases:
            if AnnotationContainsTType(base, 'CUnion'):
                type_kind = 'union'
                break
            elif AnnotationContainsTType(base, 'CEnum'):
                type_kind = 'enum'
                break
            elif AnnotationContainsTType(base, 'REnum'):
                type_kind = 'renum'
                break
            elif AnnotationContainsTType(base, 'CStruct'):
                type_kind = 'struct'
                break
            elif AnnotationContainsTType(base, 'Object'):
                is_cpython_object = True
                type_kind = 'struct'
                break
            elif isinstance(base, ast.Name) and base.id == 'Exception':
                type_kind = 'exception'
                break
            elif isinstance(base, ast.Name):
                base_entry: CTypeInfo | None = self._symtab.get(base.id)
                if base_entry and base_entry.IsExceptionClass:
                    type_kind = 'exception'
                    break

        # 提取枚举成员
        enum_members: list[tuple[str, int, int]] = []
        if type_kind == 'enum':
            next_enum_value: int = 0
            for item in node.body:
                if isinstance(item, ast.Assign):
                    if len(item.targets) == 1 and isinstance(item.targets[0], ast.Name):
                        member_name: str = item.targets[0].id
                        if isinstance(item.value, ast.Constant) and isinstance(item.value.value, int):
                            next_enum_value = item.value.value + 1
                        else:
                            next_enum_value += 1
                        enum_members.append((member_name, next_enum_value - 1, item.lineno))
                elif isinstance(item, ast.AnnAssign):
                    if isinstance(item.target, ast.Name):
                        member_name: str = item.target.id
                        if item.value:
                            if isinstance(item.value, ast.Constant) and isinstance(item.value.value, int):
                                next_enum_value = item.value.value + 1
                            else:
                                next_enum_value += 1
                        else:
                            next_enum_value += 1
                        enum_members.append((member_name, next_enum_value - 1, item.lineno))

        # REnum 变体提取（嵌套 ClassDef + Assign 显式值）
        # 与 HandlesClassDef._EmitREnumLlvm 的 tag 分配逻辑对齐：
        #   - 先收集 Assign 显式值
        #   - ClassDef 变体从 max(explicit)+1 开始顺序编号
        #   - Assign 无显式值时用独立计数器 variant_index
        if type_kind == 'renum':
            explicit_values: dict[str, int] = {}
            for item in node.body:
                if isinstance(item, ast.Assign):
                    for target in item.targets:
                        if isinstance(target, ast.Name):
                            if item.value and isinstance(item.value, ast.Constant) and isinstance(item.value.value, int):
                                explicit_values[target.id] = item.value.value
            next_class_tag: int = max(explicit_values.values()) + 1 if explicit_values else 0
            variant_index: int = 0
            for item in node.body:
                if isinstance(item, ast.ClassDef):
                    variant_name_r: str = item.name
                    tag: int = next_class_tag
                    next_class_tag += 1
                    enum_members.append((variant_name_r, tag, item.lineno))
                elif isinstance(item, ast.Assign):
                    for target in item.targets:
                        if isinstance(target, ast.Name):
                            var_name: str = target.id
                            if item.value and isinstance(item.value, ast.Constant) and isinstance(item.value.value, int):
                                value: int = item.value.value
                            else:
                                value = variant_index
                                variant_index += 1
                            enum_members.append((var_name, value, item.lineno))

        return ClassSymbolData(
            name=node.name,
            type_kind=type_kind,
            lineno=node.lineno,
            members=members_info,
            is_cpython_object=is_cpython_object,
            is_packed=is_packed,
            enum_members=enum_members,
        )

    # ------------------------------------------------------------------
    #  AnnAssign 提取（typedef / define）
    # ------------------------------------------------------------------

    def _extract_ann_assign(self, node: ast.AnnAssign, result: ModuleSymbols) -> None:
        """从 AnnAssign 节点提取 typedef 或 define 数据"""
        var_name: str = node.target.id
        has_cdefine: bool = AnnotationContainsTType(node.annotation, 'CDefine') if node.annotation else False
        has_postdef: bool = AnnotationContainsTType(node.annotation, 'Postdefinition') if node.annotation else False
        has_ctypedef: bool = AnnotationContainsTType(node.annotation, 'CTypedef') if node.annotation else False

        if has_cdefine:
            result.defines.append(DefineSymbolData(
                name=var_name,
                lineno=node.lineno,
                value_node=node.value,
            ))
        else:
            result.typedefs.append(TypedefSymbolData(
                name=var_name,
                lineno=node.lineno,
                annotation=node.annotation,
                value=node.value,
                has_cdefine=has_cdefine,
                has_postdef=has_postdef,
                has_ctypedef=has_ctypedef,
            ))

    # ------------------------------------------------------------------
    #  函数提取
    # ------------------------------------------------------------------

    def _extract_function(self, node: ast.FunctionDef) -> FuncSymbolData:
        """从 FunctionDef 节点提取函数符号数据"""
        params: list[tuple[str, ast.AST | None]] = [(arg.arg, arg.annotation) for arg in node.args.args]
        is_variadic: bool = node.args.vararg is not None
        is_inline: bool = CheckAnnotationHasCInline(node.returns)

        return FuncSymbolData(
            name=node.name,
            lineno=node.lineno,
            returns_node=node.returns,
            params=params,
            is_variadic=is_variadic,
            is_inline=is_inline,
        )
