import gradio as gr
import numpy as np
import random
import spaces
import torch
import time
from diffusers import DiffusionPipeline, AutoencoderTiny
from custom_pipeline import FluxWithCFGPipeline
from transformers import pipeline  # Added for translation

# --- Torch Optimizations ---
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.benchmark = True

# --- Constants ---
MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1024
DEFAULT_WIDTH = 1024
DEFAULT_HEIGHT = 576
ASPECT_RATIOS = {
    "16:9": (1024, 576),
    "1:1": (1024, 1024),
    "9:16": (576, 1024)
}
INFERENCE_STEPS = 8

# --- Device and Model Setup ---
dtype = torch.float16
device = "cuda" if torch.cuda.is_available() else "cpu"

pipe = FluxWithCFGPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype)
pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=dtype)
pipe.to(device)

# --- Translation Pipeline ---
try:
    translator = pipeline("translation", model="Helsinki-NLP/opus-mt-sq-en", force_download=True)
except Exception as e:
    print(f"Warning: Failed to load translation model: {e}. Using input prompt directly.")
    translator = None

# --- Inference Function ---
@spaces.GPU
def generate_image(prompt: str, seed: int = 42, aspect_ratio: str = "16:9", randomize_seed: bool = False):
    if pipe is None:
        raise gr.Error("Pipelinei nuk u ngarkua.")

    if not prompt or prompt.strip() == "":
        return None, seed, "Gabim: Plotësoni përshkrimin."

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    width, height = ASPECT_RATIOS.get(aspect_ratio, (DEFAULT_WIDTH, DEFAULT_HEIGHT))

    # Translate Albanian prompt to English if available
    prompt_final = prompt.strip()
    if translator and prompt_final:
        try:
            prompt_final = translator(prompt_final)[0]['translation_text']
        except Exception as e:
            print(f"Warning: Translation failed: {e}. Using input prompt directly.")
            prompt_final = prompt.strip()
    
    prompt_final += ", ultra realistic, sharp, 8k resolution, cinematic lighting"

    try:
        generator = torch.Generator(device=device).manual_seed(int(float(seed)))
        start_time = time.time()
        with torch.inference_mode():
            image = pipe(
                prompt=prompt_final,
                width=width,
                height=height,
                num_inference_steps=INFERENCE_STEPS,
                generator=generator,
                output_type="pil",
                return_dict=False
            )[0][0]
        latency = time.time() - start_time
        return image, seed, f"Koha e përpunimit: {latency:.2f} sekonda"
    except Exception as e:
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        raise gr.Error(f"Gabim gjatë gjenerimit: {e}")

examples = [
    "Qytet futuristik natën me drita neon",
    "Një mace e bardhë që mban një tabelë përshëndetëse",
    "Një astronaut që del nga një vezë në Hënë",
    "Pamje nga një shtëpi moderne në stilin Minecraft"
]

# --- App Layout ---
with gr.Blocks(css="""
body::before, body::after {
    content: "";
    display: block;
    height: 640px;
    background-color: inherit;
}
button[aria-label="Download"] {
    transform: scale(3);
    transform-origin: top right;
    margin: 0 !important;
    padding: 6px !important;
}
""") as app:
    gr.Markdown("# 🖼️ Gjenerues Imazhesh FLUX")
    gr.Markdown("Përdor modelin FLUX për të krijuar imazhe fantastike.")

    with gr.Row():
        with gr.Column(scale=2):
            output_image = gr.Image(label="Imazhi i Gjeneruar", interactive=False, show_download_button=True)
        with gr.Column(scale=1):
            prompt = gr.Text(label="Përshkrimi (Shqip ose Anglisht)", placeholder="Shkruani se çfarë doni të krijoni...", lines=3)
            generate_btn = gr.Button("🎨 Gjenero")
            aspect_ratio = gr.Radio(label="Raporti i Imazhit", choices=list(ASPECT_RATIOS.keys()), value="16:9")
            randomize_seed = gr.Checkbox(label="Përdor numër të rastësishëm", value=True)
            latency = gr.Text(label="Koha", interactive=False)

    gr.Markdown("### 📌 Shembuj Frymëzues")
    gr.Examples(
        examples=examples,
        fn=generate_image,
        inputs=[prompt],
        outputs=[output_image, gr.Number(visible=False), latency],
        cache_examples=True,
        cache_mode="eager"
    )

    generate_btn.click(
        fn=generate_image,
        inputs=[prompt, gr.Number(value=42, visible=False), aspect_ratio, randomize_seed],
        outputs=[output_image, gr.Number(visible=False), latency],
        show_progress="full"
    )

app.launch(share=True)