# Image Ratio & Resolution Catalog Pro — enhanced
# Python 3.10+; gradio >= 3.32; pillow >= 9.0

from __future__ import annotations

import math
import re
import os
import io
from math import gcd, lcm
from typing import List, Tuple, Literal, Dict, Optional

import gradio as gr
from PIL import Image

# ──────────────────────────────────────────────────────────────────────────────
# Базовые утилиты
# ──────────────────────────────────────────────────────────────────────────────

RoundingMode = Literal["round", "floor", "ceil"]
ModelFamily = Literal["SDXL (native~1024)", "SD 1.x (native~512)", "SD 2.x (native~768)"]

DEFAULT_MODEL_FAMILY: ModelFamily = "SDXL (native~1024)"

def ensure_multiple(value: int, multiple: int, mode: RoundingMode = "round") -> int:
    if multiple <= 1:
        return max(1, value)
    q, r = divmod(value, multiple)
    if r == 0:
        return max(multiple, value)
    if mode == "floor":
        return max(multiple, q * multiple)
    if mode == "ceil":
        return (q + 1) * multiple
    lower = q * multiple
    upper = (q + 1) * multiple
    return upper if (value - lower) >= (upper - value) else lower

def simplified_ratio(width: int, height: int) -> str:
    if width <= 0 or height <= 0:
        return "—"
    g = gcd(width, height)
    return f"{width // g}:{height // g}"

def megapixels(width: int, height: int) -> float:
    return round((width * height) / 1_000_000.0, 3)

def parse_dims(text: str) -> List[Tuple[int, int]]:
    """Вытягивает все пары 'WxH' из произвольного текста."""
    dims: List[Tuple[int, int]] = []
    for w, h in re.findall(r"(\d{3,4})x(\d{3,4})", text):
        dims.append((int(w), int(h)))
    # удалим дубликаты, сохраняя порядок
    seen = set()
    uniq: List[Tuple[int, int]] = []
    for d in dims:
        if d not in seen:
            seen.add(d)
            uniq.append(d)
    return uniq

# ──────────────────────────────────────────────────────────────────────────────
# Автоподбор кратности (multiple)
# ──────────────────────────────────────────────────────────────────────────────

def pick_base_multiple_by_model(model_family: ModelFamily, target_max_side: Optional[int] = None) -> int:
    """
    Эвристика:
      - SDXL: кратность 64 (родной 1024, датасеты шагом 64).
      - SD 2.x: минимум 16..32, чаще 32 в HD.
      - SD 1.x: минимум 8..16, чаще 32 в HD.
    """
    if model_family == "SDXL (native~1024)":
        base = 64
    elif model_family == "SD 2.x (native~768)":
        base = 32
    else:  # SD 1.x
        base = 32
    if target_max_side and target_max_side >= 1536:
        base = max(base, 64)
    return base

def combine_multiples(base_multiple: int, diffusion_tile: int = 0, vae_tile: int = 0) -> int:
    """
    Объединяем ограничения: НОК(base_multiple, diffusion_tile, vae_tile).
    Нули игнорируются. Жёстко ограничим верх до 256.
    """
    m = base_multiple
    for x in (diffusion_tile, vae_tile):
        if isinstance(x, int) and x >= 8:
            m = lcm(m, x)
    return max(8, min(m, 256))

def explain_auto_multiple(model_family: ModelFamily, max_side: Optional[int], m_base: int, m_final: int,
                          diffusion_tile: int, vae_tile: int) -> str:
    reasons = [f"модель: {model_family} -> base={m_base}"]
    if max_side:
        reasons.append(f"размер цели: ~{max_side}px")
    if diffusion_tile:
        reasons.append(f"diffusion tile={diffusion_tile}")
    if vae_tile:
        reasons.append(f"VAE tile={vae_tile}")
    if m_final != m_base:
        reasons.append(f"НОК -> {m_final}")
    return " | ".join(reasons)

