#!/usr/bin/python3

import argparse
import logging
import sys

from .instructions import Instruction
from .instructions.inc import Inc
from .arguments import ArgumentType, Argument

from typing import Callable, Dict, List, Optional

logger = logging.getLogger(__name__)

COMMENT_CHAR = '#'
LABEL_SUFFIX = ':'

GB_INSTRUCTIONS = [
    Inc()
]

def build_instruction_map() -> Dict[str, Instruction]:
    d = {}  # type: Dict[str, Instruction]
    for i in GB_INSTRUCTIONS:
        d[i.token] = i
    return d

def try_parse_arguments(args: List[str],
                        arg_types: List[ArgumentType]) -> Optional[List[Argument]]:
    if len(args) != len(arg_types):
        return None

    out_args = []

    for (arg, arg_type) in zip(args, arg_types):
        try:
            out_args.append(arg_type.parse(arg))
        except ValueError:
            return None

    return out_args

def parse_line_size(instruction: Instruction,
               arguments: List[str]) -> bytes:

    for argtype_list in instruction.argument_specs:
        args = try_parse_arguments(arguments, argtype_list)
        if args is not None:
            return instruction.num_bytes(args)

    raise ValueError("Failed to parse line.")

def parse_line_bytes(instruction: Instruction,
               arguments: List[str],
               label_resolver: Callable[[str], int]) -> bytes:

    for argtype_list in instruction.argument_specs:
        args = try_parse_arguments(arguments, argtype_list)
        if args is not None:
            return instruction.to_bytes(args, label_resolver)

    raise ValueError("Failed to parse line.")

def assemble_file(infile) -> bytes:
    program = infile.readlines()
    return assemble(infile)

def assemble(lines: str) -> bytes:

    instruction_map = build_instruction_map()
    logger.debug("Instruction map: {}".format(instruction_map))

    byte_offset = 0
    instruction_count = 0
    labels = {}  # type: Dict[str, int]
    program = bytes()

    def label_resolver(label: str) -> int:
        nonlocal labels
        return labels[label]

    for step in ["SIZE", "CONTENT"]:
        logger.debug("Starting step: {}".format(step))

        for line_num, line in enumerate(lines):
            # Remove comments
            line = line.split(COMMENT_CHAR)[0]

            # Tokenize
            tokens = line.split()
            logging.info("Line:", line)
            logging.info("Tokens:", tokens)

            if len(tokens) == 0:
                continue

            instruction_name = tokens[0]
            args = tokens[1:]
            try:
                instruction = instruction_map[instruction_name]
            except KeyError:
                if instruction_name[-1] == LABEL_SUFFIX:
                    if step == 'SIZE':
                        label = instruction_name[:-1]
                        logger.debug("Found label '{}' at {}"
                                     .format(label, byte_offset))
                        if label in labels.keys():
                            raise KeyError("Label '{}' defined at {} and {}"
                                           .format(label, labels[label], line_num))
                        labels[label] = byte_offset
                    continue
                raise KeyError("Unknown instruction \"{}\" on line {}"
                               .format(instruction_name, line_num))

            if step == 'SIZE':
                byte_offset += parse_line_size(instruction, args)
                instruction_count += 1
            if step == 'CONTENT':
                try:
                    program += parse_line_bytes(instruction, args, label_resolver)
                except ValueError:
                    raise ValueError("Failed to parse line {},\n{}"
                                     .format(line_num, line))
        if step == 'SIZE':
            logger.info("Program size: {} bytes, {} instructions"
                        .format(byte_offset, instruction_count))
            logger.debug("Found labels: {}".format(labels))

    return program

def main() -> None:
    parser = argparse.ArgumentParser(
        description= "An assembler for Gameboy assembly")

    parser.add_argument("--infile", "-i", type=argparse.FileType("r"), default=sys.stdin)
    parser.add_argument("--outfile", "-o", type=argparse.FileType("wb"), default=sys.stdout)
    parser.add_argument("--verbose", "-v", action='store_true')

    args = parser.parse_args()

    logging.basicConfig(format="%(levelname)s: %(message)s")
    logger.setLevel(logging.INFO)

    if args.verbose:
        logging.basicConfig(format="%(levelname)s: %(filename)s:%(lineno)d: %(message)s")
        logger.setLevel(logging.DEBUG)

    program = assemble_file(args.infile)
    outfile.write(program)


if __name__ == "__main__":
    main()