"""
Nexa Language — LLVM Backend (v23)
Compiles Nexa bytecode (.nexac) → LLVM IR text (.ll) → native binary.

Pipeline:
  1. Load .nexac JSON
  2. Walk instruction list and emit LLVM IR
  3. Optionally invoke `llc` to produce object file, then `clang` to link

Usage:
  nexa build file.nexa --target llvm      → file.ll
  nexa build file.nexa --target llvm-bin  → file (native binary via clang)

Note: Requires llvmlite for the IR-builder path.
      Falls back to hand-written IR text if llvmlite is not available.
"""
from __future__ import annotations
import os, json, subprocess, tempfile
from typing import List, Any

# ── IR Text Emitter (no deps) ─────────────────────────────────────────────────
LLVM_PRELUDE = """\
; Nexa v23 — LLVM IR
target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
declare i32 @printf(i8*, ...)
@.fmt_int  = private constant [4 x i8] c"%d\\0A\\00"
@.fmt_str  = private constant [4 x i8] c"%s\\0A\\00"
@.fmt_flt  = private constant [4 x i8] c"%g\\0A\\00"
"""


class LLVMEmitter:
    """Emit LLVM IR from Nexa bytecode instructions."""

    def __init__(self, source: str = "nexa"):
        self.source = source
        self._lines: List[str] = [LLVM_PRELUDE]
        self._reg = 0        # virtual register counter
        self._stack: List[str] = []   # operand stack (register names)
        self._globals: dict = {}
        self._locals: dict = {}
        self._func_depth = 0

    # ── Helpers ───────────────────────────────────────────────────────────────
    def _new_reg(self) -> str:
        self._reg += 1
        return f"%r{self._reg}"

    def _emit(self, line: str):
        indent = "  " if self._func_depth > 0 else ""
        self._lines.append(indent + line)

    def _push(self, reg: str): self._stack.append(reg)
    def _pop(self) -> str:     return self._stack.pop() if self._stack else "%r0"

    # ── Main entry ────────────────────────────────────────────────────────────
    def emit_module(self, instructions: list) -> str:
        """Emit a self-contained LLVM IR module."""
        self._emit("define i32 @main() {")
        self._emit("entry:")
        self._func_depth = 1
        for instr in instructions:
            op   = instr[0] if isinstance(instr, list) else instr.get("op", "NOP")
            arg  = instr[1] if isinstance(instr, list) else instr.get("arg")
            self._dispatch(op, arg)
        self._emit("  ret i32 0")
        self._emit("}")
        return "\n".join(self._lines)

    def _dispatch(self, op: str, arg: Any):
        if op == "LOAD_CONST":
            r = self._new_reg()
            if isinstance(arg, (int, float)):
                ty = "i64" if isinstance(arg, int) else "double"
                val = int(arg) if isinstance(arg, int) else float(arg)
                self._emit(f"{r} = add {ty} 0, {val}")
            elif isinstance(arg, str):
                slen = len(arg) + 1
                gname = f"@.str{self._reg}"
                escaped = arg.replace("\\", "\\\\").replace('"', '\\"')
                self._lines.insert(4, f'{gname} = private constant [{slen} x i8] c"{escaped}\\00"')
                self._emit(f"{r} = getelementptr [{slen} x i8], [{slen} x i8]* {gname}, i64 0, i64 0")
            elif arg is None:
                r = "null"
            elif isinstance(arg, bool):
                self._emit(f"{r} = add i1 0, {1 if arg else 0}")
            else:
                self._emit(f"; unhandled const: {repr(arg)}")
                r = "%r0"
            self._push(r)

        elif op == "STORE_NAME":
            val = self._pop()
            self._locals[arg] = val
            self._emit(f"; store {arg} = {val}")

        elif op == "LOAD_NAME":
            if arg in self._locals:
                self._push(self._locals[arg])
            else:
                r = self._new_reg()
                self._emit(f"{r} = add i64 0, 0  ; undefined: {arg}")
                self._push(r)

        elif op == "BINARY_OP":
            b = self._pop(); a = self._pop()
            r = self._new_reg()
            op_map = {"+": "add", "-": "sub", "*": "mul", "/": "sdiv",
                      "%": "srem", "==": "icmp eq", "!=": "icmp ne",
                      "<": "icmp slt", ">": "icmp sgt",
                      "<=": "icmp sle", ">=": "icmp sge"}
            llvm_op = op_map.get(arg, "add")
            self._emit(f"{r} = {llvm_op} i64 {a}, {b}")
            self._push(r)

        elif op == "PRINT":
            val = self._pop()
            r = self._new_reg()
            self._emit(f"{r} = call i32 (i8*, ...) @printf(i8* getelementptr inbounds ([4 x i8], [4 x i8]* @.fmt_int, i64 0, i64 0), i64 {val})")

        elif op == "RETURN":
            val = self._pop() if self._stack else "0"
            self._emit(f"ret i64 {val}")

        elif op == "HALT": pass
        elif op == "POP":
            if self._stack: self._pop()
        else:
            self._emit(f"; {op} {repr(arg)}")  # emit as comment for unsupported opcodes


# ── CLI Entry ──────────────────────────────────────────────────────────────────
def build_llvm(nexac_path: str, output_path: str = None, link: bool = False) -> str:
    """
    Load a .nexac file and emit LLVM IR.
    Returns the path to the emitted .ll file.
    """
    with open(nexac_path, encoding="utf-8") as f:
        data = json.load(f)

    instrs = data.get("code", [])
    source = data.get("source", os.path.basename(nexac_path))

    emitter = LLVMEmitter(source)
    ir_text = emitter.emit_module(instrs)

    if output_path is None:
        output_path = os.path.splitext(nexac_path)[0] + ".ll"

    with open(output_path, "w", encoding="utf-8") as f:
        f.write(ir_text)

    print(f"  ✅ LLVM IR → {output_path}  ({len(instrs)} instructions)")

    if link:
        _link_native(output_path)

    return output_path


def _link_native(ll_path: str):
    """Invoke clang/llc to produce a native binary from LLVM IR."""
    bin_path = os.path.splitext(ll_path)[0]
    # Try clang first
    for compiler in ["clang", "gcc"]:
        try:
            result = subprocess.run([compiler, ll_path, "-o", bin_path], capture_output=True)
            if result.returncode == 0:
                print(f"  ✅ Linked native binary → {bin_path}")
                return
        except FileNotFoundError:
            continue
    print(f"  ⚠️  Native linker not found. LLVM IR available at: {ll_path}")
    print(f"  💡 Install clang: https://llvm.org/")
