#!/usr/bin/env python3
import argparse
import os
import re
import glob
from collections import Counter
from PIL import Image, ImageDraw, ImageFont

def detect_prefix(in_dir):
    """
    Возвращает определённый префикс по файлам в папке.
    Ищем имена по шаблону: <prefix>_e######_##_*.png
    Если найден ровно один префикс — возвращаем его.
    Если найдено несколько — возвращаем наиболее частый и предупреждаем.
    Если ничего не найдено — возвращаем None.
    """
    pat = re.compile(r"^([^_]+)_e(\d{6})_(\d{2})_.+\.png$", re.IGNORECASE)
    prefixes = []
    for p in glob.glob(os.path.join(in_dir, "*.png")):
        m = pat.match(os.path.basename(p))
        if m:
            prefixes.append(m.group(1))
    if not prefixes:
        return None
    cnt = Counter(prefixes)
    if len(cnt) == 1:
        return prefixes[0]
    # несколько префиксов — выбрать самый частый
    most_common, freq = cnt.most_common(1)[0]
    print(f"Найдено несколько префиксов; выбран наиболее частый: '{most_common}' (встречается {freq} раз).")
    return most_common

def build_grid_for_prefix(in_dir, prefix, out_name, args):
    """
    Основная логика: собирает картинки с данным префиксом в сетку и сохраняет.
    """
    if out_name is None:
        out_name = f"{prefix}_grid.png"
    out = os.path.join(in_dir, out_name)

    pat = re.compile(rf"^{re.escape(prefix)}_e(\d{{6}})_(\d{{2}})_.+\.png$", re.IGNORECASE)

    paths = [p for p in glob.glob(os.path.join(in_dir, "*.png")) if pat.search(os.path.basename(p))]
    if not paths:
        print(f"[{prefix}] Не нашёл файлов по шаблону {prefix}_e******_**_*.png в {in_dir}")
        return

    items = []
    for p in paths:
        m = pat.search(os.path.basename(p))
        epoch = int(m.group(1))
        prompt_idx = int(m.group(2))
        items.append((epoch, prompt_idx, p))

    epochs = sorted({e for e, _, _ in items})
    rows = sorted({i for _, i, _ in items})
    table = {(e, i): p for e, i, p in items}

    # Определяем максимальный размер ячеек
    max_w = max_h = 0
    for _, _, p in items:
        with Image.open(p) as im:
            max_w = max(max_w, im.width)
            max_h = max(max_h, im.height)
    cell_w, cell_h = max_w, max_h

    # Параметры (можно переопределить через args)
    MARGIN = args.margin
    PAD = args.padding
    GX = args.gapx
    GY = args.gapy
    HEADER = args.header
    ROWLBL = args.rowlbl
    BG = (12, 12, 12)
    FG = (230, 230, 230)

    W = MARGIN + ROWLBL + len(epochs) * (cell_w + 2 * PAD) + (len(epochs) - 1) * GX + MARGIN
    H = MARGIN + HEADER + len(rows) * (cell_h + 2 * PAD) + (len(rows) - 1) * GY + MARGIN

    canvas = Image.new("RGB", (W, H), BG)
    draw = ImageDraw.Draw(canvas)
    try:
        font = ImageFont.truetype("arial.ttf", 18)
        font_small = ImageFont.truetype("arial.ttf", 14)
    except Exception:
        font = ImageFont.load_default()
        font_small = ImageFont.load_default()

    # Заголовки колонок (эпохи)
    for c, e in enumerate(epochs):
        x0 = MARGIN + ROWLBL + c * (cell_w + 2 * PAD + GX)
        xc = x0 + (cell_w + 2 * PAD) // 2
        draw.text((xc, MARGIN + HEADER // 2), f"e{e:06d}", fill=FG, anchor="mm", font=font)

    # Подписи строк (prompts)
    for r, i in enumerate(rows):
        y0 = MARGIN + HEADER + r * (cell_h + 2 * PAD + GY)
        yc = y0 + (cell_h + 2 * PAD) // 2
        draw.text((MARGIN + ROWLBL // 2, yc), f"p{i:02d}", fill=FG, anchor="mm", font=font)

    # Размещение изображений
    for r, i in enumerate(rows):
        for c, e in enumerate(epochs):
            x0 = MARGIN + ROWLBL + c * (cell_w + 2 * PAD + GX) + PAD
            y0 = MARGIN + HEADER + r * (cell_h + 2 * PAD + GY) + PAD
            box = (x0, y0, x0 + cell_w, y0 + cell_h)
            draw.rectangle(box, outline=(60, 60, 60), width=1)
            p = table.get((e, i))
            if not p:
                draw.text((x0 + cell_w // 2, y0 + cell_h // 2), "N/A", fill=(160, 160, 160), anchor="mm", font=font_small)
                continue
            with Image.open(p) as im:
                scale = min(cell_w / im.width, cell_h / im.height, 1.0)
                nw, nh = int(im.width * scale), int(im.height * scale)
                if scale != 1.0:
                    im = im.resize((nw, nh), Image.LANCZOS)
                ox = x0 + (cell_w - nw) // 2
                oy = y0 + (cell_h - nh) // 2
                canvas.paste(im, (ox, oy))

    canvas.save(out, "PNG")
    print(f"[{prefix}] Сохранено: {out}")
    print(f"[{prefix}] Эпох: {len(epochs)}  | Промптов: {len(rows)}  | Картинок: {len(items)}")

def main():
    ap = argparse.ArgumentParser(description="Склейка превью в сетку: строки=промпты (00..), столбцы=эпохи (e000001..)")
    ap.add_argument("--dir", default=".", help="папка с PNG (по умолчанию текущая)")
    ap.add_argument("--prefix", default=None, help="префикс имён файлов (до '_'). Если не указан — будет определён автоматически")
    ap.add_argument("--out", default=None, help="имя итогового PNG (по умолчанию: <prefix>_grid.png)")
    ap.add_argument("--margin", type=int, default=16)
    ap.add_argument("--padding", type=int, default=10)
    ap.add_argument("--gapx", type=int, default=8)
    ap.add_argument("--gapy", type=int, default=8)
    ap.add_argument("--header", type=int, default=40)
    ap.add_argument("--rowlbl", type=int, default=60)
    args = ap.parse_args()

    in_dir = os.path.abspath(args.dir)

    if args.prefix:
        prefix = args.prefix
    else:
        prefix = detect_prefix(in_dir)
        if prefix is None:
            print(f"Не нашёл файлов, соответствующих шаблону <prefix>_e######_##_*.png в папке {in_dir}")
            return
        print("Авто-определён префикс:", prefix)

    build_grid_for_prefix(in_dir, prefix, args.out, args)

if __name__ == "__main__":
    main()