from __future__ import annotations

import shutil
from pathlib import Path

from PIL import Image, ImageEnhance, ImageFilter, ImageOps


TRAIN_SIZES = (1024, 768, 512)


def reset_dir(path: Path) -> None:
    if path.exists():
        shutil.rmtree(path)
    path.mkdir(parents=True, exist_ok=True)


def normalize_image(image_path: str | Path) -> Image.Image:
    image = Image.open(image_path)
    image = ImageOps.exif_transpose(image)
    if image.mode in {"RGBA", "LA"}:
        background = Image.new("RGB", image.size, (255, 255, 255))
        background.paste(image.convert("RGBA"), mask=image.convert("RGBA").getchannel("A"))
        image = background
    return image.convert("RGB")


def fit_to_square(image: Image.Image, size: int) -> Image.Image:
    image.thumbnail((size, size), Image.Resampling.LANCZOS)
    canvas = Image.new("RGB", (size, size), (255, 255, 255))
    left = (size - image.width) // 2
    top = (size - image.height) // 2
    canvas.paste(image, (left, top))
    return canvas


def align_to_size(image: Image.Image, target_size: tuple[int, int]) -> Image.Image:
    """Fit an image onto the same canvas size as the target image.

    The comparison image should already be composition-matched. This step fixes
    small aspect-ratio drift from external image generators before both images
    are padded to the square training buckets.
    """

    target_width, target_height = target_size
    if image.size == target_size:
        return image.convert("RGB")

    image_aspect = image.width / max(image.height, 1)
    target_aspect = target_width / max(target_height, 1)
    if abs(image_aspect - target_aspect) <= 0.02:
        resized = image.resize(target_size, Image.Resampling.LANCZOS)
        return resized.convert("RGB")

    # Use cover+center-crop when the ratio is meaningfully different. This
    # preserves the target canvas and avoids adding new border regions to only
    # one side of the LoRA subtraction pair.
    return ImageOps.fit(image, target_size, method=Image.Resampling.LANCZOS, centering=(0.5, 0.5)).convert("RGB")


def grayscale_as_rgb(image: Image.Image) -> Image.Image:
    return ImageOps.grayscale(image).convert("RGB")


def prepare_dataset_from_image(source: Image.Image, caption: str, dataset_root: Path) -> Path:
    image_dir = dataset_root / "4000_style"
    reset_dir(image_dir)

    for size in TRAIN_SIZES:
        prepared = fit_to_square(source.copy(), size)
        stem = str(size)
        prepared.save(image_dir / f"{stem}.png")
        (image_dir / f"{stem}.txt").write_text(caption.strip(), encoding="utf-8")

    return dataset_root


def prepare_dataset(image_path: str | Path, caption: str, dataset_root: Path) -> Path:
    source = normalize_image(image_path)
    return prepare_dataset_from_image(source, caption, dataset_root)


def prepare_matched_datasets(
    target_image_path: str | Path,
    pair_image_path: str | Path,
    target_caption: str,
    pair_caption: str,
    run_dir: Path,
    grayscale_pair: bool = False,
) -> tuple[Path, Path]:
    target = normalize_image(target_image_path)
    pair = normalize_image(pair_image_path)
    pair = align_to_size(pair, target.size)
    if grayscale_pair:
        pair = grayscale_as_rgb(pair)
    target_data = prepare_dataset_from_image(target, target_caption, run_dir / "target_data")
    pair_data = prepare_dataset_from_image(pair, pair_caption, run_dir / "reference_data")
    return target_data, pair_data


def create_auto_reference_image(image_path: str | Path, output_path: Path) -> Path:
    """Create a content-preserving, style-muted comparison image.

    This is intentionally simple and local: it does not use any pretrained LoRA.
    The goal is to keep composition while weakening color, texture, and line style.
    """

    image = normalize_image(image_path)
    neutral = ImageOps.autocontrast(image, cutoff=1)
    neutral = ImageEnhance.Color(neutral).enhance(0.45)
    neutral = ImageEnhance.Contrast(neutral).enhance(0.85)
    neutral = neutral.filter(ImageFilter.SMOOTH_MORE)
    edges = image.convert("L").filter(ImageFilter.FIND_EDGES).filter(ImageFilter.SMOOTH)
    edges = ImageOps.invert(edges)
    edges = ImageEnhance.Contrast(edges).enhance(1.6).convert("RGB")
    neutral = Image.blend(neutral, edges, 0.18)

    output_path.parent.mkdir(parents=True, exist_ok=True)
    neutral.save(output_path)
    return output_path
