import os
import traceback
from datetime import datetime
import torch, gc
from PIL import Image
import gradio as gr

from inference import generate_with_lora
from background_edit import run_background_removal_and_inpaint

# ───────────────────── Helpers ─────────────────────
def _print_trace():
    traceback.print_exc()

def unload_models():
    torch.cuda.empty_cache()
    gc.collect()

def safe_generate_and_inpaint(
    image,
    prompt_1, neg_1, strength_1, guidance_1,
    prompt_2, neg_2, guidance_2
):
    try:
        if image is None:
            raise gr.Error("Please upload an image first.")

        # Step 1: Headshot Refinement
        print("[INFO] Step 1: Refining headshot...", flush=True)
        refined = generate_with_lora(
            image=image,
            prompt=prompt_1,
            negative_prompt=neg_1,
            strength=strength_1,
            guidance_scale=guidance_1,
        )

        # Save intermediate result to disk
        os.makedirs("./outputs", exist_ok=True)
        ts = datetime.now().strftime("%Y%m%d_%H%M%S")
        path = f"./outputs/step1_result_{ts}.png"
        refined.save(path)

        # Step 2: Background Inpainting
        print("[INFO] Step 2: Inpainting background...", flush=True)
        unload_models()
        result = run_background_removal_and_inpaint(
            image_path=path,
            prompt=prompt_2,
            negative_prompt=neg_2,
            guidance_scale=guidance_2
        )

        return refined, result, ""

    except gr.Error as e:
        return None, None, f"🛑 {str(e)}"
    except Exception as e:
        _print_trace()
        return None, None, f"❌ Unexpected Error: {type(e).__name__}: {str(e)}"

# ───────────────────── Gradio UI ─────────────────────
with gr.Blocks() as demo:
    gr.Markdown("## 🧠 Headshot + Background Generator (Full Prompt Control)")

    with gr.Row():
        input_image = gr.Image(type="pil", label="Upload Headshot")

    gr.Markdown("### Step 1: Headshot Refinement (LoRA)")

    with gr.Row():
        prompt_1 = gr.Textbox(label="Headshot Prompt", value="a professional headshot of a confident woman in her 30s with blonde hair")
        neg_1 = gr.Textbox(label="Headshot Negative Prompt", value="deformed, cartoon, anime, sketch, blurry, low quality")

    with gr.Row():
        strength_1 = gr.Slider(0.1, 1.0, value=0.2, step=0.05, label="Refinement Strength")
        guidance_1 = gr.Slider(1, 20, value=17, step=0.5, label="Guidance Scale (Headshot)")

    gr.Markdown("### Step 2: Background Inpainting (SDXL)")

    with gr.Row():
        prompt_2 = gr.Textbox(label="Background Prompt", value="modern hospital background, clean, soft lighting")
        neg_2 = gr.Textbox(label="Background Negative Prompt", value="fantasy, cartoon, cluttered, sketch")

    with gr.Row():
        guidance_2 = gr.Slider(1, 20, value=10, step=0.5, label="Guidance Scale (Background)")

    go_btn = gr.Button("✨ Generate Refined Headshot + Background")

    with gr.Row():
        output_refined = gr.Image(type="pil", label="Step 1: Refined Headshot")
        output_final = gr.Image(type="pil", label="Step 2: Final Image with Background")

    error_box = gr.Markdown(label="Error", value="", visible=True)

    go_btn.click(
        fn=safe_generate_and_inpaint,
        inputs=[
            input_image, prompt_1, neg_1, strength_1, guidance_1,
            prompt_2, neg_2, guidance_2
        ],
        outputs=[output_refined, output_final, error_box]
    )

demo.launch(debug=True)