import os, torch, shutil
from diffusers import StableDiffusionXLPipeline
from huggingface_hub import HfApi, hf_hub_download, create_repo

# ====================== CONFIG ======================
BASE_MODEL = "stabilityai/stable-diffusion-xl-base-1.0"
LORA_REPO  = "CoolKrishh/Comic-SDXL-LoRA"
LORA_FILE  = "Comic-SDXL.safetensors"
DEST_REPO  = "CoolKrishh/mythic-sdxl"
OUTPUT_CKPT = "mythic-sdxl.safetensors"
OUTPUT_DIR  = "/app/out"
# ===================================================

TOKEN = os.getenv("HF_TOKEN")
if not TOKEN:
    raise ValueError("❌ HF_TOKEN not set!")

api = HfApi(token=TOKEN)
create_repo(repo_id=DEST_REPO, exist_ok=True, token=TOKEN)

print("⬇ Downloading LoRA...")
lora_path = hf_hub_download(LORA_REPO, LORA_FILE, token=TOKEN)

print("🚀 Loading SDXL base...")
pipe = StableDiffusionXLPipeline.from_pretrained(
    BASE_MODEL, torch_dtype=torch.float16
)

print("🧬 Merging LoRA...")
pipe.load_lora_weights(lora_path, adapter_name="default")
pipe.fuse_lora()

# Ensure output folder exists
os.makedirs(OUTPUT_DIR, exist_ok=True)

print("💾 Saving merged pipeline...")
pipe.save_pretrained(OUTPUT_DIR, safe_serialization=True)

# ---- FIND REAL SAVED SAFETENSORS ----
print("🔍 Scanning saved files...")
found = None
for f in os.listdir(OUTPUT_DIR):
    if f.endswith(".safetensors"):
        found = f
        break

if not found:
    raise FileNotFoundError("❌ No .safetensors file was created by diffusers!")

src = os.path.join(OUTPUT_DIR, found)
dst = os.path.join(OUTPUT_DIR, OUTPUT_CKPT)

print(f"📌 Renaming:\n  {src}\n  → {dst}")
shutil.move(src, dst)

# ---- VERIFY FINAL FILE EXISTS ----
if not os.path.exists(dst):
    raise FileNotFoundError(f"❌ Still missing after rename: {dst}")

print(f"✅ Final model file ready at:\n{dst}")

# ---- UPLOAD ----
print("☁ Uploading to Hugging Face...")
api.upload_file(
    path_or_fileobj=dst,
    path_in_repo=OUTPUT_CKPT,
    repo_id=DEST_REPO,
    token=TOKEN
)

print("\n🎉 SUCCESS! Your model is now uploaded & inference-ready.")
print(f"🔗 https://huggingface.co/{DEST_REPO}")