"""
Merged grid image splitter.

This node is designed for AI-generated contact sheets such as 3x3 or 2x3
grids. Auto mode scores common layouts by looking for strong seams or flat
separator bands near the expected grid lines, then crops each cell.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import List, Sequence, Tuple

import numpy as np
import torch
from PIL import Image

from ..utils.image_utils import pil_to_tensor, tensor_to_pil


_AUTO_LAYOUTS: Sequence[Tuple[int, int]] = (
    (3, 3),
    (2, 3),
    (3, 2),
    (2, 2),
    (1, 2),
    (2, 1),
    (1, 3),
    (3, 1),
    (4, 4),
    (3, 4),
    (4, 3),
)

_LAYOUTS = [
    "auto",
    "1x2",
    "2x1",
    "1x3",
    "3x1",
    "2x2",
    "2x3",
    "3x2",
    "3x3",
    "3x4",
    "4x3",
    "4x4",
    "custom",
]


@dataclass(frozen=True)
class _AxisCut:
    seam: int
    span_start: int
    span_end: int
    score: float


@dataclass(frozen=True)
class _AxisPlan:
    intervals: List[Tuple[int, int]]
    cuts: List[_AxisCut]
    score: float


def _to_float_array(image: Image.Image) -> np.ndarray:
    if image.mode != "RGB":
        image = image.convert("RGB")
    return np.asarray(image).astype(np.float32) / 255.0


def _axis_texture(arr: np.ndarray, axis: str) -> np.ndarray:
    if axis == "x":
        profile = arr.std(axis=(0, 2))
    else:
        profile = arr.std(axis=(1, 2))
    high = np.percentile(profile, 95) + 1e-6
    return np.clip(profile / high, 0.0, 1.0)


def _axis_edge(arr: np.ndarray, axis: str) -> np.ndarray:
    if axis == "x":
        diff = np.abs(np.diff(arr, axis=1)).mean(axis=(0, 2))
        length = arr.shape[1]
    else:
        diff = np.abs(np.diff(arr, axis=0)).mean(axis=(1, 2))
        length = arr.shape[0]

    padded = np.zeros(length, dtype=np.float32)
    if diff.size:
        padded[1:] = diff
    high = np.percentile(padded, 95) + 1e-6
    return np.clip(padded / high, 0.0, 1.5)


def _smooth(profile: np.ndarray, radius: int = 2) -> np.ndarray:
    if radius <= 0 or profile.size < radius * 2 + 1:
        return profile
    kernel = np.ones(radius * 2 + 1, dtype=np.float32) / float(radius * 2 + 1)
    return np.convolve(profile, kernel, mode="same")


def _separator_span(
    texture: np.ndarray,
    seam: int,
    search_px: int,
    min_separator_px: int,
) -> Tuple[int, int]:
    length = texture.size
    if length <= 1:
        return 0, length

    limit = max(1, min(search_px, length // 8))
    threshold = max(0.08, min(0.28, float(np.percentile(texture, 12)) * 1.8))

    left = seam
    while left > 0 and seam - left < limit and texture[left - 1] <= threshold:
        left -= 1

    right = seam
    while right < length and right - seam < limit and texture[right] <= threshold:
        right += 1

    if right - left >= max(1, min_separator_px):
        return left, right

    return seam, seam


def _edge_trim(texture: np.ndarray, search_px: int, min_cell: int) -> Tuple[int, int]:
    length = texture.size
    if length <= 2:
        return 0, length

    max_trim = max(0, min(search_px * 2, min_cell // 3, length // 6))
    if max_trim <= 0:
        return 0, length

    threshold = max(0.08, min(0.24, float(np.percentile(texture, 12)) * 1.6))

    start = 0
    while start < max_trim and texture[start] <= threshold:
        start += 1

    end = length
    while length - end < max_trim and end > start + min_cell and texture[end - 1] <= threshold:
        end -= 1

    return start, end


def _axis_plan(
    arr: np.ndarray,
    cells: int,
    axis: str,
    search_px: int,
    crop_separators: bool,
    trim_outer: bool,
    min_separator_px: int,
) -> _AxisPlan:
    length = arr.shape[1] if axis == "x" else arr.shape[0]
    if cells <= 1:
        return _AxisPlan(intervals=[(0, length)], cuts=[], score=0.0)

    raw_texture = _axis_texture(arr, axis)
    raw_edge = _axis_edge(arr, axis)
    texture = _smooth(raw_texture, radius=2)
    edge = _smooth(raw_edge, radius=1)
    evidence = np.maximum(edge, (1.0 - texture) * 0.75)
    exact_evidence = np.maximum(raw_edge, (1.0 - raw_texture) * 0.75)

    cuts: List[_AxisCut] = []
    scores: List[float] = []
    for idx in range(1, cells):
        expected = round(length * idx / cells)
        start = max(1, expected - search_px)
        end = min(length - 1, expected + search_px)
        if start >= end:
            seam = expected
            score = 0.0
        else:
            window = evidence[start:end + 1]
            offset = int(window.argmax())
            coarse = start + offset
            fine_start = max(start, coarse - 2)
            fine_end = min(end, coarse + 2)
            fine_window = exact_evidence[fine_start:fine_end + 1]
            seam = fine_start + int(fine_window.argmax())
            score = float(window[offset])

        span_start, span_end = _separator_span(
            raw_texture,
            seam,
            search_px=search_px,
            min_separator_px=min_separator_px,
        )
        cuts.append(_AxisCut(seam=seam, span_start=span_start, span_end=span_end, score=score))
        scores.append(score)

    min_cell = max(1, length // cells)
    outer_start, outer_end = _edge_trim(raw_texture, search_px, min_cell) if trim_outer else (0, length)

    intervals: List[Tuple[int, int]] = []
    cursor = outer_start
    for cut in cuts:
        split_start = cut.span_start if crop_separators else cut.seam
        split_end = cut.span_end if crop_separators else cut.seam
        intervals.append((cursor, split_start))
        cursor = split_end
    intervals.append((cursor, outer_end))

    cleaned: List[Tuple[int, int]] = []
    for start, end in intervals:
        start = max(0, min(length - 1, int(start)))
        end = max(start + 1, min(length, int(end)))
        cleaned.append((start, end))

    return _AxisPlan(
        intervals=cleaned,
        cuts=cuts,
        score=float(np.mean(scores)) if scores else 0.0,
    )


def _parse_layout(layout: str, custom_rows: int, custom_cols: int) -> Tuple[int, int]:
    if layout == "custom":
        return max(1, int(custom_rows)), max(1, int(custom_cols))
    rows_text, cols_text = layout.split("x", 1)
    return int(rows_text), int(cols_text)


def _fallback_layout(width: int, height: int) -> Tuple[int, int]:
    aspect = width / max(1, height)
    if 0.82 <= aspect <= 1.22:
        return 3, 3
    if aspect > 1.22:
        return 2, 3
    return 3, 2


def _choose_auto_layout(
    arr: np.ndarray,
    search_px: int,
    crop_separators: bool,
    trim_outer: bool,
    min_separator_px: int,
) -> Tuple[int, int, _AxisPlan, _AxisPlan, float, bool]:
    height, width = arr.shape[:2]
    best = None

    for rows, cols in _AUTO_LAYOUTS:
        x_plan = _axis_plan(arr, cols, "x", search_px, crop_separators, trim_outer, min_separator_px)
        y_plan = _axis_plan(arr, rows, "y", search_px, crop_separators, trim_outer, min_separator_px)
        score = (x_plan.score + y_plan.score) / 2.0

        # Prefer common 3x3 / 2x3 / 3x2 layouts when the image gives weak signals.
        if (rows, cols) in ((3, 3), (2, 3), (3, 2)):
            score += 0.025

        if best is None or score > best[0]:
            best = (score, rows, cols, x_plan, y_plan)

    assert best is not None
    score, rows, cols, x_plan, y_plan = best
    confident = score >= 0.22

    if confident:
        return rows, cols, x_plan, y_plan, score, True

    rows, cols = _fallback_layout(width, height)
    x_plan = _axis_plan(arr, cols, "x", search_px, crop_separators, trim_outer, min_separator_px)
    y_plan = _axis_plan(arr, rows, "y", search_px, crop_separators, trim_outer, min_separator_px)
    return rows, cols, x_plan, y_plan, score, False


def _normalize_sizes(crops: List[Image.Image]) -> List[Image.Image]:
    min_w = min(crop.width for crop in crops)
    min_h = min(crop.height for crop in crops)
    normalized = []
    for crop in crops:
        left = max(0, (crop.width - min_w) // 2)
        top = max(0, (crop.height - min_h) // 2)
        normalized.append(crop.crop((left, top, left + min_w, top + min_h)))
    return normalized


def _split_one(
    image: Image.Image,
    layout: str,
    custom_rows: int,
    custom_cols: int,
    search_px: int,
    crop_separators: bool,
    trim_outer: bool,
    min_separator_px: int,
) -> Tuple[List[Image.Image], str]:
    arr = _to_float_array(image)

    if layout == "auto":
        rows, cols, x_plan, y_plan, confidence, confident = _choose_auto_layout(
            arr,
            search_px=search_px,
            crop_separators=crop_separators,
            trim_outer=trim_outer,
            min_separator_px=min_separator_px,
        )
        mode_note = "auto" if confident else "auto-low-confidence-fallback"
    else:
        rows, cols = _parse_layout(layout, custom_rows, custom_cols)
        x_plan = _axis_plan(arr, cols, "x", search_px, crop_separators, trim_outer, min_separator_px)
        y_plan = _axis_plan(arr, rows, "y", search_px, crop_separators, trim_outer, min_separator_px)
        confidence = (x_plan.score + y_plan.score) / 2.0
        mode_note = "manual"

    crops: List[Image.Image] = []
    for y0, y1 in y_plan.intervals:
        for x0, x1 in x_plan.intervals:
            crops.append(image.crop((x0, y0, x1, y1)))

    crops = _normalize_sizes(crops)
    info = (
        f"{mode_note}: {rows}x{cols}, cells={len(crops)}, "
        f"confidence={confidence:.3f}, "
        f"x={x_plan.intervals}, y={y_plan.intervals}"
    )
    return crops, info


class O1keyGridSplitter:
    """Split AI-generated grid/contact-sheet images into individual cells."""

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "图像": ("IMAGE",),
                "布局": (_LAYOUTS, {"default": "auto"}),
                "自定义行数": ("INT", {"default": 3, "min": 1, "max": 12, "step": 1}),
                "自定义列数": ("INT", {"default": 3, "min": 1, "max": 12, "step": 1}),
                "搜索范围px": ("INT", {"default": 32, "min": 0, "max": 256, "step": 1}),
                "裁掉分隔线": ("BOOLEAN", {"default": True}),
                "裁掉外边距": ("BOOLEAN", {"default": True}),
                "最小分隔线px": ("INT", {"default": 2, "min": 0, "max": 64, "step": 1}),
                "最大输出张数": ("INT", {"default": 16, "min": 1, "max": 144, "step": 1}),
            }
        }

    RETURN_TYPES = ("IMAGE", "STRING")
    RETURN_NAMES = ("切割图像", "检测信息")
    FUNCTION = "split_grid"
    CATEGORY = "o1key/image"
    DESCRIPTION = (
        "智能切割 AI 生成的九宫格、六宫格等合并图。"
        "自动模式会检测常见布局；没有明显分隔线时建议手动选择布局。"
    )

    def split_grid(
        self,
        图像: torch.Tensor,
        布局: str = "auto",
        自定义行数: int = 3,
        自定义列数: int = 3,
        搜索范围px: int = 32,
        裁掉分隔线: bool = True,
        裁掉外边距: bool = True,
        最小分隔线px: int = 2,
        最大输出张数: int = 16,
    ):
        source_images = tensor_to_pil(图像)
        all_crops: List[Image.Image] = []
        info_lines: List[str] = []

        for batch_index, image in enumerate(source_images, start=1):
            crops, info = _split_one(
                image=image,
                layout=布局,
                custom_rows=自定义行数,
                custom_cols=自定义列数,
                search_px=搜索范围px,
                crop_separators=裁掉分隔线,
                trim_outer=裁掉外边距,
                min_separator_px=最小分隔线px,
            )
            if len(crops) > 最大输出张数:
                raise ValueError(
                    f"合并图切割：检测到 {len(crops)} 张，超过最大输出张数 {最大输出张数}。"
                    "请调大最大输出张数，或检查布局设置。"
                )
            all_crops.extend(crops)
            info_lines.append(f"batch {batch_index}: {info}")

        if not all_crops:
            raise ValueError("合并图切割：没有生成任何切片。")

        all_crops = _normalize_sizes(all_crops)
        print("[o1key 合并图切割] " + " | ".join(info_lines))
        return (pil_to_tensor(all_crops), "\n".join(info_lines))