#!/usr/bin/env python3

import logging
from random import choices
from string import ascii_lowercase, digits
from typing import List, Optional, Sequence, Union

from ichor import isa


log = logging.getLogger(__name__)


def gensym(prefix = None) -> isa.Label:
    frag = "".join(choices(ascii_lowercase + digits, k=8))
    return isa.Label(f"{prefix or 'gensym'}_{frag}")


class FuncBuilder(object):
    def __init__(self) -> None:
        self._opcodes = []
        self._slots = {}
        self._stack = 0
        self._straightline = True

    def _write(self, op):
        self._opcodes.append(op)

        # Bookkeeping for statically analyzing the stack
        if not self._straightline:
            return

        match op:
            case isa.Label(_) | isa.GOTO(_) | isa.ATEST(_):
                # Labels and control ops abort
                self._straightline = False

            case isa.CLOSUREC(n) | isa.CLOSUREF(n):
                self._stack -= n

            case isa.ARMREF():
                self._stack -= 1

            case isa.DROP(n):
                self._stack -= n

            case isa.DUP(n):
                self._stack += n

            case isa.IDENTIFIERC(_) | isa.SLOT(_) | isa.CALLF(_) | isa.CALLC(_) | isa.ARM(_):
                self._stack -= getattr(op, "nargs", 0)
                self._stack += 1

            case isa.ROT(_) | isa.FUNREF(_) | isa.TYPEREF() | isa.ALOAD():
                pass

            case _:
                log.debug("No stack manipulation known for %s, assuming no change", op)

    def write(self, op: Union[isa.Opcode, isa.Label, Sequence[isa.Opcode]]):

        def flatten(o):
            for e in o:
                if isinstance(e, (isa.Opcode, isa.Label)):
                    yield e
                else:
                    yield from flatten(e)

        if isinstance(op, (isa.Opcode, isa.Label)):
            self._write(op)

        else:
            for e in flatten(op):
                self.write(e)

    def make_label(self, prefix: Optional[str] = ""):
        return gensym(prefix)

    def mark_argument(self, prefix: Optional[str] = ""):
        if not self._opcodes:
            s = gensym(prefix)
            self._stack += 1
            self._slots[s] = self._stack - 1
            return s
        else:
            raise Exception("Cannot mark arguments after writing opcodes")

    def mark_label(self, prefix: Optional[str] = ""):
        l = self.make_label(prefix=prefix)
        self.write(l)
        return l

    def mark_slot(self, target: Optional[int] = None, prefix: Optional[str] = ""):
        if target is not None:
            s = gensym(prefix)
            self._slots[s] = target
            return s

        elif self._straightline:
            s = gensym(prefix)
            assert self._stack > 0
            self._slots[s] = self._stack - 1
            return s

        else:
            raise Exception("Can't automatically generate slots after control flow (label/goto), provide `target=`.")


    def build(self) -> List[isa.Opcode]:
        """Assemble the written body into fully resolved opcodes."""

        # The trivial two-pass assembler. First pass removes labels from the
        # opcode stream, marking where they occurred.

        labels = {}
        unassembled = []
        for op in self._opcodes:
            match op:
                case isa.Label(_) as l:
                    assert l not in labels  # isa.Label marks must be unique.
                    labels[l] = len(unassembled)
                case o:
                    unassembled.append(o)

        # Second pass rewrites instructions (which can reference forwards OR
        # backwards labels) with real targets instead of labels.
        assembled = []
        for op in unassembled:
            match op:
                case isa.GOTO(isa.Label(_) as l):
                    assembled.append(isa.GOTO(labels[l]))

                case isa.ATEST(isa.Label(_) as l):
                    assembled.append(isa.ATEST(labels[l]))

                case isa.SLOT(isa.Label(_) as l):
                    assembled.append(isa.SLOT(self._slots[l]))

                case o:
                    assembled.append(o)

        return assembled


def assemble(builder_cls=FuncBuilder, /, **opcodes: List[isa.Opcode]) -> List[isa.Opcode]:
    builder = builder_cls()
    for o in opcodes:
        builder.write(o)
    return builder.build()
