import sys, os, gc, torch, spaces, tempfile
import gradio as gr
from PIL import Image

# PATCH GRADIO
try:
    import gradio_client.utils as client_utils
    if not hasattr(client_utils, "_old_json_schema_to_python_type"):
        client_utils._old_json_schema_to_python_type = client_utils._json_schema_to_python_type
    def patched_json_schema_to_python_type(schema, defs=None):
        if isinstance(schema, bool): return "Any"
        return client_utils._old_json_schema_to_python_type(schema, defs)
    client_utils._json_schema_to_python_type = patched_json_schema_to_python_type
except: pass

def flush():
    gc.collect()
    if torch.cuda.is_available(): torch.cuda.empty_cache()

MODELS = {"Pony Diffusion V6 XL": "cyberdelia/CyberRealisticPony"}
LORAS = {
    "Ninguno": "",
    "💎 NSFW: Real Nudity": "Lora-Daddy/Ltx2.3-real-nudity-early-alpha-30k-steps",
    "📜 DOCS: ID Card": "j0rdan/passport-sdxl",
    "🔫 WEAPONS: Tactical": "Ostris/SDXL_LoRA_Test"
}

@spaces.GPU(duration=120)
def generate(prompt, lora_name, w, h, init_img=None, strength=0.6):
    flush()
    from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline
    p = f"score_9, score_8_up, score_7_up, {prompt}"
    pipe = StableDiffusionXLPipeline.from_pretrained("cyberdelia/CyberRealisticPony", torch_dtype=torch.float16, variant="fp16", low_cpu_mem_usage=True).to("cuda")
    
    lid = LORAS.get(lora_name)
    if lid:
        try: pipe.load_lora_weights(lid)
        except: pass

    if init_img:
        pipe_i2i = StableDiffusionXLImg2ImgPipeline.from_pipe(pipe)
        res = pipe_i2i(prompt=p, image=init_img, strength=strength, num_inference_steps=25).images[0]
        del pipe_i2i
    else:
        res = pipe(prompt=p, num_inference_steps=30, width=int(w), height=int(h)).images[0]
    
    del pipe
    flush()
    return res

@spaces.GPU(duration=250)
def video(prompt, init_img):
    flush()
    from diffusers import LTXPipeline
    from diffusers.utils import export_to_video
    pipe = LTXPipeline.from_pretrained("Lightricks/LTX-Video", torch_dtype=torch.bfloat16, low_cpu_mem_usage=True).to("cuda")
    kw = {"prompt": prompt, "num_inference_steps": 20, "num_frames": 25, "width": 704, "height": 480}
    if init_img: kw["image"] = init_img
    out = pipe(**kw)
    tmp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
    export_to_video(out.frames[0], tmp.name, fps=16)
    del pipe
    flush()
    return tmp.name

with gr.Blocks() as demo:
    gr.HTML("<h1 style='text-align:center;'>🌌 Omni-Studio v3.2</h1>")
    with gr.Tabs():
        with gr.Tab("🖼 Imagen"):
            with gr.Row():
                with gr.Column():
                    p = gr.Textbox(label="Prompt")
                    l = gr.Dropdown(choices=list(LORAS.keys()), value="Ninguno", label="LoRA")
                    w = gr.Slider(512, 1024, 832, step=64)
                    h = gr.Slider(512, 1024, 1216, step=64)
                    img = gr.Image(label="Base", type="pil")
                    st = gr.Slider(0.1, 0.9, 0.6, label="Mod Strength")
                    btn = gr.Button("GENERAR")
                out = gr.Image(label="Resultado")
        with gr.Tab("🎥 Video"):
            with gr.Row():
                with gr.Column():
                    vp = gr.Textbox(label="Prompt")
                    vi = gr.Image(label="Base", type="pil")
                    vbtn = gr.Button("GENERAR VIDEO")
                vout = gr.Video(label="Resultado")

        with gr.Tab("🛒 Dropshipping"):
            gr.Markdown("## 🚀 Top 10 Winning Products - USA Summer 2026")
            gr.Markdown("These products are selected for high conversion and low friction. Sync them with Zendrop/Shopify.")
            
            products_data = [
                ["Pro Neck Fan (Cooling Chip)", "$12.00", "$34.99", "$22.99"],
                ["Solar Pest Repeller (4-Pack)", "$12.00", "$39.99", "$27.99"],
                ["Orthopedic Cooling Pet Bed", "$18.00", "$49.99", "$31.99"],
                ["Portable Electric Espresso", "$22.00", "$59.99", "$37.99"],
                ["Smart Battery Security Cam", "$15.00", "$44.99", "$29.99"],
                ["Bluetooth Sleep Mask", "$9.00", "$29.99", "$20.99"],
                ["Cordless Handheld Vacuum", "$14.00", "$39.99", "$25.99"],
                ["Electric Spin Scrubber", "$18.00", "$49.99", "$31.99"],
                ["Portable Digital Tire Inflator", "$22.00", "$59.99", "$37.99"],
                ["Digital Grip Trainer", "$6.00", "$24.99", "$18.99"]
            ]
            gr.Dataframe(headers=["Product", "Cost", "Target Price", "Margin"], value=products_data)
            
            with gr.Row():
                csv_download = gr.File(label="Download Shopify Import CSV", value="products_import_usa.csv")
                guide_link = gr.Markdown("[📖 View Registration Guide](https://huggingface.co/spaces/cobramv12/image-processor-v2/blob/main/registration_guide.md)")

    btn.click(generate, [p, l, w, h, img, st], out)
    vbtn.click(video, [vp, vi], vout)

demo.queue().launch(show_api=False, server_name="0.0.0.0", server_port=7860)