import re
from analyzer import StackItem, StackEffect, Instruction, Uop, PseudoInstruction
from dataclasses import dataclass
from cwriter import CWriter
from typing import Iterator

UNUSED = {"unused"}


def maybe_parenthesize(sym: str) -> str:
    """Add parentheses around a string if it contains an operator
       and is not already parenthesized.

    An exception is made for '*' which is common and harmless
    in the context where the symbolic size is used.
    """
    if sym.startswith("(") and sym.endswith(")"):
        return sym
    if re.match(r"^[\s\w*]+$", sym):
        return sym
    else:
        return f"({sym})"


def var_size(var: StackItem) -> str:
    if var.condition:
        # Special case simplifications
        if var.condition == "0":
            return "0"
        elif var.condition == "1":
            return var.get_size()
        elif var.condition == "oparg & 1" and not var.size:
            return f"({var.condition})"
        else:
            return f"(({var.condition}) ? {var.get_size()} : 0)"
    elif var.size:
        return var.size
    else:
        return "1"


@dataclass
class Local:
    item: StackItem
    cached: bool
    in_memory: bool
    defined: bool

    @staticmethod
    def unused(defn: StackItem) -> "Local":
        return Local(defn, False, defn.is_array(), False)

    @staticmethod
    def local(defn: StackItem) -> "Local":
        array = defn.is_array()
        return Local(defn, not array, array, True)

    @staticmethod
    def redefinition(var: StackItem, prev: "Local") -> "Local":
        assert var.is_array() == prev.is_array()
        return Local(var, prev.cached, prev.in_memory, True)

    @property
    def size(self) -> str:
        return self.item.size

    @property
    def name(self) -> str:
        return self.item.name

    @property
    def condition(self) -> str | None:
        return self.item.condition

    def is_array(self) -> bool:
        return self.item.is_array()


@dataclass
class StackOffset:
    "The stack offset of the virtual base of the stack from the physical stack pointer"

    popped: list[str]
    pushed: list[str]

    @staticmethod
    def empty() -> "StackOffset":
        return StackOffset([], [])

    def copy(self) -> "StackOffset":
        return StackOffset(self.popped[:], self.pushed[:])

    def pop(self, item: StackItem) -> None:
        self.popped.append(var_size(item))

    def push(self, item: StackItem) -> None:
        self.pushed.append(var_size(item))

    def __sub__(self, other: "StackOffset") -> "StackOffset":
        return StackOffset(self.popped + other.pushed, self.pushed + other.popped)

    def __neg__(self) -> "StackOffset":
        return StackOffset(self.pushed, self.popped)

    def simplify(self) -> None:
        "Remove matching values from both the popped and pushed list"
        if not self.popped:
            self.pushed.sort()
            return
        if not self.pushed:
            self.popped.sort()
            return
        # Sort the list so the lexically largest element is last.
        popped = sorted(self.popped)
        pushed = sorted(self.pushed)
        self.popped = []
        self.pushed = []
        while popped and pushed:
            pop = popped.pop()
            push = pushed.pop()
            if pop == push:
                pass
            elif pop > push:
                # if pop > push, there can be no element in pushed matching pop.
                self.popped.append(pop)
                pushed.append(push)
            else:
                self.pushed.append(push)
                popped.append(pop)
        self.popped.extend(popped)
        self.pushed.extend(pushed)
        self.pushed.sort()
        self.popped.sort()

    def to_c(self) -> str:
        self.simplify()
        int_offset = 0
        symbol_offset = ""
        for item in self.popped:
            try:
                int_offset -= int(item)
            except ValueError:
                symbol_offset += f" - {maybe_parenthesize(item)}"
        for item in self.pushed:
            try:
                int_offset += int(item)
            except ValueError:
                symbol_offset += f" + {maybe_parenthesize(item)}"
        if symbol_offset and not int_offset:
            res = symbol_offset
        else:
            res = f"{int_offset}{symbol_offset}"
        if res.startswith(" + "):
            res = res[3:]
        if res.startswith(" - "):
            res = "-" + res[3:]
        return res

    def clear(self) -> None:
        self.popped = []
        self.pushed = []


class StackError(Exception):
    pass


