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 llvmlite.ir as ir


class WhileHandle(BaseHandle):
    def _HandleWhileLlvm(self, Node: ast.While) -> None:
        Gen: "Translator.LlvmGen" = self.Trans.LlvmGen
        if Gen._direct_values:
            for var_name in list(Gen._direct_values.keys()):
                OldVal: ir.Value = Gen._direct_values.pop(var_name)
                if var_name not in Gen.variables or Gen.variables.get(var_name) is None:
                    saved_block: ir.Block = Gen.builder.block
                    entry_block: ir.Block = Gen.func.entry_basic_block
                    Gen.builder.position_at_start(entry_block)
                    var: ir.AllocaInstr = Gen.builder.alloca(OldVal.type, name=var_name)
                    Gen._store(OldVal, var)
                    Gen.builder.position_at_end(saved_block)
                    Gen.variables[var_name] = var
        CondBB: ir.Block = Gen.func.append_basic_block(name="while.cond")
        BodyBB: ir.Block = Gen.func.append_basic_block(name="while.body")
        ElseBB: ir.Block | None = Gen.func.append_basic_block(name="while.else") if Node.orelse else None
        AfterBB: ir.Block = Gen.func.append_basic_block(name="while.end")
        Gen.loop_break_targets.append(AfterBB)
        Gen.loop_continue_targets.append(CondBB)
        Gen.builder.branch(CondBB)
        Gen.builder.position_at_start(CondBB)
        Cond: ir.Value | None = self.HandleExprLlvm(Node.test)
        if Cond:
            if not isinstance(Cond.type, ir.IntType) or Cond.type.width != 1:
                if isinstance(Cond.type, (ir.FloatType, ir.DoubleType)):
                    Cond = Gen.builder.fcmp_ordered('!=', Cond, ir.Constant(Cond.type, 0.0), name="whilecond")
                else:
                    Cond = Gen.builder.icmp_signed('!=', Cond, Gen._ZeroConst(Cond.type), name="whilecond")
            if not Gen.builder.block.is_terminated:
                Gen.builder.cbranch(Cond, BodyBB, ElseBB if ElseBB else AfterBB)
        else:
            if not Gen.builder.block.is_terminated:
                Gen.builder.branch(ElseBB if ElseBB else AfterBB)
        Gen.builder.position_at_start(BodyBB)
        self.HandleBodyLlvm(Node.body)
        if not Gen.builder.block.is_terminated:
            Gen.builder.branch(CondBB)
        Gen.loop_break_targets.pop()
        Gen.loop_continue_targets.pop()
        if ElseBB:
            Gen.builder.position_at_start(ElseBB)
            self.HandleBodyLlvm(Node.orelse)
            if not Gen.builder.block.is_terminated:
                Gen.builder.branch(AfterBB)
        Gen.builder.position_at_start(AfterBB)