#!/usr/bin/env python3

from __future__ import annotations

import argparse
import re
from dataclasses import dataclass, field
from pathlib import Path

import mcschematic


# ============================================================
# CONFIGURATION
# ============================================================

MAX_LINES = 64

# 1 bloc VIDE entre chaque bit
BIT_GAP = 1

# 3 blocs VIDES entre deux groupes de bits (champs)
FIELD_GAP = 3

# 4 blocs VIDES entre deux lignes
ROW_GAP = 4


# ============================================================
# BLOCS MINECRAFT
# ============================================================

LAMP_OFF = "minecraft:redstone_lamp[lit=false]"
LAMP_ON = "minecraft:redstone_lamp[lit=true]"

# Le levier est SUR la face avant de la lampe.
#
# La lampe est à Z = 0
# Le levier est à Z = 1
#
LEVER_TEMPLATE = (
    "minecraft:lever"
    "[face=wall,facing=north,powered={powered}]"
)


# ============================================================
# OPCODES
# ============================================================

OPCODES = {
    "NOP":  "0000",
    "LOAD": "0001",
    "STORE": "0010",
    "MOV":  "0011",
    "ADD":  "0100",
    "0101": "0101",
    "AND":  "0110",
    "OR":   "0111",
    "XOR":  "1000",
    "NOT":  "1001",
    "SHL":  "1010",
    "SHR":  "1011",
    "CMP":  "1100",
    "JMP":  "1101",
    "JZ":   "1110",
    "HALT": "1111",
    "HLT":  "1111",
}

REGISTER_REGEX = re.compile(r"^[Rr]([0-7])$")


# ============================================================
# STRUCTURES
# ============================================================

@dataclass
class Instruction:
    mnemonic: str
    arguments: list[str]
    source_line: int

    # Liste des champs binaires de l'instruction, dans l'ordre
    # logique (avant rotation 180°), ex :
    #   [opcode, arg1, arg2, arg3]
    # ou, pour STORE :
    #   [opcode, colonne, ligne(one-hot), registre]
    #
    # Les champs peuvent avoir des largeurs différentes selon
    # l'instruction.
    fields: list[str] = field(default_factory=list)

    @property
    def binary(self) -> str:
        return "".join(self.fields)


class AssemblerError(Exception):
    pass


# ============================================================
# COMMENTAIRES
# ============================================================

def remove_comment(line: str) -> str:
    """
    Supprime les commentaires commençant par :
        ;
        #

    Exemple :

        LOAD R2 1 ; test

    devient :

        LOAD R2 1
    """

    for marker in (";", "#"):
        position = line.find(marker)

        if position != -1:
            line = line[:position]

    return line.strip()


# ============================================================
# PARSER
# ============================================================

def parse_assembly(path: Path) -> list[Instruction]:

    if not path.exists():
        raise AssemblerError(
            f"Fichier introuvable : {path}"
        )

    if not path.is_file():
        raise AssemblerError(
            f"Ce chemin n'est pas un fichier : {path}"
        )

    try:
        content = path.read_text(
            encoding="utf-8"
        )
    except UnicodeDecodeError as exc:
        raise AssemblerError(
            "Le fichier assembleur doit être "
            "encodé en UTF-8."
        ) from exc

    instructions: list[Instruction] = []

    for line_number, raw_line in enumerate(
        content.splitlines(),
        start=1,
    ):

        line = remove_comment(raw_line)

        if not line:
            continue

        parts = line.split()

        mnemonic = parts[0].upper()
        arguments = parts[1:]

        if mnemonic not in OPCODES:
            raise AssemblerError(
                f"Ligne {line_number}: "
                f"instruction inconnue "
                f"'{parts[0]}'."
            )

        if len(arguments) > 3:
            raise AssemblerError(
                f"Ligne {line_number}: "
                "maximum 3 arguments."
            )

        instructions.append(
            Instruction(
                mnemonic=mnemonic,
                arguments=arguments,
                source_line=line_number,
            )
        )

    if not instructions:
        raise AssemblerError(
            "Le fichier ne contient aucune instruction."
        )

    if len(instructions) > MAX_LINES:
        raise AssemblerError(
            f"Le programme contient "
            f"{len(instructions)} instructions. "
            f"Maximum : {MAX_LINES}."
        )

    # ========================================================
    # HALT OBLIGATOIRE
    # ========================================================

    if instructions[-1].mnemonic not in (
        "HALT",
        "HLT",
    ):
        raise AssemblerError(
            "Le programme doit obligatoirement "
            "se terminer par HALT ou HLT."
        )

    # Aucun HALT avant la dernière instruction
    for instruction in instructions[:-1]:

        if instruction.mnemonic in (
            "HALT",
            "HLT",
        ):
            raise AssemblerError(
                f"Ligne {instruction.source_line}: "
                "HALT doit être la dernière instruction."
            )

    return instructions