class Stack:
    def __init__(self) -> None:
        self.top_offset = StackOffset.empty()
        self.base_offset = StackOffset.empty()
        self.variables: list[Local] = []
        self.defined: set[str] = set()

    def pop(self, var: StackItem, extract_bits: bool = False) -> tuple[str, Local]:
        self.top_offset.pop(var)
        indirect = "&" if var.is_array() else ""
        if self.variables:
            popped = self.variables.pop()
            if popped.size != var.size:
                raise StackError(
                    f"Size mismatch when popping '{popped.name}' from stack to assign to {var.name}. "
                    f"Expected {var.size} got {popped.size}"
                )
            if var.name in UNUSED:
                if popped.name not in UNUSED and popped.name in self.defined:
                    raise StackError(
                        f"Value is declared unused, but is already cached by prior operation"
                    )
                return "", popped
            if not var.used:
                return "", popped
            self.defined.add(var.name)
            if popped.defined:
                if popped.name == var.name:
                    return "", popped
                else:
                    defn = f"{var.name} = {popped.name};\n"
            else:
                if var.is_array():
                    defn = f"{var.name} = &stack_pointer[{self.top_offset.to_c()}];\n"
                else:
                    defn = f"{var.name} = stack_pointer[{self.top_offset.to_c()}];\n"
            return defn, Local.redefinition(var, popped)

        self.base_offset.pop(var)
        if var.name in UNUSED or not var.used:
            return "", Local.unused(var)
        self.defined.add(var.name)
        cast = f"({var.type})" if (not indirect and var.type) else ""
        bits = ".bits" if cast and not extract_bits else ""
        assign = f"{var.name} = {cast}{indirect}stack_pointer[{self.base_offset.to_c()}]{bits};"
        if var.condition:
            if var.condition == "1":
                assign = f"{assign}\n"
            elif var.condition == "0":
                return "", Local.unused(var)
            else:
                assign = f"if ({var.condition}) {{ {assign} }}\n"
        else:
            assign = f"{assign}\n"
        in_memory = var.is_array() or var.peek
        return assign, Local(var, not var.is_array(), in_memory, True)

    def push(self, var: Local) -> None:
        self.variables.append(var)
        self.top_offset.push(var.item)
        if var.item.used:
            self.defined.add(var.name)
            var.defined = True

    def define_output_arrays(self, outputs: list[StackItem]) -> str:
        res = []
        top_offset = self.top_offset.copy()
        for var in outputs:
            if var.is_array() and var.used and not var.peek:
                c_offset = top_offset.to_c()
                top_offset.push(var)
                res.append(f"{var.name} = &stack_pointer[{c_offset}];\n")
            else:
                top_offset.push(var)
        return "\n".join(res)

    @staticmethod
    def _do_emit(
        out: CWriter,
        var: StackItem,
        base_offset: StackOffset,
        cast_type: str = "uintptr_t",
        extract_bits: bool = False,
    ) -> None:
        cast = f"({cast_type})" if var.type else ""
        bits = ".bits" if cast and not extract_bits else ""
        if var.condition == "0":
            return
        if var.condition and var.condition != "1":
            out.emit(f"if ({var.condition}) ")
        out.emit(f"stack_pointer[{base_offset.to_c()}]{bits} = {cast}{var.name};\n")

    @staticmethod
    def _do_flush(
        out: CWriter,
        variables: list[Local],
        base_offset: StackOffset,
        top_offset: StackOffset,
        cast_type: str = "uintptr_t",
        extract_bits: bool = False,
    ) -> None:
        out.start_line()
        for var in variables:
            if (
                var.cached
                and not var.in_memory
                and not var.item.peek
                and not var.name in UNUSED
            ):
                Stack._do_emit(out, var.item, base_offset, cast_type, extract_bits)
            base_offset.push(var.item)
        if base_offset.to_c() != top_offset.to_c():
            print("base", base_offset, "top", top_offset)
            assert False
        number = base_offset.to_c()
        if number != "0":
            out.emit(f"stack_pointer += {number};\n")
            out.emit("assert(WITHIN_STACK_BOUNDS());\n")
        out.start_line()

    def flush_locally(
        self, out: CWriter, cast_type: str = "uintptr_t", extract_bits: bool = False
    ) -> None:
        self._do_flush(
            out,
            self.variables[:],
            self.base_offset.copy(),
            self.top_offset.copy(),
            cast_type,
            extract_bits,
        )

    def flush(
        self, out: CWriter, cast_type: str = "uintptr_t", extract_bits: bool = False
    ) -> None:
        self._do_flush(
            out,
            self.variables,
            self.base_offset,
            self.top_offset,
            cast_type,
            extract_bits,
        )
        self.variables = []
        self.base_offset.clear()
        self.top_offset.clear()

    def flush_single_var(
        self,
        out: CWriter,
        var_name: str,
        outputs: list[StackItem],
        cast_type: str = "uintptr_t",
        extract_bits: bool = False,
    ) -> None:
        assert any(var.name == var_name for var in outputs)
        base_offset = self.base_offset.copy()
        top_offset = self.top_offset.copy()
        for var in self.variables:
            base_offset.push(var.item)
        for output in outputs:
            if any(output == v.item for v in self.variables):
                # The variable is already on the stack, such as a peeked value
                # in the tier1 generator
                continue
            if output.name == var_name:
                Stack._do_emit(out, output, base_offset, cast_type, extract_bits)
            base_offset.push(output)
            top_offset.push(output)
        if base_offset.to_c() != top_offset.to_c():
            print("base", base_offset, "top", top_offset)
            assert False

    def peek_offset(self) -> str:
        return self.top_offset.to_c()

    def as_comment(self) -> str:
        return f"/* Variables: {[v.name for v in self.variables]}. Base offset: {self.base_offset.to_c()}. Top offset: {self.top_offset.to_c()} */"


def get_stack_effect(inst: Instruction | PseudoInstruction) -> Stack:
    stack = Stack()

    def stacks(inst: Instruction | PseudoInstruction) -> Iterator[StackEffect]:
        if isinstance(inst, Instruction):
            for uop in inst.parts:
                if isinstance(uop, Uop):
                    yield uop.stack
        else:
            assert isinstance(inst, PseudoInstruction)
            yield inst.stack

    for s in stacks(inst):
        locals: dict[str, Local] = {}
        for var in reversed(s.inputs):
            _, local = stack.pop(var)
            if var.name != "unused":
                locals[local.name] = local
        for var in s.outputs:
            if var.name in locals:
                local = locals[var.name]
            else:
                local = Local.unused(var)
            stack.push(local)
    return stack