"""Compiler: ``.hdcasm`` source -> hypervector program image.

Two passes, because labels are forward references:

1. **scan** — strip comments, record label -> cell index, collect the alphabets
   the decoder will later be allowed to use (integer literals, string literals,
   tool names, target cells).  These alphabets are *cleanup memory*: they are
   what makes decoding a bounded argmax instead of an open search.
2. **emit** — turn each line into one composed hypervector and, on install,
   write it into HDC storage as a cell named ``code/<n>``.

Nothing about an instruction survives as text in the image.  The symbolic
parts survive only as *names*, from which their vectors are regenerated
deterministically; the composed instruction vectors are what get stored.  That
is why an image is a genuine program in HDC and not a text file with extra
steps.

    .data tbl0 42           ; preload a storage cell with an integer
    .str  hello "hi there"  ; a string literal
    start:                  ; a label is a cell index
        ldi R0, 5
        bind R1, R0, R0
        prints hello
        halt
"""

from __future__ import annotations

import re
from dataclasses import dataclass, field

from .hypervec import DEFAULT_DIM, hvc, int_to_hv
from .isa import (
    OPCODES,
    REG_NAMES,
    DecodeContext,
    DecodedInstr,
    Opcode,
    decode_instr,
    encode_instr,
)


class AsmError(Exception):
    """Assembly failed; message carries the line number."""


_COMMENT = re.compile(r"(;|#|//).*$")
_LABEL = re.compile(r"^([A-Za-z_.][\w.]*)\s*:\s*(.*)$")
_DIRECTIVE = re.compile(r"^\.(\w+)\s*(.*)$", re.S)
_INT = re.compile(r"^-?(0[xX][0-9a-fA-F]+|\d+)$")


@dataclass
class SourceLine:
    number: int
    text: str
    #: Labels that sit on or immediately above this instruction.  A label-only
    #: line carries forward to the next instruction, because that is the cell
    #: the label actually names.
    labels: tuple[str, ...] = ()


@dataclass
class Program:
    """An assembled, hypervector-encoded program."""

    d: int = DEFAULT_DIM
    instrs: list[int] = field(default_factory=list)
    source: list[SourceLine] = field(default_factory=list)
    labels: dict[str, int] = field(default_factory=dict)
    literals: list[int] = field(default_factory=list)
    strings: list[str] = field(default_factory=list)
    tools: list[str] = field(default_factory=list)
    cell_names: list[str] = field(default_factory=list)
    data: dict[str, str] = field(default_factory=dict)
    #: Decoded view, filled by :func:`assembler` for the IDE's listing.
    listing: list[DecodedInstr] = field(default_factory=list)

    # ------------------------------------------------------------- install
    def install(self, storage, prefix: str = "code") -> list[str]:
        """Write the program into HDC storage.  Returns the code cell names."""
        names: list[str] = []
        for i, hv in enumerate(self.instrs):
            name = "%s/%d" % (prefix, i)
            storage.write(name, hv)
            names.append(name)
        storage.write("program/length", int_to_hv(len(self.instrs), self.d))
        for name, spec in self.data.items():
            kind, _, value = spec.partition(":")
            if kind == "str":
                storage.write(name, hvc("str/" + value, self.d), symbol=True)
            else:
                storage.write(name, int_to_hv(int(value), self.d), symbol=True)
        for name in self.cell_names:
            if name not in storage:
                storage.write(name, 0, symbol=True)
        for label, idx in self.labels.items():
            storage.define("label/" + label)
        return names

    def context(self) -> DecodeContext:
        return DecodeContext(
            d=self.d,
            imms=tuple(self.literals),
            tgts=tuple(range(len(self.instrs))),
            names=tuple(self.cell_names),
            strs=tuple(self.strings),
            tools=tuple(self.tools),
        )

    def listing_text(self) -> str:
        lines = []
        for i, line in enumerate(self.source):
            for lbl in line.labels:
                lines.append("%s:" % lbl)
            body = "    " + line.text
            dec = self.listing[i] if i < len(self.listing) else None
            conf = "" if dec is None else "   ; conf %.2f" % dec.confidence
            lines.append("%-44s%s" % (body, conf))
        return "\n".join(lines)


# --------------------------------------------------------------------------
# parsing helpers
# --------------------------------------------------------------------------
def split_operands(text: str) -> list[str]:
    """Split an operand list on commas, respecting double-quoted strings."""
    if not text.strip():
        return []
    parts: list[str] = []
    cur: list[str] = []
    quoted = False
    for ch in text:
        if ch == '"':
            quoted = not quoted
            cur.append(ch)
        elif ch == "," and not quoted:
            parts.append("".join(cur).strip())
            cur = []
        else:
            cur.append(ch)
    tail = "".join(cur).strip()
    if tail:
        parts.append(tail)
    return parts


