#!/usr/bin/env python3
import argparse
import os
import re
import subprocess
import sys
import tempfile


DEFAULT_REFERENCE_SIZE = (2560, 1440)
DEFAULT_REFERENCE_CROP = (1050, 620, 100, 60)
DEFAULT_FONT = "Roboto-Condensed"
DEFAULT_POINTSIZE = 43
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))


def run_bytes(args):
    try:
        return subprocess.run(args, check=True, stdout=subprocess.PIPE).stdout
    except FileNotFoundError:
        sys.exit(f"missing required command: {args[0]}")
    except subprocess.CalledProcessError as exc:
        sys.exit(f"command failed ({exc.returncode}): {' '.join(args)}")


def ppm_from_magick(args):
    data = run_bytes(["magick", *args, "-alpha", "off", "-depth", "8", "ppm:-"])
    idx = 0

    def next_token():
        nonlocal idx
        while idx < len(data) and data[idx] in b" \t\r\n":
            idx += 1
        if idx < len(data) and data[idx] == ord("#"):
            while idx < len(data) and data[idx] not in b"\r\n":
                idx += 1
            return next_token()
        start = idx
        while idx < len(data) and data[idx] not in b" \t\r\n":
            idx += 1
        return data[start:idx]

    magic = next_token()
    if magic != b"P6":
        sys.exit("ImageMagick did not return binary PPM data")

    width = int(next_token())
    height = int(next_token())
    max_value = int(next_token())
    if max_value != 255:
        sys.exit(f"unsupported PPM max value: {max_value}")

    while idx < len(data) and data[idx] in b" \t\r\n":
        idx += 1

    return width, height, data[idx:]


def image_size(image_path):
    output = run_bytes(["magick", "identify", "-format", "%w %h", image_path]).decode("ascii")
    width, height = output.split()
    return int(width), int(height)


def orange_digit_mask(width, height, raw):
    mask = [[False] * width for _ in range(height)]
    for y in range(height):
        row_offset = y * width * 3
        for x in range(width):
            offset = row_offset + x * 3
            r, g, b = raw[offset], raw[offset + 1], raw[offset + 2]
            mask[y][x] = r >= 115 and 55 <= g <= 205 and b <= 100 and r - g >= 20
    return mask


def white_mask(width, height, raw):
    mask = [[False] * width for _ in range(height)]
    for y in range(height):
        row_offset = y * width * 3
        for x in range(width):
            offset = row_offset + x * 3
            r, g, b = raw[offset], raw[offset + 1], raw[offset + 2]
            mask[y][x] = r >= 120 and g >= 120 and b >= 120
    return mask


def connected_components(mask):
    height = len(mask)
    width = len(mask[0]) if height else 0
    seen = [[False] * width for _ in range(height)]
    components = []

    for y in range(height):
        for x in range(width):
            if seen[y][x] or not mask[y][x]:
                continue

            stack = [(x, y)]
            seen[y][x] = True
            points = set()

            while stack:
                cx, cy = stack.pop()
                points.add((cx, cy))
                for ny in range(cy - 1, cy + 2):
                    for nx in range(cx - 1, cx + 2):
                        if nx < 0 or ny < 0 or nx >= width or ny >= height:
                            continue
                        if seen[ny][nx] or not mask[ny][nx]:
                            continue
                        seen[ny][nx] = True
                        stack.append((nx, ny))

            xs = [point[0] for point in points]
            ys = [point[1] for point in points]
            bbox = (min(xs), min(ys), max(xs), max(ys))
            box_width = bbox[2] - bbox[0] + 1
            box_height = bbox[3] - bbox[1] + 1

            if len(points) >= 40 and box_width >= 5 and box_height >= 15:
                components.append({"points": points, "bbox": bbox, "area": len(points)})

    components.sort(key=lambda component: component["bbox"][0])
    return components


def normalize(component, width=24, height=36):
    x1, y1, x2, y2 = component["bbox"]
    source_width = x2 - x1 + 1
    source_height = y2 - y1 + 1
    points = component["points"]
    output = []

    for y in range(height):
        source_y = y1 + int((y + 0.5) * source_height / height)
        row = []
        for x in range(width):
            source_x = x1 + int((x + 0.5) * source_width / width)
            row.append((source_x, source_y) in points)
        output.append(row)

    return output