def autopick_multiple(model_family: ModelFamily,
                      diffusion_tile: int,
                      vae_tile: int,
                      width_hint: Optional[int],
                      height_hint: Optional[int]) -> Tuple[int, str]:
    max_side = None
    if width_hint and height_hint:
        max_side = max(width_hint, height_hint)
    m_base = pick_base_multiple_by_model(model_family, max_side)
    m_final = combine_multiples(m_base, diffusion_tile, vae_tile)
    return m_final, explain_auto_multiple(model_family, max_side, m_base, m_final, diffusion_tile, vae_tile)

# ──────────────────────────────────────────────────────────────────────────────
# Ресайзы по изображению
# ──────────────────────────────────────────────────────────────────────────────

def resize_by_shorter_side(image: Image.Image, target_shorter: int, multiple: int, mode: RoundingMode) -> Tuple[int, int]:
    w, h = image.size
    if w <= 0 or h <= 0 or target_shorter <= 0:
        return w, h
    scale = target_shorter / min(w, h)
    new_w = ensure_multiple(int(round(w * scale)), multiple, mode)
    new_h = ensure_multiple(int(round(h * scale)), multiple, mode)
    return new_w, new_h

def resize_to_pixel_count(image: Image.Image, target_pixels: int, multiple: int, mode: RoundingMode) -> Tuple[int, int]:
    w, h = image.size
    if w <= 0 or h <= 0 or target_pixels <= 0:
        return w, h
    scale = math.sqrt(target_pixels / (w * h))
    new_w = ensure_multiple(int(round(w * scale)), multiple, mode)
    new_h = ensure_multiple(int(round(h * scale)), multiple, mode)
    return new_w, new_h

# ──────────────────────────────────────────────────────────────────────────────
# Каталог разрешений: загрузка/генерация/фильтрация
# ──────────────────────────────────────────────────────────────────────────────

def load_catalog_from_file(path: str) -> List[Tuple[int, int]]:
    if not os.path.isfile(path):
        return []
    try:
        with open(path, "r", encoding="utf-8", errors="ignore") as f:
            text = f.read()
        return parse_dims(text)
    except Exception:
        return []

_GRID_CACHE: Dict[Tuple[int, int, int], List[Tuple[int, int]]] = {}

def generate_grid_catalog(min_size: int, max_size: int, step: int) -> List[Tuple[int, int]]:
    """Генерация полной решётки кратных 'step' размеров в диапазоне [min_size, max_size], с кэшем."""
    key = (min_size, max_size, step)
    cached = _GRID_CACHE.get(key)
    if cached is not None:
        return cached
    sizes = list(range(min_size, max_size + 1, step))
    out: List[Tuple[int, int]] = []
    for w in sizes:
        for h in sizes:
            out.append((w, h))
    _GRID_CACHE[key] = out
    return out

def orientation_of(w: int, h: int) -> str:
    if w == h:
        return "square"
    return "landscape" if w > h else "portrait"