# ============================================================
# REGISTRES
# ============================================================

def encode_register(
    value: str,
    line: int,
) -> str:

    match = REGISTER_REGEX.fullmatch(value)

    if not match:
        raise AssemblerError(
            f"Ligne {line}: "
            f"registre invalide '{value}'. "
            "Les registres vont de R0 à R7."
        )

    number = int(match.group(1))

    return f"{number:04b}"


# ============================================================
# ARGUMENT 4 BITS
# ============================================================

def encode_4bit(
    value: str,
    line: int,
) -> str:

    # --------------------------------------------------------
    # Registre
    # --------------------------------------------------------

    if REGISTER_REGEX.fullmatch(value):

        return encode_register(
            value,
            line,
        )

    # --------------------------------------------------------
    # Binaire direct
    # --------------------------------------------------------

    if (
        len(value) == 4
        and all(
            character in "01"
            for character in value
        )
    ):
        return value

    # --------------------------------------------------------
    # Décimal
    # --------------------------------------------------------

    if value.isdigit():

        number = int(value)

        if not 0 <= number <= 15:
            raise AssemblerError(
                f"Ligne {line}: "
                f"'{value}' ne tient pas "
                "sur 4 bits."
            )

        return f"{number:04b}"

    raise AssemblerError(
        f"Ligne {line}: "
        f"argument invalide '{value}'."
    )


# ============================================================
# ARGUMENT 8 BITS
# ============================================================

def encode_8bit(
    value: str,
    line: int,
) -> tuple[str, str]:
    """
    Convertit une valeur en 8 bits.

    Exemple :

        24
        ↓
        00011000
        ↓
        0001 | 1000

    IMPORTANT :
    on ne sépare PAS "2" et "4".
    """

    # Binaire direct
    if (
        len(value) == 8
        and all(
            character in "01"
            for character in value
        )
    ):
        binary = value

    # Décimal
    elif value.isdigit():

        number = int(value)

        if not 0 <= number <= 255:
            raise AssemblerError(
                f"Ligne {line}: "
                f"'{value}' ne tient pas "
                "sur 8 bits."
            )

        binary = f"{number:08b}"

    else:
        raise AssemblerError(
            f"Ligne {line}: "
            f"valeur 8 bits invalide "
            f"'{value}'."
        )

    # Coupe littéralement les 8 bits
    return binary[:4], binary[4:]


# ============================================================
# ASSEMBLAGE D'UNE INSTRUCTION
# ============================================================

def assemble_instruction(
    instruction: Instruction,
) -> list[str]:
    """
    Retourne la liste des champs binaires de l'instruction,
    dans l'ordre logique (avant rotation 180°).

    Toutes les instructions ont 4 champs de 4 bits
    (opcode + 3 arguments), soit 16 bits au total.
    """

    mnemonic = instruction.mnemonic
    args = instruction.arguments
    line = instruction.source_line

    opcode = OPCODES[mnemonic]

    # ========================================================
    # STORE <adresse mémoire / colonne> <registre>
    #
    # L'adresse (8 bits, comme avant) désigne la colonne en
    # mémoire. La ligne allumée à l'écran dépend de la valeur
    # stockée dans le registre : c'est un octet où chaque bit
    # correspond à une ligne, positionnellement (pas un
    # comptage binaire classique) :
    #
    #   registre = 00000001 -> pixel tout en bas allumé
    #   registre = 10000000 -> pixel tout en haut allumé
    #
    # LOAD n'écrit que sur 4 bits (0-15) : pour une valeur plus
    # grande, on la construit avec des ADD (ex : LOAD R7 <a>,
    # LOAD R6 <b>, ADD -> résultat dans R5), puis on l'utilise
    # directement ou on la déplace avec MOV avant le STORE.
    # ========================================================

    if mnemonic == "STORE":

        if len(args) != 2:
            raise AssemblerError(
                f"Ligne {line}: "
                "STORE attend exactement 2 arguments : "
                "STORE <adresse 0-255> <registre>."
            )

        address_high, address_low = encode_8bit(
            args[0],
            line,
        )

        register_bits = encode_register(
            args[1],
            line,
        )

        return [
            opcode,
            address_high,
            address_low,
            register_bits,
        ]

    # ========================================================
    # JMP <ligne>
    # ========================================================

    if mnemonic == "JMP":

        if len(args) != 1:
            raise AssemblerError(
                f"Ligne {line}: "
                "JMP attend exactement "
                "1 argument."
            )

        if not args[0].isdigit():
            raise AssemblerError(
                f"Ligne {line}: "
                "JMP attend un numéro de ligne."
            )

        target = int(args[0])

        if not 0 <= target <= 63:
            raise AssemblerError(
                f"Ligne {line}: "
                f"JMP cible {target}. "
                "Les lignes vont de 0 à 63."
            )

        arg1, arg2 = encode_8bit(
            args[0],
            line,
        )

        return [
            opcode,
            arg1,
            arg2,
            "0000",
        ]

    # ========================================================
    # JZ <ligne> <condition>
    # ========================================================

    if mnemonic == "JZ":

        if len(args) != 2:
            raise AssemblerError(
                f"Ligne {line}: "
                "JZ attend exactement "
                "2 arguments : "
                "JZ <ligne> <condition>."
            )

        if not args[0].isdigit():
            raise AssemblerError(
                f"Ligne {line}: "
                "JZ attend un numéro de ligne."
            )

        target = int(args[0])

        if not 0 <= target <= 63:
            raise AssemblerError(
                f"Ligne {line}: "
                f"JZ cible {target}. "
                "Les lignes vont de 0 à 63."
            )

        # Adresse 8 bits :
        # ARG1 + ARG2
        arg1, arg2 = encode_8bit(
            args[0],
            line,
        )

        # Condition :
        # ARG3
        arg3 = encode_4bit(
            args[1],
            line,
        )

        return [
            opcode,
            arg1,
            arg2,
            arg3,
        ]

    # ========================================================
    # INSTRUCTIONS NORMALES
    # ========================================================

    encoded_args = []

    for argument in args:

        encoded_args.append(
            encode_4bit(
                argument,
                line,
            )
        )

    # Toujours exactement 3 arguments
    while len(encoded_args) < 3:
        encoded_args.append("0000")

    return [
        opcode,
        encoded_args[0],
        encoded_args[1],
        encoded_args[2],
    ]