def _parse_int(token: str) -> int:
    if not _INT.match(token):
        raise ValueError("not an integer literal: %r" % (token,))
    return int(token, 0) if token.lower().startswith(("0x", "-0x")) else int(token)


def _parse_reg(token: str) -> int:
    t = token.strip().upper()
    if t.startswith("R"):
        t = t[1:]
    if not t.isdigit():
        raise ValueError("not a register: %r" % (token,))
    idx = int(t)
    if not 0 <= idx < len(REG_NAMES):
        raise ValueError("register out of range: %r" % (token,))
    return idx


def _strip_string(token: str) -> str:
    t = token.strip()
    if len(t) >= 2 and t[0] == '"' and t[-1] == '"':
        return t[1:-1]
    raise ValueError("expected a quoted string, got %r" % (token,))


def _resolve_text(token: str, data: dict) -> str:
    """A string operand is either quoted inline or the name of a .str cell.

    Deferring this to after pass 1 is what lets ``prints greeting`` refer to a
    ``.str greeting "..."`` declared further down the file.
    """
    t = token.strip()
    if len(t) >= 2 and t[0] == '"' and t[-1] == '"':
        return t[1:-1]
    spec = data.get(t, "")
    if isinstance(spec, str) and spec.startswith("str:"):
        return spec[4:]
    raise ValueError(
        "unquoted string operand %r is not a declared .str cell" % (t,)
    )


def _scan_operand(
    literals: dict,
    strings: dict,
    names: dict,
    tools: dict,
    str_tokens: list,
    kind: str,
    token: str,
) -> None:
    """Record a token into the alphabet pass-2 decoding will need.

    These four sets become the cleanup alphabets for the whole program, so
    they must be complete before the first instruction is composed.  Labels
    need no entry: pass 1 has already seen every definition.
    """
    if kind == "imm":
        literals.setdefault(_parse_int(token), None)
    elif kind == "str":
        # deferred: .str declarations anywhere in the file are visible
        str_tokens.append(token)
    elif kind == "name":
        names.setdefault(token, None)
    elif kind == "tool":
        tools.setdefault(token, None)