def template_distance(left, right):
    total = len(left) * len(left[0])
    differences = 0
    intersection = 0
    union = 0

    for left_row, right_row in zip(left, right):
        for left_value, right_value in zip(left_row, right_row):
            differences += left_value != right_value
            intersection += left_value and right_value
            union += left_value or right_value

    hamming = differences / total
    iou = intersection / union if union else 0.0
    return hamming, iou


def build_templates(font, pointsize):
    width, height, raw = ppm_from_magick(
        [
            "-background",
            "black",
            "-fill",
            "white",
            "-font",
            font,
            "-pointsize",
            str(pointsize),
            "label:0123456789",
        ]
    )
    components = connected_components(white_mask(width, height, raw))
    if len(components) < 10:
        sys.exit(f"template rendering produced only {len(components)} digit components")
    return {str(index): normalize(component) for index, component in enumerate(components[:10])}


def classify(component, templates):
    sample = normalize(component)
    ranked = []
    for digit, template in templates.items():
        hamming, iou = template_distance(sample, template)
        ranked.append((hamming, -iou, digit, iou))
    ranked.sort()
    hamming, negative_iou, digit, iou = ranked[0]
    return digit, hamming, iou


def helper_path():
    return os.path.join(SCRIPT_DIR, "kwin_capture_screen")


def capture_named_screen(output, screen):
    helper = helper_path()
    if not os.path.exists(helper):
        sys.exit(f"missing {helper}; run `make` in {SCRIPT_DIR}")
    subprocess.run(
        [helper, screen, output],
        check=True,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )


def parse_geometry(geometry):
    match = re.match(r"(-?\d+),(-?\d+),(\d+)x(\d+)", geometry)
    if not match:
        raise ValueError(f"could not parse screen geometry: {geometry}")
    return tuple(int(group) for group in match.groups())


def screen_geometry(screen_name):
    screens = kwin_screens()
    for screen in screens:
        if screen.get("name") == screen_name:
            if not screen.get("enabled"):
                sys.exit(f"screen {screen_name} is disabled")
            geometry = screen.get("geometry")
            if not geometry:
                sys.exit(f"screen {screen_name} has no geometry")
            return parse_geometry(geometry)
    available = ", ".join(screen.get("name", "unknown") for screen in screens)
    sys.exit(f"unknown screen {screen_name}; available screens: {available}")


def virtual_desktop_origin():
    enabled_geometries = [
        parse_geometry(screen["geometry"])
        for screen in kwin_screens()
        if screen.get("enabled") and screen.get("geometry")
    ]
    if not enabled_geometries:
        return 0, 0
    return min(geometry[0] for geometry in enabled_geometries), min(geometry[1] for geometry in enabled_geometries)


def capture_screen_via_fullscreen_crop(output, screen):
    x, y, width, height = screen_geometry(screen)
    origin_x, origin_y = virtual_desktop_origin()
    crop_x = x - origin_x
    crop_y = y - origin_y

    fd, full_path = tempfile.mkstemp(prefix="reforger-fullscreen-", suffix=".png")
    os.close(fd)
    try:
        capture_spectacle(full_path, "fullscreen")
        run_bytes(["magick", full_path, "-crop", f"{width}x{height}+{crop_x}+{crop_y}", "+repage", output])
    finally:
        try:
            os.unlink(full_path)
        except OSError:
            pass


def capture_spectacle(output, mode):
    if mode == "fullscreen":
        capture_arg = "--fullscreen"
    elif mode == "current":
        capture_arg = "--current"
    else:
        raise ValueError(f"unsupported capture mode: {mode}")

    subprocess.run(
        ["spectacle", "--background", capture_arg, "--nonotify", "--output", output],
        check=True,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )


def capture_screenshot(output, mode, screen=None):
    if screen:
        try:
            capture_named_screen(output, screen)
        except subprocess.CalledProcessError:
            capture_screen_via_fullscreen_crop(output, screen)
        return
    capture_spectacle(output, mode)


