import json
import logging
from argparse import ArgumentParser
from pathlib import Path
from zipfile import ZipFile

from util import is_valid_zip, select_zip

# Costs
# fmt: off
CONSTANT_COST = 1
REGISTER_COST = 1
OP_COST = {
    "SUB" :  4,
    "ADD" :  4,
    "INC" :  4,
    "DEC" :  4,
    "MUL" : 10,
    "DIV" : 10,
    "MOD" : 10,
    "OR"  :  8,
    "AND" :  8,
    "XOR" :  8,
    "INV" :  8,
    "SL"  :  5,
    "SR"  :  5,
    "SLU" :  5,
    "SRU" :  5,
    "ROTL":  7,
    "ROTR":  7,
}
ALG_LINE_COST = 0.5
BASE_OPERATIONS = ("A_ADD_B", "B_SUB_A", "TRANS_A", "TRANS_B")
# fmt: on

# ansi escape codes
ul = "\033[4m"  # underline
end = "\033[0m"  # reset
ylw = "\033[33m"  # yellow


def is_empty_row(row: dict, pedantic: bool = False) -> bool:
    # Check if signal table is non-empty
    if "signal" in row and len(row["signal"]) != 0:
        return False

    # Only check if keys are set if pedantic is true
    if not pedantic:
        return True

    # Check if "unconditional-jump", "conditional-jump" or "label" is set
    keys = ("unconditional-jump", "conditional-jump", "label")
    return not any(key in row for key in keys)


def evaluate(file: Path, pedantic: bool = False) -> float:
    with ZipFile(file, "r") as savefile:
        with savefile.open("machine.json", "r") as machinefile, savefile.open("signal.json", "r") as signalfile:
            machine = json.load(machinefile)
            signal = json.load(signalfile)

    # Load lines of code (rows)
    total_rows = signal["signaltable"]["row"]
    # Remove rows without effect (used for formatting for example)
    rows = [row for row in total_rows if not is_empty_row(row, pedantic=pedantic)]

    logging.debug(f"{file.name} :: Total number of rows: {len(total_rows)}")
    logging.debug(f"{file.name} :: Number of rows after excluding empty: {len(rows)}")

    # Check if IR or PC register was used
    pc_used = any((signal["name"] == "PC.W" and signal["value"] == "1" for row in rows for signal in row["signal"]))
    ir_used = any((signal["name"] == "IR.W" and signal["value"] == "1" for row in rows for signal in row["signal"]))

    if pc_used:
        logging.debug(f"{file.name} :: PC Register was used in signal table row.")
    if ir_used:
        logging.debug(f"{file.name} :: IR Register was used in signal table row.")

    # Load used multiplexer constants
    try:
        mux_input_a = next(filter(lambda mux: mux["muxType"] == "A", machine["machine"]["muxInputs"]))["input"]
        mux_consts_a = [int(mux_input["value"]) for mux_input in mux_input_a if mux_input["type"] == "constant"]
    except StopIteration:
        logging.error(f"{file.name} :: Couldn't find input for multiplexer A. Is the file corrupted?")
        exit(1)
    try:
        mux_input_b = next(filter(lambda mux: mux["muxType"] == "B", machine["machine"]["muxInputs"]))["input"]
        mux_consts_b = [int(mux_input["value"]) for mux_input in mux_input_b if mux_input["type"] == "constant"]
    except StopIteration:
        logging.error(f"{file.name} :: Couldn't find input for multiplexer B. Is the file corrupted?")
        exit(1)

    # Base machine has constants 0 and 1 at multiplexer A. All other constants are extensions.
    constants = set(mux_consts_a + mux_consts_b) - set((0, 1))

    logging.debug(f"{file.name} :: Found {len(mux_consts_a)} constants for multiplexer A: {mux_consts_a}")
    logging.debug(f"{file.name} :: Found {len(mux_consts_b)} constants for multiplexer B: {mux_consts_b}")
    logging.debug(
        f"{file.name} :: Found {len(constants)} total unique constants: [{', '.join([str(c) for c in constants])}]"
    )

    # Load used registers
    registers = machine["machine"]["registers"]["register"]
    registers = [register["name"] for register in registers]

    if pc_used:
        registers.append("PC_ALT")
    if ir_used:
        registers.append("IR_ALT")

    logging.debug(f"{file.name} :: Found {len(registers)} additional registers: {registers}")

    # Load used operations
    operations = machine["machine"]["alu"]["operation"]

    for base_op in BASE_OPERATIONS:
        try:
            operations.remove(base_op)
        except ValueError:
            pass

    # Extract operation (remove operands)
    def get_op(op_str):
        return next(filter(lambda substr: substr not in ("A", "B"), op_str.split("_")))

    operations = list(map(get_op, operations))

    logging.debug(f"{file.name} :: Found {len(operations)} additional operations: {operations}")

    # Sum points
    alg_line_costs = len(rows) * ALG_LINE_COST  # every line of code
    constant_costs = len(constants) * CONSTANT_COST  # constants at both multiplexers
    register_costs = len(registers) * REGISTER_COST  # registers
    operation_costs = 0
    for operation in operations:  # operations
        operation_costs += OP_COST[operation]

    total = alg_line_costs + constant_costs + register_costs + operation_costs

    # Summarize
    costs = (alg_line_costs, constant_costs, register_costs, operation_costs, total)
    precision = max([len(str(float(cost)).split(".")[1].lstrip("0")) for cost in costs])
    #           ^ unreadable but works ¯\_(ツ)_/¯

    logging.debug("")

    logging.info(f"{ul}Summary for {file.name}:{end}\n")
    logging.info(f"   {alg_line_costs:5.{min(precision, 2)}f} LINES (excluding empty lines)")
    logging.info(f"+  {constant_costs:5.{min(precision, 2)}f} CONSTANTS")
    logging.info(f"+  {register_costs:5.{min(precision, 2)}f} REGISTERS")
    logging.info(f"+  {operation_costs:5.{min(precision, 2)}f} OPERATIONS")
    logging.info(f"-------------")
    logging.info(f"=  {ylw}{total:5.{min(precision, 2)}f} TOTAL{end}\n\n")

    return total


