#!/usr/bin/env python3
import argparse
import os
import subprocess
import torch
from PIL import Image, ImageDraw
import shutil
import tempfile

from diffusers import (
    StableDiffusionPipeline,
    DPMSolverMultistepScheduler,
    StableDiffusionInpaintPipeline
)

def shift_image(img: Image.Image, shift: int) -> Image.Image:
    w, h = img.size
    out = Image.new("RGB", (w, h))
    out.paste(img.crop((shift, 0, w, h)), (0, 0))
    out.paste(img.crop((0, 0, shift, h)), (w - shift, 0))
    return out

def create_mask(width: int, height: int, mask_w: int) -> Image.Image:
    mask = Image.new("L", (width, height), 0)
    draw = ImageDraw.Draw(mask)
    left = (width - mask_w) // 2
    draw.rectangle([left, 0, left + mask_w, height], fill=255)
    return mask

def unshift_image(img: Image.Image, shift: int) -> Image.Image:
    w, h = img.size
    out = Image.new("RGB", (w, h))
    out.paste(img.crop((w - shift, 0, w, h)), (0, 0))
    out.paste(img.crop((0, 0, w - shift, h)), (shift, 0))
    return out

def generate_equirect(prompt: str, output: str, work_dir: str | None = None) -> str | None:
    # Output-Ordner (bleibt wie gehabt)
    output_abs = os.path.abspath(output)
    work_dir = work_dir or os.path.dirname(os.path.abspath(__file__))

    # Zwischenschritte landen im eigenem temp-Ordner:
    with tempfile.TemporaryDirectory(dir=work_dir) as tempdir:
        print(f"→ Using tempdir: {tempdir}")

        model_path = "/Volumes/SD/ML-Models/diffusers/hdri-panorama-v1-diffusers"
        topaz_cli = "/Applications/Topaz Photo AI.app/Contents/MacOS/Topaz Photo AI"
        steps = 20
        scale = 7.0
        width, height = 1024, 512

        if torch.backends.mps.is_available():
            device = "mps"
        elif torch.cuda.is_available():
            device = "cuda"
        else:
            device = "cpu"

        # 1) Generate base HDRI
        gen_pipe = StableDiffusionPipeline.from_pretrained(
            model_path,
            torch_dtype=torch.float32
        ).to(device)
        gen_pipe.scheduler = DPMSolverMultistepScheduler.from_config(gen_pipe.scheduler.config)
        gen_pipe.enable_attention_slicing()

        print("→ Generating equirectangular HDRI…")
        image = gen_pipe(
            prompt=prompt,
            num_inference_steps=steps,
            guidance_scale=scale-1.5,
            width=width,
            height=height
        ).images[0]

        gen_path = os.path.join(tempdir, f"base_{width}x{height}.png")
        image.save(gen_path)
        print(f"→ Saved initial image to {gen_path}")

        # 2) Make it seamless
        shift_amt = width // 2
        mask_w = width // 8

        shifted = shift_image(image, shift_amt)
        mask = create_mask(width, height, mask_w)

        inpaint_pipe = StableDiffusionInpaintPipeline.from_pretrained(
            "Lykon/dreamshaper-8-inpainting",
            torch_dtype=torch.float32
        ).to(device)
        inpaint_pipe.enable_attention_slicing()

        print("→ Inpainting seam for seamless tiling…")
        inpainted = inpaint_pipe(
            prompt=prompt,
            image=shifted,
            mask_image=mask,
            num_inference_steps=steps,
            guidance_scale=scale,
            width=width,
            height=height
        ).images[0]

        seamless_path = os.path.join(tempdir, os.path.basename(output_abs))
        inpainted = unshift_image(inpainted, shift_amt)
        inpainted.save(seamless_path)
        print(f"→ Crafted seamless image: {seamless_path}")

        # 3) Upscale with Topaz Photo AI CLI
        print("→ Upscaling with Topaz Photo AI CLI…")
        result = subprocess.run(
            [topaz_cli, "--cli", seamless_path, "-o", tempdir],
            check=True
        )

        # Finde das letzte erstellte PNG im tempdir (das ist das hochskalierte!)
        # Topaz kann einen Suffix anhängen, falls der Name schon existiert.
        upscaled_files = sorted(
            [os.path.join(tempdir, f) for f in os.listdir(tempdir) if f.lower().endswith(".png")],
            key=os.path.getmtime,
            reverse=True
        )
        if not upscaled_files:
            print("→ No PNG output found in tempdir after Topaz run!")
            return None

        upscaled = upscaled_files[0]
        shutil.move(upscaled, output_abs)
        print(f"→ Upscaled image moved to {output_abs}")
        return output_abs


def main():
    parser = argparse.ArgumentParser(
        description="Generate an equirectangular HDRI, make it seamless, and upscale it with Topaz Photo AI CLI."
    )
    parser.add_argument("--prompt", required=True,
                        help="Text prompt for generation and inpainting")
    parser.add_argument("--output", required=True,
                        help="Filename for the final upscaled image (e.g. seamless.png)")
    parser.add_argument("--work-dir", default=os.path.dirname(os.path.abspath(__file__)),
                        help="Working directory for intermediates and final outputs")
    args = parser.parse_args()
    generate_equirect(args.prompt, args.output, args.work_dir)

if __name__ == "__main__":
    main()