def kwin_screens():
    output = run_bytes(["qdbus6", "org.kde.KWin", "/KWin", "supportInformation"]).decode("utf-8", "replace")
    screens = []
    current = None

    for line in output.splitlines():
        if re.match(r"^Screen \d+:", line):
            if current:
                screens.append(current)
            current = {}
            continue
        if current is None:
            continue

        stripped = line.strip()
        if stripped.startswith("Name:"):
            current["name"] = stripped.split(":", 1)[1].strip()
        elif stripped.startswith("Enabled:"):
            current["enabled"] = stripped.split(":", 1)[1].strip() == "1"
        elif stripped.startswith("Geometry:"):
            current["geometry"] = stripped.split(":", 1)[1].strip()

    if current:
        screens.append(current)

    return screens


def print_screens():
    for screen in kwin_screens():
        enabled = "enabled" if screen.get("enabled") else "disabled"
        geometry = screen.get("geometry", "no geometry")
        print(f"{screen.get('name', 'unknown')}\t{enabled}\t{geometry}")


def load_queue_image(image_path, crop, already_cropped=False):
    if already_cropped:
        return ppm_from_magick([image_path, "+repage"])

    x, y, width, height = crop
    return ppm_from_magick([image_path, "+repage", "-crop", f"{width}x{height}+{x}+{y}", "+repage"])


def read_queue_number(image_path, crop, font, pointsize, already_cropped=False):
    width, height, raw = load_queue_image(image_path, crop, already_cropped)
    components = connected_components(orange_digit_mask(width, height, raw))
    templates = build_templates(font, pointsize)

    digits = []
    details = []
    for component in components:
        digit, hamming, iou = classify(component, templates)
        digits.append(digit)
        details.append((digit, component["bbox"], component["area"], hamming, iou))

    return "".join(digits), details


def parse_crop(value):
    parts = value.replace(",", " ").split()
    if len(parts) != 4:
        raise argparse.ArgumentTypeError("crop must be four numbers: x,y,width,height")
    try:
        return tuple(int(part) for part in parts)
    except ValueError as exc:
        raise argparse.ArgumentTypeError("crop values must be integers") from exc


def parse_size(value):
    parts = value.lower().replace(",", "x").split("x")
    if len(parts) != 2:
        raise argparse.ArgumentTypeError("size must be WIDTHxHEIGHT")
    try:
        return int(parts[0]), int(parts[1])
    except ValueError as exc:
        raise argparse.ArgumentTypeError("size values must be integers") from exc


def scale_crop(reference_crop, reference_size, target_size, scale_mode):
    reference_width, reference_height = reference_size
    target_width, target_height = target_size
    x, y, width, height = reference_crop

    if scale_mode == "width":
        x_scale = target_width / reference_width
        y_scale = x_scale
    elif scale_mode == "independent":
        x_scale = target_width / reference_width
        y_scale = target_height / reference_height
    else:
        raise ValueError(f"unsupported scale mode: {scale_mode}")

    return (
        round(x * x_scale),
        round(y * y_scale),
        max(1, round(width * x_scale)),
        max(1, round(height * y_scale)),
    )


def resolve_crop(image_path, explicit_crop, reference_crop, reference_size, scale_mode, already_cropped):
    if already_cropped:
        return None
    if explicit_crop:
        return explicit_crop
    return scale_crop(reference_crop, reference_size, image_size(image_path), scale_mode)


