import sys
import os

# --- INYECCIÓN ATÓMICA REFORZADA ---
try:
    import huggingface_hub
    class MockHfFolder:
        @staticmethod
        def get_token(): return os.getenv("HF_TOKEN")
        @staticmethod
        def save_token(token): pass
        @staticmethod
        def delete_token(): pass
    huggingface_hub.HfFolder = MockHfFolder
    sys.modules["huggingface_hub.HfFolder"] = MockHfFolder
    setattr(huggingface_hub, "HfFolder", MockHfFolder)
except: pass

try:
    import audioop_lts
    sys.modules["audioop"] = audioop_lts
except:
    from unittest.mock import MagicMock
    sys.modules["audioop"] = MagicMock()
# ---------------------------------------------

import gradio as gr

# --- SILENCIADOR DE API ---
def fake_get_api_info(self, *args, **kwargs):
    return {"components": [], "endpoints": []}
gr.Blocks.get_api_info = fake_get_api_info
# ----------------------------------------------------

import spaces
import torch
import numpy as np
from PIL import Image
import tempfile

# CONFIG
BASE_MODEL = "cyberdelia/CyberRealisticPony"
LTX_MODEL  = "Lightricks/LTX-Video"
LTX_NSFW_LORA = "Lora-Daddy/Ltx2.3-real-nudity-early-alpha-30k-steps"
NEG_DEFAULT = "blurry, low quality, bad anatomy, deformed, ugly, watermark, text"

def load_t2i(is_img2img=False):
    from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline
    cls = StableDiffusionXLImg2ImgPipeline if is_img2img else StableDiffusionXLPipeline
    pipe = cls.from_pretrained(
        BASE_MODEL, torch_dtype=torch.float16, use_safetensors=True, variant="fp16"
    )
    return pipe

def load_video():
    from diffusers import LTXPipeline
    pipe = LTXPipeline.from_pretrained(LTX_MODEL, torch_dtype=torch.bfloat16)
    try:
        pipe.load_lora_weights(LTX_NSFW_LORA)
    except: pass
    return pipe

@spaces.GPU(duration=100)
def generate_t2i(prompt, neg, lora_id, lora_scale, w, h, init_img):
    is_img2img = init_img is not None
    pipe = load_t2i(is_img2img).to("cuda")
    
    if lora_id and len(lora_id.strip()) > 5:
        try:
            pipe.load_lora_weights(lora_id.strip())
            pipe.fuse_lora(lora_scale=lora_scale)
        except: pass

    kwargs = {
        "prompt": prompt, "negative_prompt": neg, "num_inference_steps": 30, 
        "guidance_scale": 7.0, "generator": torch.Generator("cuda").manual_seed(42)
    }
    
    if is_img2img:
        if isinstance(init_img, dict):
            init_img = init_img["composite"] if "composite" in init_image else init_img["background"]
        kwargs["image"] = Image.fromarray(init_img).convert("RGB").resize((int(w), int(h)))
        kwargs["strength"] = 0.6 # Balance entre original y prompt
    else:
        kwargs["width"] = int(w)
        kwargs["height"] = int(h)

    img = pipe(**kwargs).images[0]
    return img

@spaces.GPU(duration=200)
def generate_video(prompt, init_image, lora_scale):
    from diffusers.utils import export_to_video
    pipe = load_video().to("cuda")
    kwargs = {"prompt": prompt, "negative_prompt": NEG_DEFAULT, "num_frames": 49, 
              "num_inference_steps": 30, "generator": torch.Generator("cuda").manual_seed(42)}
    if init_image is not None:
        if isinstance(init_image, dict):
            init_image = init_image["composite"] if "composite" in init_image else init_image["background"]
        kwargs["image"] = Image.fromarray(init_image).convert("RGB").resize((768, 512))
    if lora_scale > 0:
        kwargs["cross_attention_kwargs"] = {"scale": lora_scale}
    output = pipe(**kwargs)
    tmp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
    export_to_video(output.frames[0], tmp.name, fps=24)
    return tmp.name

with gr.Blocks(title="Image Utility v2.1") as demo:
    gr.HTML("<h1 style='text-align:center;'>🛠 Image Processing Utility v2.1.4</h1>")
    with gr.Tabs():
        with gr.Tab("D-Processor (Image/T2I)"):
            with gr.Row():
                with gr.Column():
                    t2i_p = gr.Textbox(label="Input Data String", lines=3)
                    t2i_img = gr.Image(label="Base Reference (Optional)", type="numpy", sources=["upload", "clipboard"])
                    t2i_n = gr.Textbox(label="Excluded Data", value=NEG_DEFAULT)
                    t2i_lora = gr.Textbox(label="Extension ID")
                    t2i_ls = gr.Slider(0, 1.5, 0.8, label="Extension Weight")
                    with gr.Row():
                        t2i_w = gr.Slider(512, 1024, 1024, step=64, label="X-Axis")
                        t2i_h = gr.Slider(512, 1024, 1024, step=64, label="Y-Axis")
                    t2i_btn = gr.Button("Execute Process", variant="primary")
                t2i_out = gr.Image(label="Output Preview")
            t2i_btn.click(generate_t2i, [t2i_p, t2i_n, t2i_lora, t2i_ls, t2i_w, t2i_h, t2i_img], t2i_out)

        with gr.Tab("M-Sequence (Video)"):
            with gr.Row():
                with gr.Column():
                    v_p = gr.Textbox(label="Motion Vector String", lines=3)
                    v_img = gr.Image(label="Source Buffer", type="numpy", sources=["upload", "clipboard"])
                    v_ls = gr.Slider(0, 1.5, 0.8, label="Motion Weight")
                    v_btn = gr.Button("Process Sequence", variant="primary")
                v_out = gr.Video(label="Sequence Output")
            v_btn.click(generate_video, [v_p, v_img, v_ls], v_out)

demo.queue().launch(show_api=False)