from __future__ import annotations

import os
import tempfile
import zipfile
from codecs import encode
from pathlib import Path
from textwrap import wrap

import numpy as np
from PIL import Image


def _setpress(pressure: float) -> str:
    pressure_str = str(int(pressure * 10)).zfill(4)
    command_bytes = bytes("08PS  " + pressure_str, "utf-8")
    hex_command = encode(command_bytes, "hex").decode("utf-8")
    format_command = "\\x" + "\\x".join(
        hex_command[i : i + 2] for i in range(0, len(hex_command), 2)
    )

    hex_pairs = wrap(hex_command, 2)
    decimal_sum = sum(int(pair, 16) for pair in hex_pairs)
    checksum_bin = bin(decimal_sum % 256)[2:].zfill(8)
    inverted = int("".join("1" if c == "0" else "0" for c in checksum_bin), 2) + 1
    checksum_hex = hex(inverted)[2:].upper()
    format_checksum = "\\x" + "\\x".join(
        checksum_hex[i : i + 2] for i in range(0, len(checksum_hex), 2)
    )

    return "b'" + "\\x05\\x02" + format_command + format_checksum + "\\x03" + "'"


def _togglepress() -> str:
    return "b'\\x05\\x02\\x30\\x34\\x44\\x49\\x20\\x20\\x43\\x46\\x03'"


def _setpress_cmd(port: str, pressure: float, start: bool) -> str:
    insert = "{preset}" if start else ""
    return f"\n\r{insert}{port}.write({_setpress(pressure)})"


def _toggle_cmd(port: str, start: bool) -> str:
    insert = "{preset}" if start else ""
    return f"\n\r{insert}{port}.write({_togglepress()})"


def _valve_cmd(valve: int, command: int) -> str:
    return f"\n{{aux_command}}WAGO_ValveCommands({valve}, {command})\n"


def _gcode_layer(
    path_img: np.ndarray,
    color_img: np.ndarray,
    output_list: list[dict],
    pixel_size: float,
    direction: int,
    layer_number: int,
) -> int:
    mask = path_img > 0
    first_nonblack = np.where(mask.any(axis=1), mask.argmax(axis=1), -1)
    last_nonblack = np.where(
        mask.any(axis=1),
        mask.shape[1] - 1 - np.fliplr(mask).argmax(axis=1),
        -1,
    )

    stored_gcode: list[dict] = []
    nonblank_rows = np.where(first_nonblack != -1)[0]

    for idx, i in enumerate(nonblank_rows):
        f_idx, l_idx = int(first_nonblack[i]), int(last_nonblack[i])
        if f_idx == -1:
            continue

        if direction < 0:
            rng = range(f_idx, l_idx + 1)
        else:
            rng = range(l_idx, f_idx - 1, -1)
        direction *= -1

        prev_color = None
        color_len = 0
        buffer = direction
        stored_gcode.append({"X": buffer * pixel_size, "Y": 0, "Color": 0})

        for j in rng:
            this_color = int(color_img[i, j])
            if prev_color is None:
                prev_color = this_color
                color_len = 1
            elif this_color == prev_color:
                color_len += 1
            else:
                stored_gcode.append(
                    {
                        "X": direction * color_len * pixel_size,
                        "Y": 0,
                        "Color": prev_color,
                    }
                )
                color_len = 1
                prev_color = this_color

        if color_len > 0:
            stored_gcode.append(
                {
                    "X": direction * color_len * pixel_size,
                    "Y": 0,
                    "Color": prev_color,
                }
            )

        stored_gcode.append({"X": buffer * pixel_size, "Y": 0, "Color": 0})

        curr_x = l_idx if direction > 0 else f_idx
        curr_x += buffer

        if idx + 1 < len(nonblank_rows):
            next_i = int(nonblank_rows[idx + 1])
            y_travel_dist = next_i - int(i)
            nf, nl = int(first_nonblack[next_i]), int(last_nonblack[next_i])
            if nf == -1:
                continue
            next_start = nf if direction < 0 else nl
            travel_x = (next_start + buffer) - curr_x
            y_dir = -1 if layer_number % 2 == 1 else 1
            stored_gcode.append(
                {
                    "X": travel_x * pixel_size,
                    "Y": y_travel_dist * pixel_size * y_dir,
                    "Color": 0,
                }
            )

    output_list.extend(stored_gcode)
    return direction


def _sort_key(filename: str) -> int:
    digits = "".join(filter(str.isdigit, filename))
    return int(digits) if digits else 2**31