def aspect_tuple(w: int, h: int) -> Tuple[int, int]:
    g = gcd(w, h)
    return (w // g, h // g)

def _parse_aspect_tokens(aspect_str: str) -> List[Tuple[int, int, bool, float]]:
    """
    Поддерживаем токены:
      - 'A:B'         → точное совпадение
      - '~A:B'        → приблизительное (±0.02 по умолчанию)
      - 'A:B@0.03'    → точный A:B, но с допуском 0.03 (редкий кейс)
      - '~A:B@0.015'  → приблизительный матч с явным допуском
    Возвращает список (ax, ay, approx, tol).
    """
    out: List[Tuple[int, int, bool, float]] = []
    if not aspect_str.strip():
        return out
    tokens = [t.strip() for t in re.split(r"[;,]+", aspect_str) if t.strip()]
    for t in tokens:
        approx = t.startswith("~")
        if approx:
            t = t[1:].strip()
        tol = 0.02 if approx else 0.0
        m = re.match(r"(\d+)\s*:\s*(\d+)(?:\s*@\s*(0\.\d+))?$", t)
        if not m:
            # проигнорируем кривые токены
            continue
        ax, ay = int(m.group(1)), int(m.group(2))
        if m.group(3):
            tol = float(m.group(3))
        out.append((ax, ay, approx or tol > 0.0, tol))
    return out

def filter_catalog(catalog: List[Tuple[int, int]],
                   min_width: int, max_width: int,
                   min_height: int, max_height: int,
                   min_mp: float, max_mp: float,
                   multiple: int,
                   orientation: str,           # any|square|landscape|portrait
                   aspect_str: str,            # "", "16:9, 4:3", "~3:2", "21:9@0.03"
                   limit: int,
                   sort_key: str,             # "MP", "Width", "Height"
                   sort_desc: bool) -> List[Dict]:
    aspect_tokens = _parse_aspect_tokens(aspect_str)

    rows: List[Dict] = []
    for (w, h) in catalog:
        if w < min_width or w > max_width:
            continue
        if h < min_height or h > max_height:
            continue
        if multiple > 1 and (w % multiple != 0 or h % multiple != 0):
            continue
        mp = (w * h) / 1_000_000.0
        if mp < min_mp or mp > max_mp:
            continue
        ori = orientation_of(w, h)
        if orientation != "any" and ori != orientation:
            continue

        # аспект-фильтр
        if aspect_tokens:
            w0, h0 = aspect_tuple(w, h)
            ok = False
            for ax, ay, approx, tol in aspect_tokens:
                if approx:
                    if abs((w / h) - (ax / ay)) <= (tol if tol > 0 else 0.02):
                        ok = True
                        break
                else:
                    if (w0, h0) == (ax, ay):
                        ok = True
                        break
            if not ok:
                continue

        rows.append({
            "Width": w,
            "Height": h,
            "Aspect": f"{aspect_tuple(w,h)[0]}:{aspect_tuple(w,h)[1]}",
            "Orientation": ori,
            "MP": round(mp, 3)
        })

    # сортировка
    key = {
        "MP": lambda r: r["MP"],
        "Width": lambda r: r["Width"],
        "Height": lambda r: r["Height"],
    }.get(sort_key, lambda r: (r["MP"], r["Width"], r["Height"]))
    rows.sort(key=key, reverse=sort_desc)

    return rows[:limit]

def stringify_wh_list(rows: List[Dict]) -> str:
    return "\n".join(f"{r['Width']}x{r['Height']}" for r in rows)

# ──────────────────────────────────────────────────────────────────────────────
# Gradio UI
# ──────────────────────────────────────────────────────────────────────────────

DEFAULT_SHORT_PRESETS = ["512", "768", "896"]
DEFAULT_PIXEL_BASES = ["512", "640", "768", "832", "896", "1024", "1152", "1216", "1344", "1536"]

ROUNDING_CHOICES: List[RoundingMode] = ["round", "floor", "ceil"]

def _parse_int_list(str_list: List[str]) -> List[int]:
    out: List[int] = []
    for s in str_list or []:
        m = re.search(r"\d+", s)
        if m:
            out.append(int(m.group(0)))
    return out

# ——— Вычисления по загруженному изображению ——————————————————————————————

def calculate_all(
    image: Image.Image,
    shorters_raw: List[str],
    bases_raw: List[str],
    use_autopick: bool,
    model_family: ModelFamily,
    diffusion_tile: int,
    vae_tile: int,
    manual_multiple: int,
    rounding_mode: RoundingMode,
):
    if image is None:
        raise gr.Error("Загрузите изображение слева.")

    w, h = image.size

    if use_autopick:
        m_auto, reason = autopick_multiple(model_family, diffusion_tile, vae_tile, w, h)
        rounding_multiple = m_auto
        reason_text = f"Auto multiple = **{m_auto}** — {reason}"
    else:
        rounding_multiple = manual_multiple
        reason_text = f"Manual multiple = **{manual_multiple}**"

    shorters = _parse_int_list(shorters_raw) or _parse_int_list(DEFAULT_SHORT_PRESETS)
    bases = _parse_int_list(bases_raw) or _parse_int_list(DEFAULT_PIXEL_BASES)

    rows = []

    # Короткая сторона
    for s in shorters:
        nw, nh = resize_by_shorter_side(image, s, rounding_multiple, rounding_mode)
        rows.append({
            "Preset": f"short={s}",
            "Width": nw,
            "Height": nh,
            "Aspect": simplified_ratio(nw, nh),
            "Pixels": nw * nh,
            "MP": megapixels(nw, nh),
        })

    # По количеству пикселей (N^2)
    for base in bases:
        target_pixels = base * base
        nw, nh = resize_to_pixel_count(image, target_pixels, rounding_multiple, rounding_mode)
        rows.append({
            "Preset": f"{base}² ({target_pixels:,})",
            "Width": nw,
            "Height": nh,
            "Aspect": simplified_ratio(nw, nh),
            "Pixels": nw * nh,
            "MP": megapixels(nw, nh),
        })

    original_info = f"Original: {w}x{h}  |  Aspect {simplified_ratio(w, h)}  |  {megapixels(w, h)} MP"
    return rows, stringify_wh_list(rows), original_info, reason_text

# ——— Каталог: загрузка/фильтр/экспорт ——————————————————————————————————————

CATALOG_PATH = os.path.join(os.getcwd(), "sdxl_resolutions_by_aspect.txt")  # ищем и в /mnt/data
_LOADED_CATALOG: Optional[List[Tuple[int, int]]] = None

def _load_or_generate(default_step: int) -> Tuple[List[Tuple[int, int]], str]:
    global _LOADED_CATALOG
    if _LOADED_CATALOG is None:
        candidates = [
            CATALOG_PATH,
            os.path.join(os.getcwd(), "data", "sdxl_resolutions_by_aspect.txt"),
            "/mnt/data/sdxl_resolutions_by_aspect.txt",
        ]
        for p in candidates:
            cat = load_catalog_from_file(p)
            if cat:
                _LOADED_CATALOG = cat
                return _LOADED_CATALOG, f"Catalog: loaded from file ({len(cat)} items)."
        _LOADED_CATALOG = []
    return _LOADED_CATALOG, f"Catalog: no file, you can GENERATE grid (step {default_step})."

def catalog_filter_action(
    use_file_catalog: bool,
    generate_min: int,
    generate_max: int,
    generate_step: int,
    filter_min_w: int,
    filter_max_w: int,
    filter_min_h: int,
    filter_max_h: int,
    filter_min_mp: float,
    filter_max_mp: float,
    filter_multiple: int,
    filter_orientation: str,
    filter_aspects: str,
    filter_limit: int,
    sort_key: str,
    sort_desc: bool,
):
    base_list, inf = _load_or_generate(generate_step)
    if use_file_catalog and base_list:
        catalog = base_list
        source = inf.replace("Catalog:", "Source:")
    else:
        # safety
        generate_min = max(128, generate_min)
        generate_max = min(4096, max(generate_min, generate_max))
        generate_step = max(8, generate_step)
        sizes = list(range(generate_min, generate_max + 1, generate_step))
        est = len(sizes) ** 2
        if est > 20000:
            # при слишком больших объёмах — подрежем диапазон по центру
            mid = (generate_min + generate_max) // 2
            half = max(generate_step * 5, (generate_max - generate_min) // 6)
            generate_min = max(128, mid - half)
            generate_max = min(4096, mid + half)
        catalog = generate_grid_catalog(generate_min, generate_max, generate_step)
        source = f"Source: generated grid {generate_min}..{generate_max} step={generate_step} ({len(catalog)} combos)"

    rows = filter_catalog(
        catalog=catalog,
        min_width=filter_min_w, max_width=filter_max_w,
        min_height=filter_min_h, max_height=filter_max_h,
        min_mp=filter_min_mp, max_mp=filter_max_mp,
        multiple=filter_multiple,
        orientation=filter_orientation,
        aspect_str=filter_aspects,
        limit=filter_limit,
        sort_key=sort_key,
        sort_desc=sort_desc,
    )
    return rows, stringify_wh_list(rows), source

def _rows_to_csv_bytes(rows: List[Dict]) -> bytes:
    import csv
    from io import StringIO
    buf = StringIO()
    w = csv.writer(buf)
    w.writerow(["Width", "Height", "Aspect", "Orientation", "MP"])
    for r in rows:
        w.writerow([r["Width"], r["Height"], r["Aspect"], r["Orientation"], r["MP"]])
    return buf.getvalue().encode("utf-8")

def _rows_to_json_bytes(rows: List[Dict]) -> bytes:
    import json
    return json.dumps(rows, ensure_ascii=False, indent=2).encode("utf-8")

def export_rows(rows: List[Dict], kind: str):
    """Возвращает tuple(filename, bytes) для gr.File()"""
    if not rows:
        raise gr.Error("Нет данных для экспорта — сначала примените фильтр.")
    if kind == "csv":
        data = _rows_to_csv_bytes(rows)
        name = "filtered_catalog.csv"
    else:
        data = _rows_to_json_bytes(rows)
        name = "filtered_catalog.json"
    path = os.path.join(os.getcwd(), name)
    with open(path, "wb") as f:
        f.write(data)
    return path

# ──────────────────────────────────────────────────────────────────────────────
# UI
# ──────────────────────────────────────────────────────────────────────────────

def on_ui_tab_called():
    with gr.Blocks() as ui:
        gr.Markdown("### 📐 Image Ratio & Resolution Catalog Pro  \n"
                    "Используйте файл SDXL пресетов, либо генерируйте сетку.  "
                    "Аспекты поддерживают токены `16:9, ~3:2, 21:9@0.03, ~9:16@0.015`.")

        with gr.Row():
            with gr.Column(scale=1):
                image = gr.Image(type="pil", source="upload", label="Изображение")
                original_info = gr.Markdown("")

            with gr.Column(scale=1):
                # Авто/ручная кратность
                use_autopick = gr.Checkbox(value=True, label="Auto multiple (рекомендуется)")
                model_family = gr.Radio(
                    choices=["SDXL (native~1024)", "SD 1.x (native~512)", "SD 2.x (native~768)"],
                    value=DEFAULT_MODEL_FAMILY,
                    label="Model family (ручной выбор)"
                )
                with gr.Row():
                    diffusion_tile = gr.Number(value=0, precision=0, label="Diffusion tile (px, 0=нет)")
                    vae_tile = gr.Number(value=0, precision=0, label="VAE tile (px, 0=нет)")
                manual_multiple = gr.Slider(minimum=8, maximum=256, step=8, value=64, label="Manual multiple")
                rounding_mode = gr.Radio(choices=["round", "floor", "ceil"], value="round", label="Режим округления")
                auto_reason = gr.Markdown("")

        with gr.Row():
            with gr.Column():
                shorters = gr.CheckboxGroup(
                    choices=DEFAULT_SHORT_PRESETS,
                    value=DEFAULT_SHORT_PRESETS,
                    label="Цели по короткой стороне (px)",
                )
                pixel_bases = gr.CheckboxGroup(
                    choices=DEFAULT_PIXEL_BASES,
                    value=DEFAULT_PIXEL_BASES,
                    label="Цели по общему числу пикселей (берётся N²)",
                )
                run_btn = gr.Button("Рассчитать", variant="primary")

            with gr.Column():
                results = gr.Dataframe(
                    headers=["Preset", "Width", "Height", "Aspect", "Pixels", "MP"],
                    datatype=["str", "number", "number", "str", "number", "number"],
                    label="Результаты",
                    interactive=False,
                    wrap=True,
                    overflow_row_behaviour="paginate",
                )
                copy_box = gr.Textbox(label="Список WxH (для копирования)", lines=8, interactive=False)

        # Каталог разрешений (импорт/генерация/фильтры)
        gr.Markdown("---")
        gr.Markdown("#### Каталог разрешений (импорт из файла или генерация сетки)")

        with gr.Row():
            use_file_catalog = gr.Checkbox(value=True, label="Использовать файл, если найден (SDXL пресеты)")
            generate_min = gr.Slider(minimum=128, maximum=4096, step=8, value=512, label="Generate: min")
            generate_max = gr.Slider(minimum=128, maximum=4096, step=8, value=2048, label="Generate: max")
            generate_step = gr.Slider(minimum=8, maximum=256, step=8, value=64, label="Generate: step")

        with gr.Row():
            filter_min_w = gr.Slider(minimum=128, maximum=4096, step=8, value=512, label="min Width")
            filter_max_w = gr.Slider(minimum=128, maximum=4096, step=8, value=2048, label="max Width")
            filter_min_h = gr.Slider(minimum=128, maximum=4096, step=8, value=512, label="min Height")
            filter_max_h = gr.Slider(minimum=128, maximum=4096, step=8, value=2048, label="max Height")

        with gr.Row():
            filter_min_mp = gr.Number(value=0.0, label="min MP")
            filter_max_mp = gr.Number(value=8.0, label="max MP")
            filter_multiple = gr.Slider(minimum=8, maximum=256, step=8, value=64, label="Требуемая кратность")
            filter_orientation = gr.Dropdown(choices=["any", "square", "landscape", "portrait"], value="any", label="Ориентация")
            filter_aspects = gr.Textbox(value="", label="Аспекты (пример: '16:9, ~3:2, 21:9@0.03')")

        with gr.Row():
            filter_limit = gr.Slider(minimum=10, maximum=5000, step=10, value=500, label="Лимит результатов")
            sort_key = gr.Dropdown(choices=["MP", "Width", "Height"], value="MP", label="Сортировать по")
            sort_desc = gr.Checkbox(value=True, label="По убыванию")
            apply_filters = gr.Button("Применить фильтры", variant="primary")

        catalog_df = gr.Dataframe(
            headers=["Width", "Height", "Aspect", "Orientation", "MP"],
            datatype=["number", "number", "str", "str", "number"],
            label="Каталог (фильтрованный)",
            interactive=False,
            wrap=True,
            overflow_row_behaviour="paginate",
        )
        catalog_copy = gr.Textbox(label="Список WxH (для копирования)", lines=12, interactive=False)
        catalog_source = gr.Markdown("")

        with gr.Row():
            export_csv_btn = gr.Button("Скачать CSV")
            export_json_btn = gr.Button("Скачать JSON")
            exported_file = gr.File(label="Экспорт", visible=True)

        # — Автозапуски/связи —
        def _toggle_manual(m_use_auto: bool):
            return gr.update(interactive=not m_use_auto)

        use_autopick.change(fn=_toggle_manual, inputs=[use_autopick], outputs=[manual_multiple])

        for trigger in (image.upload, run_btn.click):
            trigger(
                fn=calculate_all,
                inputs=[image, shorters, pixel_bases, use_autopick, model_family, diffusion_tile, vae_tile, manual_multiple, rounding_mode],
                outputs=[results, copy_box, original_info, auto_reason],
                show_progress=False,
            )

        apply_filters.click(
            fn=catalog_filter_action,
            inputs=[
                use_file_catalog, generate_min, generate_max, generate_step,
                filter_min_w, filter_max_w, filter_min_h, filter_max_h,
                filter_min_mp, filter_max_mp, filter_multiple,
                filter_orientation, filter_aspects, filter_limit,
                sort_key, sort_desc
            ],
            outputs=[catalog_df, catalog_copy, catalog_source],
            show_progress=False
        )

        export_csv_btn.click(lambda rows: export_rows(rows, "csv"), inputs=[catalog_df], outputs=[exported_file])
        export_json_btn.click(lambda rows: export_rows(rows, "json"), inputs=[catalog_df], outputs=[exported_file])

    return (ui, "calculator+", "calculator_plus_interface"),


# ──────────────────────────────────────────────────────────────────────────────
# Регистрация в A1111 или автономный запуск
# ──────────────────────────────────────────────────────────────────────────────
try:
    from modules import script_callbacks  # type: ignore
    script_callbacks.on_ui_tabs(on_ui_tab_called)
except (ImportError, ModuleNotFoundError):
    if __name__ == "__main__":
        interface, _, _ = on_ui_tab_called()[0]
        interface.launch()