import argparse
import os.path
import re
import sys
from dataclasses import dataclass

XADD = 150
YADD = 300
ZDOWN = 62.4
ZUP = ZDOWN + 5

ANSI = "\033[0;{}m"
OFF = ANSI.format(0)
RED = ANSI.format(31)
GREEN = ANSI.format(32)

MINX = 100
MAXX = 650
MINY = 200
MAXY = 1000

OOZA_HEADER_PATH = "ooza/config.g"
NUMBER_RE = r"(\d*)?(\.)?(\d*)?"


def get_number_re(prefix):
    pattern = f"{prefix}{NUMBER_RE}"
    return re.compile(pattern)


X_RE = get_number_re("X")
Y_RE = get_number_re("Y")


@dataclass
class GCodeInfo:
    minx: float
    maxx: float
    miny: float
    maxy: float


def get_xy(line):
    xval = yval = None
    xmatch = X_RE.search(line)
    ymatch = Y_RE.search(line)

    if xmatch is not None:
        xval = float(xmatch.group().strip("X"))
    if ymatch is not None:
        yval = float(ymatch.group().strip("Y"))

    return xval, yval


def rewrite_lines(inlines):
    # Home the board after bed is setup
    outlines = "G28\n"

    minx = float("inf")
    miny = float("inf")
    maxx = float("-inf")
    maxy = float("-inf")

    for line in inlines:
        x, y = get_xy(line)
        prefix = line[:2]

        if x is None and y is None:
            if prefix != "M2":
                outlines += f"{line}\n"
            continue

        lines = prefix

        if x is not None:
            x += XADD
            lines += f" X{x}"
            minx = min(x, minx)
            maxx = max(x, maxx)

        if y is not None:
            y += YADD
            lines += f" Y{y}"
            miny = min(y, miny)
            maxy = max(y, maxy)

        postfix = "" if prefix == "G0" else " F1500"
        lines += f"{postfix}\n"

        # Raise and lower the pen before moving on fresh lines
        if prefix == "G0":
            lines = f"G0 Z{ZUP}\n{lines}G0 Z{ZDOWN}\n"

        outlines += lines

    # Raise pen as final control
    outlines += f"G0 Z{ZUP}\n"
    info = GCodeInfo(
        minx=minx,
        maxx=maxx,
        miny=miny,
        maxy=maxy,
    )

    return (outlines, info)


def check_okay(info):
    if info.minx < MINX or info.miny < MINY or info.maxx > MAXX or info.maxy > MAXY:
        return False
    return True


def format_info_error(info):
    output = f"{RED}Print is outside range of Ooza, please resize.{OFF}"
    if info.minx < MINX:
        output += f"\n  Min X: {int(info.minx)} < {MINX}"
    if info.maxx > MAXX:
        output += f"\n  Max X: {int(info.maxx)} > {MAXX}"
    if info.miny < MINY:
        output += f"\n  Min Y: {int(info.miny)} < {MINY}"
    if info.maxy > MAXY:
        output += f"\n  Max Y: {int(info.maxy)} > {MAXY}"
    return output


class ExtensionFileType(argparse.FileType):
    def __init__(self, ext=None, **kwargs):
        super().__init__(**kwargs)
        self.ext = ext

    def __call__(self, string):
        if self.ext is not None:
            _, ext = os.path.splitext(string)
            if not ext or ext != self.ext:
                message = f"Expected extension {self.ext} but got {ext}"
                raise argparse.ArgumentTypeError(message)

        return super().__call__(string)



def rewrite_gcode(infile, outfile):
    lines = infile.read().splitlines()
    lines, info = rewrite_lines(lines)

    if not check_okay(info):
        print(format_info_error(info))
        return 1

    with open(OOZA_HEADER_PATH) as fp:
        header = fp.read()

    outfile.write(header + lines)
    return 0

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("infile", type=ExtensionFileType(ext=".g", mode="r"))
    parser.add_argument("outfile", type=ExtensionFileType(ext=".g", mode="w"))
    args = parser.parse_args()
    sys.exit(rewrite_gcode(args.infile, args.outfile))