import os
import time
import re
import torch
from diffusers import StableDiffusionPipeline


def slugify(text: str, max_len: int = 40) -> str:
    text = text.lower()
    text = re.sub(r"[^a-z0-9]+", "_", text).strip("_")
    return text[:max_len] if len(text) > max_len else text


def build_pipe(
    base_model_id: str = "runwayml/stable-diffusion-v1-5",
    dtype: torch.dtype = torch.float16,
    device: str = "cuda",
    use_xformers: bool = False,  # T4에서 이슈 있으면 False 추천
):
    pipe = StableDiffusionPipeline.from_pretrained(
        base_model_id,
        torch_dtype=dtype,
        safety_checker=None,  # 과제/로컬 데모면 보통 끔
    ).to(device)

    if use_xformers:
        try:
            pipe.enable_xformers_memory_efficient_attention()
            print("xformers enabled ✅")
        except Exception as e:
            print(f"xformers not available (skip): {e}")

    return pipe


def try_set_lora_scale(pipe, lora_scale: float) -> bool:
    """
    diffusers 버전에 따라 set_adapters 지원 여부가 달라서 방어적으로 처리.
    """
    try:
        pipe.set_adapters(["default"], adapter_weights=[float(lora_scale)])
        return True
    except Exception:
        return False


def main():
    # ===== 사용자 설정 =====
    trigger_word = "shinkai_makoto_style"
    base_model_id = "runwayml/stable-diffusion-v1-5"

    # 최종 LoRA 폴더 (필요 시 checkpoint-1500으로 바꿔도 됨)
    lora_model_path = "./shinkai_lora_output"

    # 결과 개선 핵심: LoRA 강도 낮추기
    lora_scales = [0.65, 0.75, 0.85]

    # guidance는 낮춘 쪽이 형태 안정에 유리한 경우가 많음
    guidance_scales = [6.5, 7.0]

    # 디테일을 위해 steps 살짝 늘림
    num_inference_steps = 45

    # 재현성
    seed = 1337

    # T4면 False 추천(안정), L4/A100이면 True도 OK
    use_xformers = False

    # ===== 풍경 프롬프트 5개 세트 =====
    prompts = [
        (
            "rainy_neon_city",
            f"{trigger_word}, a rainy city street at dusk, neon lights reflecting on wet asphalt, "
            "small girl silhouette holding a transparent umbrella, wide shot, cinematic composition, "
            "soft lighting, detailed background, anime illustration",
        ),
        (
            "sunset_cityscape",
            f"{trigger_word}, a vast cityscape under a dramatic sunset sky, glowing clouds, "
            "warm orange and pink tones, tiny human figure on a rooftop looking at the sky, "
            "wide angle, atmospheric perspective, anime background art",
        ),
        (
            "starry_town_night",
            f"{trigger_word}, a quiet town under a starry night sky, bright stars, soft clouds drifting, "
            "cinematic lighting, calm mood, highly detailed sky, anime background illustration",
        ),
        (
            "railway_rain_perspective",
            f"{trigger_word}, a railway stretching into the distance during rainfall, wet rails reflecting city lights, "
            "empty platform, deep perspective, moody atmosphere, anime cinematic background",
        ),
        (
            "rural_after_rain",
            f"{trigger_word}, a quiet rural town with fields and houses, dramatic cloudy sky after rain, "
            "soft sunlight breaking through clouds, wide shot, peaceful mood, anime background art",
        ),
    ]

    negative_prompt = (
        "lowres, bad anatomy, bad hands, extra fingers, missing fingers, "
        "deformed face, cross-eye, distorted, ugly, blurry, noisy, watermark, text"
    )

    # ===== 체크 =====
    if not os.path.exists(lora_model_path):
        raise FileNotFoundError(f"LoRA 경로가 없습니다: {lora_model_path}")

    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.float16 if device == "cuda" else torch.float32

    # ===== 파이프라인 로드(1번만) =====
    print("Loading base model:", base_model_id)
    pipe = build_pipe(
        base_model_id=base_model_id,
        dtype=dtype,
        device=device,
        use_xformers=use_xformers,
    )

    # LoRA 로드(1번만)
    print("Loading LoRA from:", lora_model_path)
    pipe.load_lora_weights(lora_model_path)

    # ===== 출력 폴더 =====
    out_dir = "inference_outputs_landscape"
    os.makedirs(out_dir, exist_ok=True)

    ts = time.strftime("%Y%m%d_%H%M%S")
    print(f"Output dir: {out_dir}")
    print(f"Steps={num_inference_steps}, seed={seed}, lora_scales={lora_scales}, guidance_scales={guidance_scales}")
    print("Negative prompt:", negative_prompt)

    # ===== 생성 =====
    total = len(prompts) * len(lora_scales) * len(guidance_scales)
    done = 0

    for lora_scale in lora_scales:
        supported = try_set_lora_scale(pipe, lora_scale)
        if supported:
            print(f"\n✅ LoRA scale set ({lora_scale})")
        else:
            print(f"\n⚠️ set_adapters not supported; using default LoRA strength (requested {lora_scale})")

        for gs in guidance_scales:
            for name, prompt in prompts:
                done += 1
                gen = torch.Generator(device=device).manual_seed(seed)

                fname = (
                    f"{out_dir}/"
                    f"{ts}_p{name}_lora{lora_scale}_gs{gs}_steps{num_inference_steps}_seed{seed}.png"
                )

                print(f"[{done}/{total}] -> {fname}")

                if device == "cuda":
                    with torch.autocast("cuda", dtype=dtype):
                        image = pipe(
                            prompt=prompt,
                            negative_prompt=negative_prompt,
                            num_inference_steps=num_inference_steps,
                            guidance_scale=float(gs),
                            generator=gen,
                        ).images[0]
                else:
                    with torch.no_grad():
                        image = pipe(
                            prompt=prompt,
                            negative_prompt=negative_prompt,
                            num_inference_steps=num_inference_steps,
                            guidance_scale=float(gs),
                        ).images[0]

                image.save(fname)

    print("\n✅ All done!")
    print(f"Check: {out_dir}")


if __name__ == "__main__":
    main()