"""Insert spills for values that are live across yields."""

from __future__ import annotations

from mypyc.analysis.dataflow import AnalysisResult, analyze_live_regs, get_cfg
from mypyc.common import TEMP_ATTR_NAME
from mypyc.ir.class_ir import ClassIR
from mypyc.ir.func_ir import FuncIR
from mypyc.ir.ops import (
    BasicBlock,
    Branch,
    DecRef,
    GetAttr,
    IncRef,
    LoadErrorValue,
    Register,
    SetAttr,
    Value,
)


def insert_spills(ir: FuncIR, env: ClassIR) -> None:
    cfg = get_cfg(ir.blocks, use_yields=True)
    live = analyze_live_regs(ir.blocks, cfg)
    entry_live = live.before[ir.blocks[0], 0]

    entry_live = {op for op in entry_live if not (isinstance(op, Register) and op.is_arg)}
    # TODO: Actually for now, no Registers at all -- we keep the manual spills
    entry_live = {op for op in entry_live if not isinstance(op, Register)}

    ir.blocks = spill_regs(ir.blocks, env, entry_live, live, ir.arg_regs[0])


def spill_regs(
    blocks: list[BasicBlock],
    env: ClassIR,
    to_spill: set[Value],
    live: AnalysisResult[Value],
    self_reg: Register,
) -> list[BasicBlock]:
    env_reg: Value
    for op in blocks[0].ops:
        if isinstance(op, GetAttr) and op.attr == "__mypyc_env__":
            env_reg = op
            break
    else:
        # Environment has been merged into generator object
        env_reg = self_reg

    spill_locs = {}
    for i, val in enumerate(to_spill):
        name = f"{TEMP_ATTR_NAME}2_{i}"
        env.attributes[name] = val.type
        if val.type.error_overlap:
            # We can safely treat as always initialized, since the type has no pointers.
            # This way we also don't need to manage the defined attribute bitfield.
            env._always_initialized_attrs.add(name)
        spill_locs[val] = name

    for block in blocks:
        ops = block.ops
        block.ops = []

        for i, op in enumerate(ops):
            to_decref = []

            if isinstance(op, IncRef) and op.src in spill_locs:
                raise AssertionError("not sure what to do with an incref of a spill...")
            if isinstance(op, DecRef) and op.src in spill_locs:
                # When we decref a spilled value, we turn that into
                # NULLing out the attribute, but only if the spilled
                # value is not live *when we include yields in the
                # CFG*. (The original decrefs are computed without that.)
                #
                # We also skip a decref is the env register is not
                # live. That should only happen when an exception is
                # being raised, so everything should be handled there.
                if op.src not in live.after[block, i] and env_reg in live.after[block, i]:
                    # Skip the DecRef but null out the spilled location
                    null = LoadErrorValue(op.src.type)
                    block.ops.extend([null, SetAttr(env_reg, spill_locs[op.src], null, op.line)])
                continue

            if (
                any(src in spill_locs for src in op.sources())
                # N.B: IS_ERROR should be before a spill happens
                # XXX: but could we have a regular branch?
                and not (isinstance(op, Branch) and op.op == Branch.IS_ERROR)
            ):
                new_sources: list[Value] = []
                stolen = op.stolen()
                for src in op.sources():
                    if src in spill_locs:
                        read = GetAttr(env_reg, spill_locs[src], op.line)
                        block.ops.append(read)
                        new_sources.append(read)
                        if src.type.is_refcounted and src not in stolen:
                            to_decref.append(read)
                    else:
                        new_sources.append(src)

                op.set_sources(new_sources)

            block.ops.append(op)

            for dec in to_decref:
                block.ops.append(DecRef(dec))

            if op in spill_locs:
                # XXX: could we set uninit?
                block.ops.append(SetAttr(env_reg, spill_locs[op], op, op.line))

    return blocks