# ============================================================
# ASSEMBLAGE COMPLET
# ============================================================

def assemble(
    instructions: list[Instruction],
) -> None:

    for instruction in instructions:

        instruction.fields = (
            assemble_instruction(
                instruction
            )
        )


# ============================================================
# AFFICHAGE
# ============================================================

def print_program(
    instructions: list[Instruction],
) -> None:

    print()
    print(
        "LINE   INSTRUCTION       ARGS            FIELDS"
    )
    print("-" * 100)

    for index, instruction in enumerate(
        instructions
    ):

        fields_display = " ".join(
            instruction.fields
        )

        print(
            f"{index:02d}     "
            f"{instruction.mnemonic:<8} "
            f"{' '.join(instruction.arguments):<15} "
            f"{fields_display}"
        )

    print("-" * 100)
    print()


# ============================================================
# GÉNÉRATEUR SCHEMATIC
# ============================================================

class SchematicGenerator:

    def __init__(self):

        self.schematic = (
            mcschematic.MCSchematic()
        )

    # ========================================================
    # AJOUT D'UN BIT
    # ========================================================

    def add_bit(
        self,
        x: int,
        y: int,
        bit: str,
    ) -> None:

        powered = bit == "1"

        # ----------------------------------------------------
        # LAMPE
        # ----------------------------------------------------

        self.schematic.setBlock(
            (x, y, 0),
            LAMP_ON if powered else LAMP_OFF,
        )

        # ----------------------------------------------------
        # LEVIER
        #
        # Il est de l'autre côté de la lampe,
        # à 2 blocs de celle-ci.
        # ----------------------------------------------------

        self.schematic.setBlock(
            (x, y, -1),
            LEVER_TEMPLATE.format(
                powered=(
                    "true"
                    if powered
                    else "false"
                )
            ),
        )

    # ========================================================
    # UNE INSTRUCTION
    # ========================================================

    def add_instruction(
        self,
        index: int,
        fields: list[str],
    ) -> None:
        """
        `fields` est la liste des champs binaires de
        l'instruction dans l'ordre logique, ex :

            [opcode, arg1, arg2, arg3]

        ou, pour STORE :

            [opcode, colonne, ligne(16 bits), registre]

        Les champs peuvent avoir des largeurs différentes :
        cette méthode ne suppose plus que chaque champ fait
        4 bits.
        """

        # ----------------------------------------------------
        # ROTATION 180°
        #
        # Une vraie rotation 180° dans le plan de la ligne :
        # l'ordre des champs est inversé, et chaque champ
        # est lui-même inversé (bit à bit).
        # ----------------------------------------------------

        rotated_fields = [
            field_bits[::-1]
            for field_bits in reversed(fields)
        ]

        # ----------------------------------------------------
        # Position de la ligne
        #
        # Première ligne = Y = 0
        # Deuxième       = Y = -5
        # Troisième      = Y = -10
        #
        # Donc le programme descend.
        # ----------------------------------------------------

        y = -(
            index * (
                1 + ROW_GAP
            )
        )

        # ----------------------------------------------------
        # Construction de la ligne
        #
        # Chaque bit est séparé de 1 bloc :
        #
        #   B . B . B . B
        #
        # Les champs sont séparés par FIELD_GAP blocs :
        #
        #   B . B . B . B ... B . B . B . B
        #
        # ----------------------------------------------------

        x = 0

        for field_index, field_bits in enumerate(
            rotated_fields
        ):

            # Écart entre les champs
            #
            # Après le dernier bit d'un champ, x a déjà avancé
            # de BIT_GAP blocs vides (via le "x += 2" ci-dessous).
            # On ne rajoute donc que la différence pour arriver
            # à un écart total de FIELD_GAP blocs vides.
            if field_index > 0:

                x += FIELD_GAP - BIT_GAP

            for bit in field_bits:

                self.add_bit(
                    x=x,
                    y=y,
                    bit=bit,
                )

                # 1 bloc VIDE entre deux bits
                x += 2

    # ========================================================
    # PROGRAMME COMPLET
    # ========================================================

    def generate(
        self,
        instructions: list[Instruction],
    ) -> None:

        for index, instruction in enumerate(
            instructions
        ):

            self.add_instruction(
                index,
                instruction.fields,
            )


