#!/usr/bin/env python3

from __future__ import annotations

import argparse
import curses
import random
import shutil
import struct
import zlib
from dataclasses import dataclass
from pathlib import Path
from typing import Callable


DEFAULT_SIZES = "320x240,640x480,1024x768,1600x900"


@dataclass
class GeneratorConfig:
    output: Path
    sizes: list[tuple[int, int]]
    variants_per_size: int
    folder_count: int
    images_per_folder: int
    folder_prefix: str
    file_prefix: str
    seed: int
    clear_output: bool


def parse_sizes(value: str) -> list[tuple[int, int]]:
    sizes: list[tuple[int, int]] = []
    for item in value.split(","):
        item = item.strip().lower()
        if not item:
            continue
        if "x" not in item:
            raise ValueError(f"Invalid size '{item}'. Use WIDTHxHEIGHT,WIDTHxHEIGHT")
        width_text, height_text = item.split("x", 1)
        try:
            width = int(width_text)
            height = int(height_text)
        except ValueError as exc:
            raise ValueError(
                f"Invalid size '{item}'. Width and height must be integers."
            ) from exc
        if width <= 0 or height <= 0:
            raise ValueError(
                f"Invalid size '{item}'. Width and height must be positive."
            )
        sizes.append((width, height))
    if not sizes:
        raise ValueError("At least one image size is required.")
    return sizes


