from __future__ import annotations

import ast
from typing import TYPE_CHECKING
import llvmlite.ir as ir

from lib.core.ConstEvaluator import ConstEvaluator

if TYPE_CHECKING:
    from lib.core.LlvmCodeGenerator import LlvmCodeGenerator


class ConstEvalMixin:
    """常量表达式求值 Mixin

    提供数组初始化常量构建与常量表达式求值能力。
    """

    @staticmethod
    def _bswap_constant_value(val: int, width: int) -> int:
        """编译期字节序交换"""
        if width == 16:
            return ((val & 0xFF) << 8) | ((val >> 8) & 0xFF)
        elif width == 32:
            return ((val & 0xFF) << 24) | ((val & 0xFF00) << 8) | ((val >> 8) & 0xFF00) | ((val >> 24) & 0xFF)
        elif width == 64:
            return ((val & 0xFF) << 56) | ((val & 0xFF00) << 40) | ((val & 0xFF0000) << 24) | ((val & 0xFF000000) << 8) | ((val >> 8) & 0xFF000000) | ((val >> 24) & 0xFF0000) | ((val >> 40) & 0xFF00) | ((val >> 56) & 0xFF)
        return val

    def _BuildArrayInitConstants(self, value_node: ast.AST, ElemType: ir.Type, elem_type_node: ast.AST, ArrayCount: int, Gen: LlvmCodeGenerator, byte_order: str = '') -> list[ir.Constant] | None:
        _zv: ir.Constant | None = Gen._zero_value_for_type(ElemType)
        _fallback: ir.Constant = _zv if _zv is not None else ir.Constant(ElemType, ir.Undefined)
        if not isinstance(value_node, (ast.List, ast.Set)):
            return None
        if isinstance(ElemType, ir.ArrayType):
            constants: list[ir.Constant] = []
            for elt in value_node.elts:
                if isinstance(elt, (ast.List, ast.Set)):
                    inner_consts: list[ir.Constant] | None = self._BuildMultiDimArrayInitConstants(elt, ElemType, Gen)
                    if inner_consts:
                        try:
                            constants.append(ir.Constant(ElemType, inner_consts))
                        except Exception:  # fallback to _fallback on constant construction failure
                            constants.append(_fallback)
                    else:
                        constants.append(_fallback)
                else:
                    constants.append(_fallback)
            while len(constants) < ArrayCount:
                constants.append(_fallback)
            return constants[:ArrayCount]
        StructName: str | None = None
        if isinstance(elem_type_node, ast.Name) and elem_type_node.id in Gen.structs:
            StructName = elem_type_node.id
        constants = []
        for elt in value_node.elts:
            if StructName and isinstance(elt, ast.Call) and isinstance(elt.func, ast.Name) and elt.func.id == StructName:
                const: ir.Constant | None = self.ClassHandler._BuildStructConstant(elt, StructName, Gen)
                if const:
                    constants.append(const)
                else:
                    constants.append(_fallback)
            elif isinstance(elt, ast.Constant):
                # 字符串常量：根据引号类型决定 i8 或 i8*
                if isinstance(elt.value, str):
                    const: ir.Constant | None = self._BuildStringLiteralConstant(elt, ElemType, Gen)
                    if const:
                        constants.append(const)
                    else:
                        constants.append(_fallback)
                else:
                    const = self._BuildScalarConstantUnified(elt, ElemType, Gen)
                    if const:
                        constants.append(const)
                    else:
                        constants.append(_fallback)
            elif isinstance(elt, ast.UnaryOp):
                const = self._BuildScalarConstantUnified(elt, ElemType, Gen)
                if const:
                    constants.append(const)
                else:
                    constants.append(_fallback)
            elif isinstance(elt, ast.Call) and isinstance(elt.func, ast.Name) and elt.func.id == 'chr':
                if elt.args and isinstance(elt.args[0], ast.Constant) and isinstance(elt.args[0].value, int):
                    try:
                        constants.append(ir.Constant(ElemType, elt.args[0].value & 0xFF))
                    except Exception:  # fallback to _fallback on constant construction failure
                        constants.append(_fallback)
                else:
                    constants.append(_fallback)
            elif isinstance(elt, ast.Call) and isinstance(elt.func, ast.Name) and elt.func.id == 'ord':
                if elt.args and isinstance(elt.args[0], ast.Constant) and isinstance(elt.args[0].value, str) and len(elt.args[0].value) == 1:
                    try:
                        constants.append(ir.Constant(ElemType, ord(elt.args[0].value)))
                    except Exception:  # fallback to _fallback on constant construction failure
                        constants.append(_fallback)
                else:
                    constants.append(_fallback)
            else:
                const_val: int | float | str | None = self._try_eval_const_call(elt, Gen)
                if const_val is not None:
                    try:
                        constants.append(ir.Constant(ElemType, const_val))
                    except Exception:  # fallback to _fallback on constant construction failure
                        constants.append(_fallback)
                else:
                    constants.append(_fallback)
        while len(constants) < ArrayCount:
            constants.append(_fallback)
        if byte_order == 'big' and isinstance(ElemType, ir.IntType) and ElemType.width in (16, 32, 64):
            swapped: list[ir.Constant] = []
            for c in constants:
                if isinstance(c, ir.Constant) and isinstance(c.type, ir.IntType) and isinstance(c.constant, int):
                    swapped_val: int = self._bswap_constant_value(int(c.constant), ElemType.width)
                    swapped.append(ir.Constant(ElemType, swapped_val))
                else:
                    swapped.append(c)
            constants = swapped
        return constants[:ArrayCount]

    def _BuildStringLiteralConstant(self, node: ast.Constant, target_type: ir.Type, Gen: LlvmCodeGenerator) -> ir.Constant | None:
        """根据引号类型决定字符串常量的类型：

        - 单引号单字符 → i8 (char 值)
        - 双引号/单引号多字符/三重引号 → i8* (C 字符串指针)
        - 空字符串 → 0 (None)，传入 ptr 目标自动隐式转换
        """
        if not isinstance(node.value, str):
            return None
        s: str = node.value
        # 空字符串 → 0 (None)
        if len(s) == 0:
            try:
                return ir.Constant(target_type, 0)
            except Exception:
                return None
        # 获取源代码以判断引号类型
        quote_char: str | None = None
        is_triple: bool = False
        col_offset: int = getattr(node, 'col_offset', 0)
        node_lineno: int = getattr(node, 'lineno', 0)
        # AST 的 lineno 可能不准确（常量节点可能指向下一行），
        # 因此在 lineno 和 lineno-1 两行中查找引号
        for try_lineno in (node_lineno, node_lineno - 1):
            source_line: str | None = Gen._get_source_line(try_lineno)
            if source_line and col_offset < len(source_line):
                ch: str = source_line[col_offset]
                if ch in ("'", '"'):
                    quote_char = ch
                    if col_offset + 2 < len(source_line) and source_line[col_offset:col_offset + 3] in ("'''", '"""'):
                        is_triple = True
                    break
        # 判断是否为 char (i8)：单引号、非三重、单字符
        is_char: bool = (not is_triple) and (quote_char == "'") and (len(s) == 1)
        if is_char:
            # 单引号单字符 → i8
            if isinstance(target_type, ir.IntType):
                try:
                    return ir.Constant(target_type, ord(s))
                except Exception:
                    return None
            # 目标类型不是整数 → 类型不匹配
            return None
        else:
            # 双引号/多字符/三重 → i8* (C 字符串)
            if isinstance(target_type, ir.PointerType):
                str_val: str = s + '\x00'
                str_bytes: bytes = str_val.encode('utf-8')
                arr_type: ir.ArrayType = ir.ArrayType(ir.IntType(8), len(str_bytes))
                gv_name: str = f"str_const_{id(node)}"
                gv: ir.GlobalVariable = ir.GlobalVariable(Gen.module, arr_type, name=gv_name)
                gv.initializer = ir.Constant(arr_type, bytearray(str_bytes))
                gv.linkage = 'internal'
                return ir.Constant(target_type, gv.get_reference())
            # 目标类型不是指针 → 类型不匹配
            return None

    def _BuildScalarConstantUnified(self, node: ast.AST, target_type: ir.Type, Gen: LlvmCodeGenerator) -> ir.Constant | None:
        """统一的标量常量构建方法。

        合并原 HandlesClassDef._BuildScalarConstant 与 HandlesImports._BuildScalarConstant 的特性：
        - 支持 bool/int/float/str/ast.UnaryOp(USub)/ast.Name(True/False)
        - 字符串处理委托给 _BuildStringLiteralConstant
        - 显式拒绝 BaseStructType
        - 全部 try/except 保护
        """
        # 显式拒绝结构体类型
        if isinstance(target_type, ir.BaseStructType):
            return None
        # 处理常量节点
        if isinstance(node, ast.Constant):
            # bool 必须在 int 之前判断（bool 是 int 的子类）
            if isinstance(node.value, bool):
                try:
                    return ir.Constant(target_type, 1 if node.value else 0)
                except Exception:
                    return None
            elif isinstance(node.value, int):
                try:
                    return ir.Constant(target_type, node.value)
                except Exception:
                    return None
            elif isinstance(node.value, float):
                try:
                    return ir.Constant(target_type, node.value)
                except Exception:
                    return None
            elif isinstance(node.value, str):
                # 字符串处理委托给 _BuildStringLiteralConstant
                const: ir.Constant | None = self._BuildStringLiteralConstant(node, target_type, Gen)
                if const is not None:
                    return const
                # 回退：直接创建字符串全局变量（i8*）
                try:
                    return Gen._create_string_global(node.value)
                except Exception:
                    return None
        # 处理一元负号
        elif isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub):
            if isinstance(node.operand, ast.Constant) and isinstance(node.operand.value, int):
                try:
                    return ir.Constant(target_type, -node.operand.value)
                except Exception:
                    return None
        # 处理 True/False 名称节点
        elif isinstance(node, ast.Name):
            if node.id == 'True':
                try:
                    return ir.Constant(target_type, 1)
                except Exception:
                    return None
            elif node.id == 'False':
                try:
                    if isinstance(target_type, ir.PointerType):
                        return ir.Constant(target_type, None)
                    return ir.Constant(target_type, 0)
                except Exception:
                    return None
        return None

    def _try_eval_const_call(self, node: ast.AST, Gen: LlvmCodeGenerator) -> int | None:
        if isinstance(node, ast.Call):
            func_name: str | None = None
            if isinstance(node.func, ast.Name):
                func_name = node.func.id
            elif isinstance(node.func, ast.Attribute):
                func_name = node.func.attr
            if func_name:
                arg_vals: list[int | float | str | None] = []
                for arg in node.args:
                    v: int | float | str | None = self._try_eval_const_expr(arg, Gen)
                    if v is None:
                        return None
                    arg_vals.append(v)
                if func_name == 'COLOR_RGB' and len(arg_vals) == 3:
                    r: int | float | str | None
                    g: int | float | str | None
                    b: int | float | str | None
                    r, g, b = arg_vals
                    return (255 << 24) | (int(b) << 16) | (int(g) << 8) | int(r)
                if func_name in Gen._DefineConstants:
                    return Gen._DefineConstants[func_name]
                for key in Gen.functions:
                    if key.endswith(f'__{func_name}'):
                        return None
            return None
        return self._try_eval_const_expr(node, Gen)

    def _try_eval_const_expr(self, node: ast.AST, Gen: LlvmCodeGenerator) -> int | float | str | None:
        return ConstEvaluator.eval_with_symtab(node, self.SymbolTable)

    def _BuildMultiDimArrayInitConstants(self, value_node: ast.AST, ArrayType: ir.ArrayType, Gen: LlvmCodeGenerator) -> list[ir.Constant] | None:
        inner_type: ir.Type = ArrayType.element
        count: int = ArrayType.count
        elts: list[ast.AST] = list(value_node.elts) if hasattr(value_node, 'elts') else []
        _zv: ir.Constant | None = Gen._zero_value_for_type(inner_type)
        _fallback: ir.Constant = _zv if _zv is not None else ir.Constant(inner_type, ir.Undefined)
        constants: list[ir.Constant] = []
        for elt in elts:
            if isinstance(inner_type, ir.ArrayType):
                if isinstance(elt, (ast.List, ast.Set)):
                    inner_consts: list[ir.Constant] | None = self._BuildMultiDimArrayInitConstants(elt, inner_type, Gen)
                    if inner_consts:
                        constants.append(ir.Constant(inner_type, inner_consts))
                    else:
                        constants.append(_fallback)
                else:
                    constants.append(_fallback)
            elif isinstance(elt, ast.Constant):
                const: ir.Constant | None = self._BuildScalarConstantUnified(elt, inner_type, Gen)
                if const:
                    constants.append(const)
                else:
                    constants.append(_fallback)
            elif isinstance(elt, ast.UnaryOp):
                const = self._BuildScalarConstantUnified(elt, inner_type, Gen)
                if const:
                    constants.append(const)
                else:
                    constants.append(_fallback)
            else:
                constants.append(_fallback)
        while len(constants) < count:
            constants.append(_fallback)
        return constants[:count]