from __future__ import annotations
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
    from lib.core.translator import Translator
    from lib.core.LlvmCodeGenerator import LlvmCodeGenerator
from lib.core.Handles.HandlesBase import BaseHandle
import ast
import re
import llvmlite.ir as ir


class ExprFormatHandle(BaseHandle):
    def _HandleJoinedStrLlvm(self, Node: ast.JoinedStr) -> ir.Value:
        Gen: LlvmCodeGenerator = self.Trans.LlvmGen
        if not Gen or not Gen.builder:
            return ir.Constant(ir.IntType(8).as_pointer(), None)
        format_str: str = ""
        cast_args: list[Any] = []
        for part in Node.values:
            if isinstance(part, ast.Constant) and isinstance(part.value, str):
                format_str += part.value.replace('\\', '\\\\').replace('%', '%%').replace('\n', '\\n').replace('\t', '\\t').replace('\0', '\\0')
            elif isinstance(part, ast.FormattedValue):
                fmt_class: str | None = self.Trans.ExprHandler._get_var_class(part.value, Gen)
                if fmt_class and Gen._has_function(f'{fmt_class}.__str__'):
                    obj_val: Any = self.HandleExprLlvm(part.value)
                    if obj_val:
                        if isinstance(obj_val.type, ir.PointerType) and isinstance(obj_val.type.pointee, ir.IntType) and obj_val.type.pointee.width == 8:
                            if fmt_class in Gen.structs:
                                obj_val = Gen.builder.bitcast(obj_val, ir.PointerType(Gen.structs[fmt_class]), name=f"cast_{fmt_class}")
                        val: Any = Gen.builder.call(Gen._get_function(f'{fmt_class}.__str__'), [obj_val], name=f"call_{fmt_class}.__str__")
                        Gen._RegisterTempPtr(val)
                        format_str += "%s"
                        cast_args.append(val)
                    else:
                        format_str += "%s"
                        cast_args.append(ir.Constant(ir.PointerType(ir.IntType(8)), None))
                    continue
                val = self.HandleExprLlvm(part.value)
                if not val:
                    format_str += "%d"
                    cast_args.append(ir.Constant(ir.IntType(32), 0))
                    continue
                # 检查 format_spec（Python f-string 格式说明符）
                if part.format_spec is not None:
                    c_fmt: str | None = self._parse_format_spec(part.format_spec, val)
                    if c_fmt is not None:
                        cast_val: Any = self._cast_value_for_format(val, c_fmt, Gen, part.value)
                        format_str += c_fmt
                        cast_args.append(cast_val)
                        continue
                if isinstance(val.type, ir.PointerType):
                    pointee: ir.Type = val.type.pointee
                    if isinstance(pointee, ir.IntType) and pointee.width == 8:
                        format_str += "%s"
                        cast_args.append(val)
                    elif isinstance(pointee, (ir.LiteralStructType, ir.IdentifiedStructType)):
                        int_val: Any = Gen.builder.ptrtoint(val, ir.IntType(64), name="fstr_ptr2int")
                        trunc_val: Any = Gen.builder.trunc(int_val, ir.IntType(32), name="fstr_trunc")
                        format_str += "%d"
                        cast_args.append(trunc_val)
                    else:
                        Loaded: Any = Gen._load(val, name="fstr_deref")
                        if isinstance(Loaded.type, ir.IntType):
                            if Loaded.type.width == 8:
                                format_str += "%c"
                                ext: Any = Gen.builder.zext(Loaded, ir.IntType(32), name="fstr_char2int")
                                cast_args.append(ext)
                            else:
                                format_str += "%d"
                                cast_args.append(Loaded)
                        elif isinstance(Loaded.type, (ir.FloatType, ir.DoubleType)):
                            format_str += "%f"
                            cast_args.append(Loaded)
                        elif isinstance(Loaded.type, ir.PointerType):
                            format_str += "%s"
                            cast_args.append(Loaded)
                        else:
                            format_str += "%d"
                            cast_args.append(Loaded)
                elif isinstance(val.type, ir.IntType):
                    if val.type.width == 8:
                        format_str += "%c"
                        ext = Gen.builder.zext(val, ir.IntType(32), name="fstr_char2int")
                        cast_args.append(ext)
                    elif val.type.width == 1:
                        zext: Any = Gen.builder.zext(val, ir.IntType(32), name="fstr_bool2int")
                        format_str += "%d"
                        cast_args.append(zext)
                    else:
                        is_u: bool = Gen._check_node_unsigned(part.value)
                        format_str += "%u" if is_u else "%d"
                        cast_args.append(val)
                elif isinstance(val.type, (ir.FloatType, ir.DoubleType)):
                    format_str += "%f"
                    cast_args.append(val)
                else:
                    format_str += "%d"
                    cast_args.append(val)
        if not format_str:
            return Gen.emit_constant("", 'string')
        format_ptr: Any = Gen._create_string_global(format_str)
        # 统一使用 snprintf 生成字符串，返回 i8*
        # __str__ 方法内使用堆分配（返回值需持久化），其他场景使用栈分配（自动释放）
        in_str_method: bool = bool(getattr(Gen, 'func', None) and getattr(Gen.func, 'name', '').endswith('.__str__'))
        snprintf_type: ir.FunctionType = ir.FunctionType(ir.IntType(32), [
            ir.PointerType(ir.IntType(8)),
            ir.IntType(64),
            ir.PointerType(ir.IntType(8))
        ], var_arg=True)
        snprintf_func: Any = Gen.get_or_declare_c_func('snprintf', snprintf_type)
        null_ptr: ir.Constant = ir.Constant(ir.PointerType(ir.IntType(8)), None)
        # 第一遍：获取所需长度（NULL 缓冲，size=0）
        len_val: ir.Value = Gen.builder.call(snprintf_func, [null_ptr, ir.Constant(ir.IntType(64), 0), format_ptr] + cast_args, name="fstr_len")
        # 加 1 用于 null 终止符
        size_val: ir.Value = Gen.builder.add(len_val, ir.Constant(ir.IntType(32), 1), name="fstr_size")
        size_ext: ir.Value = Gen.builder.zext(size_val, ir.IntType(64), name="fstr_size_ext")
        if in_str_method:
            # 堆分配（__str__ 返回值需跨函数生命周期）
            malloc_type: ir.FunctionType = ir.FunctionType(ir.PointerType(ir.IntType(8)), [ir.IntType(64)])
            malloc_func: Any = Gen._get_or_declare_func('malloc', malloc_type)
            buf_ptr: ir.Value = Gen.builder.call(malloc_func, [size_ext], name="fstr_heap_buf")
            Gen.builder.call(snprintf_func, [buf_ptr, size_ext, format_ptr] + cast_args, name="fstr_snprintf")
            Gen._RegisterTempPtr(buf_ptr)
            return buf_ptr
        # 栈分配（局部使用，自动释放）
        buf_ptr: ir.Value = Gen._alloca(ir.IntType(8), name="fstr_stack_buf", size=size_val)
        Gen.builder.call(snprintf_func, [buf_ptr, size_ext, format_ptr] + cast_args, name="fstr_snprintf")
        return buf_ptr

    def _parse_format_spec(self, format_spec_node: ast.AST, val: ir.Value) -> str | None:
        """将 Python f-string format_spec 转换为 C printf 格式说明符。

        支持: d/x/X/o/c/f/s 及 [[fill]align][sign][#][0][width][.precision][type] 语法。
        返回 None 表示无法解析，调用方应回退到类型推断。
        """
        if isinstance(format_spec_node, ast.JoinedStr):
            parts: list[str] = []
            for v in format_spec_node.values:
                if isinstance(v, ast.Constant) and isinstance(v.value, str):
                    parts.append(v.value)
                else:
                    return None
            spec_str: str = ''.join(parts)
        elif isinstance(format_spec_node, ast.Constant) and isinstance(format_spec_node.value, str):
            spec_str = format_spec_node.value
        else:
            return None

        if not spec_str:
            return None

        m: Any = re.match(r'^(.?[<>=^])?([+\- ])?(#)?(0)?(\d+)?(,)?(\.\d+)?([bcdeEfFgGnosxX%])?$', spec_str)
        if not m:
            return None

        fill_align, sign, alt, zero, width, comma, precision, fmt_type = m.groups()

        c_flags: str = ''
        if sign == '+':
            c_flags += '+'
        elif sign == ' ':
            c_flags += ' '
        if alt:
            c_flags += '#'

        if fill_align:
            align_char: str = fill_align[-1] if len(fill_align) > 1 else fill_align[0]
            if align_char in ('<', '^'):
                c_flags += '-'
        if zero:
            c_flags += '0'

        c_width: str = width or ''
        c_precision: str = precision or ''

        length_mod: str = ''
        if isinstance(val.type, ir.IntType):
            if val.type.width == 64:
                length_mod = 'll'
            elif val.type.width == 32:
                length_mod = 'l'
            elif val.type.width == 16:
                length_mod = 'h'
            elif val.type.width == 8:
                length_mod = 'hh'

        c_specifier: str = ''
        if fmt_type in ('d', 'n'):
            c_specifier = f'{length_mod}d'
        elif fmt_type == 'x':
            c_specifier = f'{length_mod}x'
        elif fmt_type == 'X':
            c_specifier = f'{length_mod}X'
        elif fmt_type == 'o':
            c_specifier = f'{length_mod}o'
        elif fmt_type == 'b':
            return None
        elif fmt_type == 'c':
            c_specifier = 'c'
        elif fmt_type in ('f', 'F'):
            c_specifier = 'f'
        elif fmt_type in ('e', 'E', 'g', 'G'):
            c_specifier = fmt_type
        elif fmt_type == 's':
            c_specifier = 's'
        elif fmt_type == '%':
            c_specifier = 'f%%'
        else:
            if isinstance(val.type, ir.IntType):
                if val.type.width == 8:
                    c_specifier = 'c'
                elif val.type.width == 1:
                    c_specifier = 'd'
                else:
                    c_specifier = f'{length_mod}d'
            elif isinstance(val.type, (ir.FloatType, ir.DoubleType)):
                c_specifier = 'f'
            elif isinstance(val.type, ir.PointerType):
                c_specifier = 's'
            else:
                c_specifier = 'd'

        return f'%{c_flags}{c_width}{c_precision}{c_specifier}'

    def _cast_value_for_format(self, val: ir.Value, c_fmt: str, Gen: Any, arg_node: ast.AST) -> ir.Value:
        """根据 C printf 格式说明符转换值类型。"""
        fmt_type_char: str = ''
        for ch in reversed(c_fmt):
            if ch.isalpha() and ch not in ('l', 'h'):
                fmt_type_char = ch
                break
            elif ch == '%':
                break

        if fmt_type_char == 's':
            if isinstance(val.type, ir.PointerType):
                pointee: ir.Type = val.type.pointee
                if isinstance(pointee, ir.IntType) and pointee.width == 8:
                    return val
                if isinstance(pointee, ir.ArrayType) and isinstance(pointee.element, ir.IntType) and pointee.element.width == 8:
                    return Gen.builder.gep(val, [ir.Constant(ir.IntType(32), 0), ir.Constant(ir.IntType(32), 0)], name="fstr_arr2ptr")
            return val
        elif fmt_type_char == 'c':
            if isinstance(val.type, ir.IntType) and val.type.width < 32:
                return Gen.builder.zext(val, ir.IntType(32), name="fstr_zext_c")
            return val
        elif fmt_type_char in ('d', 'x', 'X', 'o', 'i'):
            if isinstance(val.type, ir.IntType):
                if val.type.width == 8:
                    is_u: bool = Gen._check_node_unsigned(arg_node)
                    if is_u:
                        return Gen.builder.zext(val, ir.IntType(32), name="fstr_zext_i8")
                    return Gen.builder.sext(val, ir.IntType(32), name="fstr_sext_i8")
                elif val.type.width == 16:
                    is_u = Gen._check_node_unsigned(arg_node)
                    if is_u:
                        return Gen.builder.zext(val, ir.IntType(32), name="fstr_zext_i16")
                    return Gen.builder.sext(val, ir.IntType(32), name="fstr_sext_i16")
                elif val.type.width == 1:
                    return Gen.builder.zext(val, ir.IntType(32), name="fstr_zext_bool")
                return val
            elif isinstance(val.type, (ir.FloatType, ir.DoubleType)):
                return Gen.builder.fptosi(val, ir.IntType(32), name="fstr_f2i")
            return val
        elif fmt_type_char in ('f', 'e', 'E', 'g', 'G'):
            if isinstance(val.type, ir.FloatType):
                return Gen.builder.fpext(val, ir.DoubleType(), name="fstr_fpext")
            elif isinstance(val.type, ir.IntType):
                return Gen.builder.sitofp(val, ir.DoubleType(), name="fstr_i2f")
            return val
        return val