if __name__ == "__main__":
    parser = ArgumentParser()
    parser.add_argument(
        "submissions",
        type=Path,
        nargs="+",
        help="One or more submission root folder",
    )
    parser.add_argument(
        "-v",
        "--verbose",
        dest="verbose",
        action="store_true",
        help="Prints additional information",
    )
    parser.add_argument(
        "-p",
        "--pedantic",
        dest="pedantic",
        action="store_true",
        help="Extra pedantic (for example when checking for empty lines)",
    )
    parser.add_argument(
        "-r",
        "--rename",
        dest="rename",
        action="store_true",
        help="Rename ZIP files to be the same name as its group directory",
    )
    args = parser.parse_args()

    # Logger setup (DEBUG used for verbose output)
    if args.verbose:
        logging.basicConfig(format="%(message)s", level=logging.DEBUG)
    else:
        logging.basicConfig(format="%(message)s", level=logging.INFO)

    # Output score for each file
    scores = []

    # Gather simulator files
    savefiles = []
    for submissions in args.submissions:
        for group in [f for f in submissions.iterdir() if f.is_dir()]:
            # Find submission ZIP file
            zips = [file for file in group.glob("*.zip") if is_valid_zip(file)]
            if len(zips) == 0:
                logging.error(f"Could not find valid ZIP file for {group.name}")
            savefile = zips[0] if len(zips) == 1 else select_zip(zips)

            # Rename if required
            if args.rename and savefile.stem != group.stem:
                savefile = savefile.rename((savefile.parent / group.stem).with_suffix(".zip"))
            savefiles.append(savefile)

    # Evaluate
    for savefile in savefiles:
        score = evaluate(savefile, pedantic=args.pedantic)
        scores.append((savefile, score))

    # Print leaderboard
    scores.sort(key=lambda x: x[1])
    if len(savefiles) > 1:
        logging.info(f"{ul}Leaderboard:{end}")
        for index, (file, score) in enumerate(scores, start=1):
            logging.info(f"#{index} - {ylw}{score:5.2f} TOTAL{end} - {file.name}")