# --------------------------------------------------------------------------
# the compiler
# --------------------------------------------------------------------------
def assembler(source: str | list[str], d: int = DEFAULT_DIM, strict: bool = True) -> Program:
    """Compile assembler text into a :class:`Program`."""
    if isinstance(source, str):
        raw_lines = source.splitlines()
    else:
        raw_lines = list(source)

    # ---------------------------------------------------------- pass 1
    prog = Program(d=d)
    labels: dict[str, int] = {}
    literals: dict[int, None] = {}
    strings: dict[str, None] = {}
    names: dict[str, None] = {}
    tools: dict[str, None] = {}
    str_tokens: list[str] = []
    pending_labels: list[str] = []
    #: Instruction lines only, so ``labels[name]`` is an instruction index and
    #: ``body[i]`` is the i-th instruction.  Directives emit no cells.
    body: list[SourceLine] = []

    for n, raw in enumerate(raw_lines, 1):
        text = _COMMENT.sub("", raw).strip()
        if not text:
            continue
        label = None
        m = _LABEL.match(text)
        if m:
            label = m.group(1)
            text = m.group(2).strip()
            if label in labels:
                raise AsmError("line %d: duplicate label %r" % (n, label))
            labels[label] = len(body)
            pending_labels.append(label)
        if not text:
            continue

        dm = _DIRECTIVE.match(text)
        if dm:
            dname, rest = dm.group(1).lower(), dm.group(2).strip()
            if dname == "str":
                # directives are whitespace separated: .str NAME "some text"
                parts = rest.split(None, 1)
                if len(parts) != 2:
                    raise AsmError('line %d: .str NAME "text"' % n)
                txt = _strip_string(parts[1])
                strings.setdefault(txt, None)
                names.setdefault(parts[0], None)
                prog.data[parts[0]] = "str:" + txt
            elif dname == "data":
                parts = rest.split(None, 1)
                if len(parts) != 2:
                    raise AsmError("line %d: .data NAME VALUE" % n)
                names.setdefault(parts[0], None)
                prog.data[parts[0]] = "int:%d" % _parse_int(parts[1])
            elif dname == "name":
                names.setdefault(rest, None)
            elif dname == "tool":
                tools.setdefault(rest, None)
            else:
                raise AsmError("line %d: unknown directive .%s" % (n, dname))
            # A label on a directive names no program cell.
            pending_labels.clear()
            continue

        mn, _, operand_text = text.partition(" ")
        mn = mn.upper()
        if mn not in OPCODES:
            raise AsmError("line %d: unknown opcode %r" % (n, mn))
        ops = split_operands(operand_text)
        spec = OPCODES[mn]
        want = spec.operand_list
        if len(ops) != len(want):
            raise AsmError(
                "line %d: %s takes %d operand(s) (%s), got %d"
                % (n, mn, len(want), spec.operands, len(ops))
            )
        try:
            for kind, tok in zip(want, ops):
                _scan_operand(literals, strings, names, tools, str_tokens, kind, tok)
        except ValueError as exc:
            raise AsmError("line %d: %s" % (n, exc)) from exc
        body.append(SourceLine(n, text, tuple(pending_labels)))
        pending_labels.clear()

    # Deferred resolution: string operands may name a .str cell declared later,
    # and .data values join the literal alphabet so a value preloaded into
    # storage can also be decoded out of an instruction field.
    for tok in str_tokens:
        try:
            strings.setdefault(_resolve_text(tok, prog.data), None)
        except ValueError as exc:
            raise AsmError(str(exc)) from exc
    for spec in prog.data.values():
        if spec.startswith("int:"):
            literals.setdefault(int(spec[4:]), None)
        elif spec.startswith("str:"):
            strings.setdefault(spec[4:], None)

    prog.labels = labels
    prog.literals = sorted(literals)
    prog.strings = sorted(strings)
    prog.tools = sorted(tools)
    prog.cell_names = sorted(names)

    # ---------------------------------------------------------- pass 2
    for line in body:
        text = line.text
        if _DIRECTIVE.match(text):
            continue
        mn, _, operand_text = text.partition(" ")
        mn = mn.upper()
        spec: Opcode = OPCODES[mn]
        ops = split_operands(operand_text)
        regs: list[int] = []
        kwargs: dict = {}
        try:
            for kind, tok in zip(spec.operand_list, ops):
                if kind == "reg":
                    regs.append(_parse_reg(tok))
                elif kind == "imm":
                    kwargs["imm"] = _parse_int(tok)
                elif kind == "label":
                    if tok.startswith("@"):
                        kwargs["tgt"] = int(tok[1:])
                    elif tok in labels:
                        kwargs["tgt"] = labels[tok]
                    else:
                        raise ValueError("undefined label %r" % (tok,))
                elif kind == "name":
                    kwargs["name"] = tok
                elif kind == "str":
                    kwargs["text"] = _resolve_text(tok, prog.data)
                elif kind == "tool":
                    kwargs["tool"] = tok
            hv = encode_instr(mn, regs=regs, d=d, **kwargs)
        except (ValueError, KeyError) as exc:
            raise AsmError("line %d: %s" % (line.number, exc)) from exc
        prog.instrs.append(hv)

    # ------------------------------------------------- decode-back check
    prog.source = body
    ctx = prog.context()
    for i, hv in enumerate(prog.instrs):
        try:
            dec = decode_instr(hv, ctx)
        except Exception as exc:  # pragma: no cover - defensive
            raise AsmError("internal: cell %d failed to decode: %s" % (i, exc)) from exc
        if strict and dec.confidence < 0.15:
            raise AsmError(
                "cell %d decoded weakly (conf %.2f); check operands" % (i, dec.confidence)
            )
        prog.listing.append(dec)
    return prog


def assemble(source: str | list[str], d: int = DEFAULT_DIM, strict: bool = True) -> Program:
    """Alias kept for the public API."""
    return assembler(source, d=d, strict=strict)


def roundtrip_report(prog: Program) -> str:
    """Self-check: does every stored cell decode back to its source line?"""
    ctx = prog.context()
    rows = []
    bad = 0
    for i, hv in enumerate(prog.instrs):
        dec = decode_instr(hv, ctx)
        src = prog.source[i].text.strip()
        rows.append("%3d  %-30s -> %-30s conf=%.2f" % (i, src, dec.op, dec.confidence))
        if dec.op not in src.upper():
            bad += 1
    head = "decode-back check: %d/%d cells round-trip" % (len(prog.instrs) - bad, len(prog.instrs))
    return "\n".join([head] + rows)