def build_png(width: int, height: int, variant_seed: int) -> bytes:
    rng = random.Random(variant_seed)
    base_r = rng.randint(20, 230)
    base_g = rng.randint(20, 230)
    base_b = rng.randint(20, 230)
    accent_r = (base_r + rng.randint(40, 120)) % 256
    accent_g = (base_g + rng.randint(40, 120)) % 256
    accent_b = (base_b + rng.randint(40, 120)) % 256
    block = max(12, min(width, height) // max(6, rng.randint(6, 12)))
    raw = bytearray()

    for y in range(height):
        raw.append(0)
        for x in range(width):
            mix_x = x / max(1, width - 1)
            mix_y = y / max(1, height - 1)
            checker = ((x // block) + (y // block) + variant_seed) % 2
            stripe = ((x // max(4, block // 2)) + variant_seed) % 3
            noise = rng.randint(-12, 12)

            r = int(base_r * (1 - mix_x) + accent_r * mix_x)
            g = int(base_g * (1 - mix_y) + accent_g * mix_y)
            b = int((base_b + accent_b) / 2)

            if checker:
                r = min(255, r + 18)
                g = max(0, g - 10)
            if stripe == 0:
                b = min(255, b + 28)
            elif stripe == 1:
                r = max(0, r - 18)

            raw.extend(
                (
                    max(0, min(255, r + noise)),
                    max(0, min(255, g - noise)),
                    max(0, min(255, b + noise // 2)),
                )
            )

    def chunk(chunk_type: bytes, data: bytes) -> bytes:
        return (
            struct.pack(">I", len(data))
            + chunk_type
            + data
            + struct.pack(">I", zlib.crc32(chunk_type + data) & 0xFFFFFFFF)
        )

    header = b"\x89PNG\r\n\x1a\n"
    ihdr = struct.pack(">IIBBBBB", width, height, 8, 2, 0, 0, 0)
    idat = zlib.compress(bytes(raw), level=6)
    return header + chunk(b"IHDR", ihdr) + chunk(b"IDAT", idat) + chunk(b"IEND", b"")


def build_variants(
    sizes: list[tuple[int, int]], variants_per_size: int, seed: int
) -> list[tuple[str, bytes]]:
    variants: list[tuple[str, bytes]] = []
    for size_index, (width, height) in enumerate(sizes):
        for variant_index in range(variants_per_size):
            variant_seed = seed + (size_index * 10_000) + variant_index
            name = f"{width}x{height}-v{variant_index + 1}.png"
            variants.append((name, build_png(width, height, variant_seed)))
    return variants


def write_collection(
    output_dir: Path,
    folder_count: int,
    images_per_folder: int,
    folder_prefix: str,
    file_prefix: str,
    variants: list[tuple[str, bytes]],
    progress: Callable[[int, int, int], None] | None = None,
) -> int:
    total_files = 0
    for folder_number in range(1, folder_count + 1):
        folder = output_dir / f"{folder_prefix}-{folder_number:03d}"
        folder.mkdir(parents=True, exist_ok=True)

        for image_number in range(1, images_per_folder + 1):
            variant_name, variant_bytes = variants[(image_number - 1) % len(variants)]
            destination = folder / f"{file_prefix}-{image_number:05d}-{variant_name}"
            destination.write_bytes(variant_bytes)
            total_files += 1

        if progress:
            progress(folder_number, folder_count, total_files)

    return total_files


def validate_config(config: GeneratorConfig) -> None:
    if config.variants_per_size <= 0:
        raise ValueError("Variants per size must be greater than 0.")
    if config.folder_count <= 0:
        raise ValueError("Folder count must be greater than 0.")
    if config.images_per_folder <= 0:
        raise ValueError("Images per folder must be greater than 0.")
    if not config.folder_prefix.strip():
        raise ValueError("Folder prefix cannot be empty.")
    if not config.file_prefix.strip():
        raise ValueError("File prefix cannot be empty.")


def generate_collection(
    config: GeneratorConfig,
    progress: Callable[[int, int, int], None] | None = None,
) -> dict[str, str | int]:
    validate_config(config)

    output_dir = config.output.expanduser().resolve()
    if config.clear_output and output_dir.exists():
        shutil.rmtree(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    variants = build_variants(config.sizes, config.variants_per_size, config.seed)
    total_files = write_collection(
        output_dir=output_dir,
        folder_count=config.folder_count,
        images_per_folder=config.images_per_folder,
        folder_prefix=config.folder_prefix,
        file_prefix=config.file_prefix,
        variants=variants,
        progress=progress,
    )

    return {
        "output_dir": str(output_dir),
        "unique_variants": len(variants),
        "folders_created": config.folder_count,
        "files_created": total_files,
        "sizes_used": ", ".join(f"{w}x{h}" for w, h in config.sizes),
    }


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="Generate fake placeholder image collections for testing."
    )
    parser.add_argument("output", nargs="?", type=Path)
    parser.add_argument("--sizes", default=DEFAULT_SIZES)
    parser.add_argument("--variants-per-size", type=int, default=3)
    parser.add_argument("--folder-count", type=int, default=5)
    parser.add_argument("--images-per-folder", type=int, default=100)
    parser.add_argument("--folder-prefix", default="batch")
    parser.add_argument("--file-prefix", default="image")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--clear-output", action="store_true")
    parser.add_argument(
        "--no-tui",
        action="store_true",
        help="Run in non-interactive mode using command-line options.",
    )
    return parser


def config_from_args(args: argparse.Namespace) -> GeneratorConfig:
    if args.output is None:
        raise ValueError("Output directory is required in --no-tui mode.")
    return GeneratorConfig(
        output=args.output,
        sizes=parse_sizes(args.sizes),
        variants_per_size=args.variants_per_size,
        folder_count=args.folder_count,
        images_per_folder=args.images_per_folder,
        folder_prefix=args.folder_prefix,
        file_prefix=args.file_prefix,
        seed=args.seed,
        clear_output=args.clear_output,
    )


def cli_mode(args: argparse.Namespace) -> int:
    try:
        config = config_from_args(args)
        summary = generate_collection(config)
    except ValueError as exc:
        raise SystemExit(str(exc)) from exc

    print(f"Output directory: {summary['output_dir']}")
    print(f"Unique variants: {summary['unique_variants']}")
    print(f"Folders created: {summary['folders_created']}")
    print(f"Files created: {summary['files_created']}")
    print(f"Sizes used: {summary['sizes_used']}")
    return 0


def truncate(text: str, width: int) -> str:
    if width <= 0:
        return ""
    if len(text) <= width:
        return text
    if width <= 3:
        return text[:width]
    return text[: width - 3] + "..."


def prompt_text(
    stdscr: curses.window, label: str, current_value: str, message: str
) -> str | None:
    curses.curs_set(1)
    height, width = stdscr.getmaxyx()
    prompt = f"{label}: "
    stdscr.move(height - 2, 0)
    stdscr.clrtoeol()
    stdscr.addstr(height - 2, 0, truncate(prompt + current_value, width - 1))
    stdscr.move(height - 1, 0)
    stdscr.clrtoeol()
    stdscr.addstr(height - 1, 0, truncate(message, width - 1))
    buffer = list(current_value)
    position = len(buffer)

    while True:
        stdscr.move(height - 2, 0)
        stdscr.clrtoeol()
        display = prompt + "".join(buffer)
        stdscr.addstr(height - 2, 0, truncate(display, width - 1))
        cursor_x = min(len(prompt) + position, max(0, width - 2))
        stdscr.move(height - 2, cursor_x)
        key = stdscr.getch()

        if key in (10, 13):
            curses.curs_set(0)
            return "".join(buffer).strip()
        if key == 27:
            curses.curs_set(0)
            return None
        if key in (curses.KEY_BACKSPACE, 127, 8):
            if position > 0:
                position -= 1
                buffer.pop(position)
            continue
        if key == curses.KEY_LEFT and position > 0:
            position -= 1
            continue
        if key == curses.KEY_RIGHT and position < len(buffer):
            position += 1
            continue
        if 32 <= key <= 126:
            buffer.insert(position, chr(key))
            position += 1


def draw_screen(
    stdscr: curses.window,
    fields: list[tuple[str, str]],
    selected: int,
    message: str,
    status_lines: list[str],
) -> None:
    stdscr.erase()
    height, width = stdscr.getmaxyx()
    title = "Fake Media Generator"
    subtitle = "Arrows to move, Enter to edit, Space to toggle clear-output, G to generate, Q to quit"
    stdscr.addstr(0, 0, truncate(title, width - 1), curses.A_BOLD)
    stdscr.addstr(1, 0, truncate(subtitle, width - 1))

    for index, (label, value) in enumerate(fields):
        y = 3 + index
        if y >= height - 5:
            break
        prefix = ">" if index == selected else " "
        line = f"{prefix} {label:<18} {value}"
        attr = curses.A_REVERSE if index == selected else curses.A_NORMAL
        stdscr.addstr(y, 0, truncate(line, width - 1), attr)

    status_start = min(height - 4, 4 + len(fields))
    for offset, line in enumerate(status_lines[-3:]):
        y = status_start + offset
        if y < height - 1:
            stdscr.addstr(y, 0, truncate(line, width - 1))

    stdscr.addstr(height - 1, 0, truncate(message, width - 1), curses.A_DIM)
    stdscr.refresh()


def tui_mode(args: argparse.Namespace) -> int:
    initial_output = (
        str(args.output) if args.output else str(Path.cwd() / "fake-media-output")
    )
    state: dict[str, str | bool] = {
        "Output Directory": initial_output,
        "Sizes": args.sizes,
        "Variants Per Size": str(args.variants_per_size),
        "Folder Count": str(args.folder_count),
        "Images Per Folder": str(args.images_per_folder),
        "Folder Prefix": args.folder_prefix,
        "File Prefix": args.file_prefix,
        "Seed": str(args.seed),
        "Clear Output": args.clear_output,
    }

    def run(stdscr: curses.window) -> int:
        curses.curs_set(0)
        stdscr.keypad(True)

        field_order = [
            "Output Directory",
            "Sizes",
            "Variants Per Size",
            "Folder Count",
            "Images Per Folder",
            "Folder Prefix",
            "File Prefix",
            "Seed",
            "Clear Output",
        ]
        selected = 0
        message = "Fill in the fields, then press G to generate."
        status_lines: list[str] = []

        while True:
            fields = []
            for key in field_order:
                value = state[key]
                if isinstance(value, bool):
                    display = "yes" if value else "no"
                else:
                    display = value
                fields.append((key, str(display)))

            draw_screen(stdscr, fields, selected, message, status_lines)
            key = stdscr.getch()

            if key in (ord("q"), ord("Q")):
                return 0
            if key in (curses.KEY_UP, ord("k")):
                selected = (selected - 1) % len(field_order)
                continue
            if key in (curses.KEY_DOWN, ord("j")):
                selected = (selected + 1) % len(field_order)
                continue

            current_field = field_order[selected]

            if key == ord(" ") and current_field == "Clear Output":
                state[current_field] = not bool(state[current_field])
                message = (
                    f"{current_field} set to {'yes' if state[current_field] else 'no'}."
                )
                continue

            if key in (10, 13):
                if current_field == "Clear Output":
                    state[current_field] = not bool(state[current_field])
                    message = f"{current_field} set to {'yes' if state[current_field] else 'no'}."
                    continue

                current_value = str(state[current_field])
                updated = prompt_text(
                    stdscr,
                    current_field,
                    current_value,
                    "Enter saves, Esc cancels",
                )
                if updated is not None and updated:
                    state[current_field] = updated
                    message = f"Updated {current_field}."
                elif updated == "":
                    message = f"{current_field} cannot be empty."
                else:
                    message = f"Canceled edit for {current_field}."
                continue

            if key in (ord("g"), ord("G")):
                try:
                    config = GeneratorConfig(
                        output=Path(str(state["Output Directory"])),
                        sizes=parse_sizes(str(state["Sizes"])),
                        variants_per_size=int(str(state["Variants Per Size"])),
                        folder_count=int(str(state["Folder Count"])),
                        images_per_folder=int(str(state["Images Per Folder"])),
                        folder_prefix=str(state["Folder Prefix"]),
                        file_prefix=str(state["File Prefix"]),
                        seed=int(str(state["Seed"])),
                        clear_output=bool(state["Clear Output"]),
                    )
                    status_lines = ["Generating files..."]

                    def progress(
                        folder_number: int, folder_total: int, total_files: int
                    ) -> None:
                        status_lines[:] = [
                            f"Output: {config.output.expanduser().resolve()}",
                            f"Progress: folder {folder_number}/{folder_total}",
                            f"Files written: {total_files}",
                        ]
                        draw_screen(stdscr, fields, selected, message, status_lines)

                    summary = generate_collection(config, progress=progress)
                    status_lines = [
                        f"Output: {summary['output_dir']}",
                        f"Unique variants: {summary['unique_variants']}",
                        f"Files created: {summary['files_created']}",
                    ]
                    message = "Generation complete. Press G to run again or Q to quit."
                except ValueError as exc:
                    message = str(exc)
                continue

        return 0

    return curses.wrapper(run)


def main() -> int:
    parser = build_parser()
    args = parser.parse_args()
    if args.no_tui:
        return cli_mode(args)
    return tui_mode(args)


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