def main():
    parser = argparse.ArgumentParser(description="Read Arma Reforger queue number from a fixed screen crop.")
    parser.add_argument("--image", help="Read from an existing image.")
    parser.add_argument("--dataset", help="Read every image in a dataset directory.")
    parser.add_argument("--list-screens", action="store_true", help="List KWin output names for use with --screen.")
    parser.add_argument("--cropped", action="store_true", help="Treat input image(s) as already cropped to the queue number.")
    parser.add_argument("--capture", action="store_true", help="Capture the current monitor with Spectacle.")
    parser.add_argument("--screen", help="KWin output name to capture explicitly, for example DP-3.")
    parser.add_argument(
        "--capture-mode",
        choices=("current", "fullscreen"),
        default="current",
        help="Spectacle capture mode. Use current for one monitor, fullscreen for the whole desktop.",
    )
    parser.add_argument("--crop", type=parse_crop, help="Absolute crop as x,y,width,height.")
    parser.add_argument(
        "--reference-crop",
        type=parse_crop,
        default=DEFAULT_REFERENCE_CROP,
        help="Reference crop as x,y,width,height, scaled to the input image size when --crop is omitted.",
    )
    parser.add_argument(
        "--reference-size",
        type=parse_size,
        default=DEFAULT_REFERENCE_SIZE,
        help="Reference screen size for --reference-crop scaling.",
    )
    parser.add_argument(
        "--scale-mode",
        choices=("width", "independent"),
        default="width",
        help="How to scale --reference-crop. Width mode preserves a 16:9 UI layout and is the default.",
    )
    parser.add_argument("--font", default=DEFAULT_FONT, help="Font used for synthetic digit templates.")
    parser.add_argument("--pointsize", type=int, default=DEFAULT_POINTSIZE, help="Point size used for templates.")
    parser.add_argument("--debug", action="store_true", help="Print component and match details.")
    parser.add_argument("--show-crop", action="store_true", help="Print the resolved crop used for each full-size image.")
    parser.add_argument(
        "--expect-filenames",
        action="store_true",
        help="In dataset mode, compare detected values to each image filename stem.",
    )
    args = parser.parse_args()

    if args.screen and not args.capture:
        parser.error("--screen can only be used with --capture")

    if sum(bool(value) for value in (args.image, args.dataset, args.capture, args.list_screens)) != 1:
        parser.error("choose exactly one of --image, --dataset, --capture, or --list-screens")

    if args.list_screens:
        print_screens()
        return 0

    if args.dataset:
        extensions = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
        paths = [
            os.path.join(args.dataset, name)
            for name in sorted(os.listdir(args.dataset))
            if os.path.splitext(name)[1].lower() in extensions
        ]
        if not paths:
            print("no dataset images found")
            return 2

        status = 0
        passed = 0
        failed = 0
        for path in paths:
            crop = resolve_crop(path, args.crop, args.reference_crop, args.reference_size, args.scale_mode, args.cropped)
            if args.show_crop and crop:
                print(f"{path}\tcrop={crop[0]},{crop[1]},{crop[2]},{crop[3]}")
            number, details = read_queue_number(path, crop, args.font, args.pointsize, args.cropped)
            if not number:
                status = 2
                failed += 1
                suffix = ""
                if args.expect_filenames:
                    expected = os.path.splitext(os.path.basename(path))[0]
                    suffix = f"\texpected={expected}\tFAIL"
                print(f"{path}\tno digits found{suffix}")
                continue

            suffix = ""
            if args.expect_filenames:
                expected = os.path.splitext(os.path.basename(path))[0]
                if number == expected:
                    passed += 1
                    suffix = f"\texpected={expected}\tPASS"
                else:
                    failed += 1
                    status = 1
                    suffix = f"\texpected={expected}\tFAIL"

            print(f"{path}\t{number}{suffix}")
            if args.debug:
                for digit, bbox, area, hamming, iou in details:
                    print(f"  digit={digit} bbox={bbox} area={area} hamming={hamming:.3f} iou={iou:.3f}")
        if args.expect_filenames:
            print(f"summary: {passed} passed, {failed} failed")
        return status

    temp_path = None
    image_path = args.image
    if args.capture:
        fd, temp_path = tempfile.mkstemp(prefix="reforger-queue-", suffix=".png")
        os.close(fd)
        image_path = temp_path
        capture_screenshot(image_path, args.capture_mode, args.screen)

    try:
        crop = resolve_crop(image_path, args.crop, args.reference_crop, args.reference_size, args.scale_mode, args.cropped)
        if args.show_crop and crop:
            print(f"crop={crop[0]},{crop[1]},{crop[2]},{crop[3]}")
        number, details = read_queue_number(image_path, crop, args.font, args.pointsize, args.cropped)
    finally:
        if temp_path:
            try:
                os.unlink(temp_path)
            except OSError:
                pass

    if not number:
        print("no digits found")
        return 2

    print(number)
    if args.debug:
        for digit, bbox, area, hamming, iou in details:
            print(f"digit={digit} bbox={bbox} area={area} hamming={hamming:.3f} iou={iou:.3f}")

    return 0


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