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


class ExprFormatHandle(BaseHandle):
    def _HandleJoinedStrLlvm(self, Node, return_str=False):
        Gen = self.Trans.LlvmGen
        if not Gen or not Gen.builder:
            return ir.Constant(ir.IntType(8).as_pointer(), None)
        format_str = ""
        cast_args = []
        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 = self.Trans.ExprHandler._get_var_class(part.value, Gen)
                if fmt_class and Gen._has_function(f'{fmt_class}.__str__'):
                    obj_val = 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 = 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
                if isinstance(val.type, ir.PointerType):
                    pointee = 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 = Gen.builder.ptrtoint(val, ir.IntType(64), name="fstr_ptr2int")
                        trunc_val = Gen.builder.trunc(int_val, ir.IntType(32), name="fstr_trunc")
                        format_str += "%d"
                        cast_args.append(trunc_val)
                    else:
                        Loaded = Gen._load(val, name="fstr_deref")
                        if isinstance(Loaded.type, ir.IntType):
                            if Loaded.type.width == 8:
                                format_str += "%c"
                                ext = 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 = Gen.builder.zext(val, ir.IntType(32), name="fstr_bool2int")
                        format_str += "%d"
                        cast_args.append(zext)
                    else:
                        is_u = 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 and not cast_args:
            return Gen.emit_constant("", 'string')
        encoded = format_str.encode('utf-8') + b'\x00'
        string_type = ir.ArrayType(ir.IntType(8), len(encoded))
        string_const = ir.Constant(string_type, bytearray(encoded))
        global_var = ir.GlobalVariable(Gen.module, string_type, name=f"str_const_{Gen.string_const_counter}")
        Gen.string_const_counter += 1
        global_var.initializer = string_const
        global_var.linkage = 'internal'
        format_ptr = Gen.builder.bitcast(global_var, ir.PointerType(ir.IntType(8)))
        if return_str:
            buf_size = ir.Constant(ir.IntType(64), 1024)
            calloc_type = ir.FunctionType(ir.PointerType(ir.IntType(8)), [ir.IntType(64), ir.IntType(64)])
            calloc_func = Gen._get_or_declare_func('calloc', calloc_type)
            buf_ptr = Gen.builder.call(calloc_func, [ir.Constant(ir.IntType(64), 1), buf_size], name="call_calloc_buf")
            sprintf_func = Gen.get_or_declare_c_func('sprintf', ir.FunctionType(ir.IntType(32), [ir.PointerType(ir.IntType(8)), ir.PointerType(ir.IntType(8))], var_arg=True))
            sprintf_args = [buf_ptr, format_ptr] + cast_args
            Gen.builder.call(sprintf_func, sprintf_args, name="call_sprintf")
            return buf_ptr
        printf_func = Gen.get_or_declare_c_func('printf', ir.FunctionType(ir.IntType(32), [ir.PointerType(ir.IntType(8))], var_arg=True))
        call_args = [format_ptr] + cast_args
        return Gen.builder.call(printf_func, call_args, name="call_fstring")