def _extract_zip_tiffs(zip_path: Path, dest: Path) -> list[Path]:
    with zipfile.ZipFile(zip_path) as archive:
        archive.extractall(dest)

    tiffs: list[Path] = []
    for root, _, files in os.walk(dest):
        for name in files:
            if name.lower().endswith((".tif", ".tiff")):
                tiffs.append(Path(root) / name)
    tiffs.sort(key=lambda p: _sort_key(p.name))
    return tiffs


def _load_grayscale(path: Path, invert: bool) -> np.ndarray:
    with Image.open(path) as image:
        array = np.array(image.convert("L"), dtype=np.uint8)
    if invert:
        array = 255 - array
    return array


def _center_on_canvas(
    img: np.ndarray, canvas_h: int, canvas_w: int, fill: int = 0
) -> np.ndarray:
    """Place `img` centred on a (canvas_h, canvas_w) canvas filled with `fill`.

    Mirrors the centring used to build the reference stack, so a shape's slice
    lines up pixel-for-pixel with the reference (motion) slice of the same layer.
    """
    h, w = img.shape[:2]
    out = np.full((canvas_h, canvas_w), fill, dtype=img.dtype)
    y_off = max(0, (canvas_h - h) // 2)
    x_off = max(0, (canvas_w - w) // 2)
    out[y_off : y_off + h, x_off : x_off + w] = img[: canvas_h, : canvas_w]
    return out


def generate_snake_path_gcode(
    zip_path: str | Path,
    shape_name: str,
    pressure: float,
    valve: int,
    port: int,
    layer_height: float = 0.8,
    fil_width: float = 0.8,
    invert: bool = True,
    increase_pressure_per_layer: float = 0.1,
    all_g1: bool = False,
    motion_tiffs: list[str] | None = None,
) -> Path:
    zip_path = Path(zip_path)
    if not zip_path.exists():
        raise FileNotFoundError(f"ZIP file not found: {zip_path}")

    work_dir = Path(tempfile.mkdtemp(prefix="tiff_gcode_"))
    extract_dir = work_dir / "tiffs"
    extract_dir.mkdir(parents=True, exist_ok=True)
    tiff_files = _extract_zip_tiffs(zip_path, extract_dir)
    if not tiff_files:
        raise ValueError("No TIFF files found in the ZIP archive.")

    off_color = 0
    com_port = f"serialPort{port}"
    color_dict: dict[int, int] = {0: 100, 255: valve}

    # Two non-flipped source image lists. The "path" images drive the nozzle
    # motion (which rows are swept, the sweep extent, the inter-layer shifts);
    # the "color" images decide the valve state (material) at each swept pixel.
    # Normally both are this shape's own slices. When reference motion tiffs are
    # supplied, motion comes from the combined reference stack while the valve is
    # still driven by this shape's slices, centred onto the reference canvas — so
    # parallel heads share one motion path but each dispenses only its geometry.
    shape_imgs = [_load_grayscale(p, invert=invert) for p in tiff_files]

    if motion_tiffs:
        motion_paths = sorted(
            (Path(p) for p in motion_tiffs), key=lambda p: _sort_key(p.name)
        )
        path_ref_list = [_load_grayscale(p, invert=invert) for p in motion_paths]
        if not path_ref_list:
            raise ValueError("No reference TIFF files provided for motion.")
        color_ref_list: list[np.ndarray] = []
        for li, motion_img in enumerate(path_ref_list):
            h_c, w_c = motion_img.shape[:2]
            if li < len(shape_imgs):
                color_ref_list.append(
                    _center_on_canvas(shape_imgs[li], h_c, w_c, fill=off_color)
                )
            else:
                # Reference is taller than this shape: move but dispense nothing.
                color_ref_list.append(np.full((h_c, w_c), off_color, dtype=np.uint8))
    else:
        path_ref_list = [im.copy() for im in shape_imgs]
        color_ref_list = [im.copy() for im in shape_imgs]

    setpress_lines = [_setpress_cmd(com_port, pressure, start=True)]
    pressure_on_lines = [_toggle_cmd(com_port, start=True)]
    pressure_off_lines = [_toggle_cmd(com_port, start=False)]

    gcode_list: list[dict] = []
    dist_sign_long = 1
    current_offsets_x: list[int] = []
    use_flip_y = False
    direction = -1

    for layers in range(len(path_ref_list)):
        current_image_ref = path_ref_list[layers]
        last_image_ref = path_ref_list[layers - 1] if layers > 0 else None
        y_ref = current_image_ref.shape[0]

        def find_first_valid_y(row: np.ndarray | None, flip: bool = False) -> int | None:
            if row is None:
                return None
            row_data = np.flip(row) if flip else row
            for j, pixel in enumerate(row_data):
                if np.any(pixel) != off_color:
                    return y_ref - 1 - j if flip else j
            return None

        last_x = last_y = None
        if current_offsets_x:
            use_flip_x = layers % 2 == 1
            last_x = current_offsets_x[-1] if use_flip_x else current_offsets_x[0]
            last_row = (
                last_image_ref[last_x] if last_image_ref is not None else None
            )
            last_y = find_first_valid_y(last_row, flip=use_flip_y)
            current_offsets_x.clear()

        current_offsets_x = [
            i for i, row in enumerate(current_image_ref) if np.any(row) != off_color
        ]

        first_x = first_y = None
        if current_offsets_x:
            use_flip_x = layers % 2 == 1
            first_x = current_offsets_x[-1] if use_flip_x else current_offsets_x[0]
            first_row = current_image_ref[first_x]
            first_y = find_first_valid_y(first_row, flip=use_flip_y)

        if None in (last_x, last_y, first_x, first_y):
            shift_x = shift_y = 0
        else:
            shift_x = (first_x - last_x) * fil_width
            shift_y = (first_y - last_y) * fil_width * dist_sign_long
            if use_flip_y:
                shift_y = -shift_y

        if len(current_offsets_x) % 2 == 1:
            use_flip_y = not use_flip_y

        if layers > 0:
            gcode_list.append(
                {"X": shift_y, "Y": shift_x, "Z": layer_height, "Color": 0}
            )

        for row in current_image_ref:
            if all(p == off_color for p in row):
                dist_sign_long = -dist_sign_long
            dist_sign_long = -dist_sign_long

        # Flip path and color together on even layers so they stay aligned.
        even_layer = (layers + 1) % 2 == 0
        ref_for_path = (
            np.flipud(current_image_ref) if even_layer else current_image_ref.copy()
        )
        current_image = (
            np.flipud(color_ref_list[layers]) if even_layer else color_ref_list[layers]
        )

        if layers == 0:
            direction = -1
        direction = _gcode_layer(
            ref_for_path,
            current_image,
            gcode_list,
            fil_width,
            direction,
            layers,
        )

    gcode_path = work_dir / f"{shape_name}_SnakePath_gcode.txt"
    pressure_cur = float(pressure)

    with open(gcode_path, "w") as f:
        f.write("G91\n")
        for line in setpress_lines:
            f.write(f"{line}\n")
        for line in pressure_on_lines:
            f.write(f"{line}\n")
        for color in color_dict:
            f.write(_valve_cmd(color_dict[color], 0))

        pressure_next: str | None = None
        for i, move in enumerate(gcode_list):
            prev_color = gcode_list[i - 1]["Color"] if i > 0 else 0
            cur_color = move["Color"]
            if prev_color != cur_color:
                if cur_color == off_color:
                    f.write(_valve_cmd(color_dict[prev_color], 0))
                else:
                    if prev_color == off_color:
                        f.write(_valve_cmd(color_dict[cur_color], 1))
                    else:
                        f.write(_valve_cmd(color_dict[cur_color], 1))
                        f.write(_valve_cmd(color_dict[prev_color], 0))

            # When all_g1 is set, every move is emitted as G1 regardless of
            # valve state; the valve commands still mark print vs travel.
            move_type = "G1" if (all_g1 or cur_color != off_color) else "G0"
            if "Z" in move:
                line = (
                    f"{move_type} X{move['X']} Y{move['Y']} Z{move['Z']} "
                    f"; Color {move['Color']}"
                )
                pressure_cur += increase_pressure_per_layer
                pressure_next = _setpress_cmd(com_port, pressure_cur, start=False)
            else:
                line = (
                    f"{move_type} X{move['X']} Y{move['Y']} ; Color {move['Color']}"
                )
                pressure_next = None

            f.write(f"{line}\n")
            if pressure_next is not None:
                f.write(f"{pressure_next}\n")
                pressure_next = None

        for color in color_dict:
            f.write(_valve_cmd(color_dict[color], 0))
        for line in pressure_off_lines:
            f.write(f"{line}\n")

    return gcode_path