# ============================================================
# SAUVEGARDE
# ============================================================

def save_schematic(
    generator: SchematicGenerator,
    output_path: Path,
) -> None:

    output_path.parent.mkdir(
        parents=True,
        exist_ok=True,
    )

    generator.schematic.save(
        str(output_path.parent),
        output_path.stem,
        mcschematic.Version.JE_1_21_1,
    )


# ============================================================
# MAIN
# ============================================================

def main() -> int:

    parser = argparse.ArgumentParser(
        description=(
            "Custom ASM -> Minecraft .schem"
        )
    )

    parser.add_argument(
        "input",
        type=Path,
        help="Fichier assembleur",
    )

    parser.add_argument(
        "-o",
        "--output",
        type=Path,
        help="Fichier .schem de sortie",
    )

    args = parser.parse_args()

    input_path = args.input

    if args.output:

        output_path = args.output

    else:

        output_path = (
            input_path.with_suffix(
                ".schem"
            )
        )

    try:

        # ====================================================
        # 1. PARSING
        # ====================================================

        print(
            f"[1/4] Parsing : {input_path}"
        )

        instructions = (
            parse_assembly(
                input_path
            )
        )

        print(
            f"      {len(instructions)}/64 instructions"
        )

        # ====================================================
        # 2. ASSEMBLAGE
        # ====================================================

        print(
            "[2/4] Assemblage..."
        )

        assemble(
            instructions
        )

        print_program(
            instructions
        )

        # ====================================================
        # 3. GÉNÉRATION
        # ====================================================

        print(
            "[3/4] Génération du schematic..."
        )

        generator = (
            SchematicGenerator()
        )

        generator.generate(
            instructions
        )

        # ====================================================
        # 4. SAUVEGARDE
        # ====================================================

        print(
            "[4/4] Sauvegarde..."
        )

        save_schematic(
            generator,
            output_path,
        )

        print()
        print(
            "=" * 55
        )
        print(
            "             SCHEMATIC GÉNÉRÉ"
        )
        print(
            "=" * 55
        )
        print(
            f"Fichier : {output_path}"
        )
        print(
            f"Lignes  : {len(instructions)}/64"
        )
        print()
        print(
            "Disposition :"
        )
        print(
            "  • rotation 180°"
        )
        print(
            "  • champs -> ordre inversé, chacun inversé bit à bit"
        )
        print(
            "  • LOAD <registre> <valeur 0-15>, "
            "ex. LOAD R0 5"
        )
        print(
            "  • STORE <adresse/colonne 0-255> <registre>, "
            "ex. STORE 5 R0"
        )
        print(
            "  • dans le registre, chaque bit = une ligne "
            "(00000001 = bas, 10000000 = haut)"
        )
        print(
            "  • valeur > 15 : LOAD R7 <a>, LOAD R6 <b>, "
            "ADD (résultat -> R5), puis utiliser R5 "
            "ou le MOV avant STORE"
        )
        print(
            "  • 1 bloc vide entre chaque bit"
        )
        print(
            "  • 3 blocs vides entre les champs"
        )
        print(
            "  • 4 blocs vides entre les lignes"
        )
        print(
            "  • les lignes descendent"
        )
        print(
            "  • levier sur la face avant"
        )
        print()

        return 0

    except AssemblerError as exc:

        print()
        print(
            f"ERREUR : {exc}"
        )
        print()

        return 1

    except Exception as exc:

        print()
        print(
            f"ERREUR INATTENDUE : "
            f"{type(exc).__name__}: {exc}"
        )
        print()

        return 2


if __name__ == "__main__":
    raise SystemExit(
        main()
    )