import gradio as gr
import numpy as np
import random
import spaces
import torch
import time
import requests
from diffusers import DiffusionPipeline, AutoencoderTiny
from custom_pipeline import FluxWithCFGPipeline

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

# Constants
MAX_SEED = np.iinfo(np.int32).max
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"

print("⏳ Loading Flux pipeline...")
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)
print("✅ Flux pipeline loaded.")

@spaces.GPU
def translate_albanian_to_english(text):
    """Translate Albanian to English using sepioo-facebook-translation API."""
    if not text.strip():
        return ""
    for attempt in range(2):
        try:
            response = requests.post(
                "https://sepioo-facebook-translation.hf.space/translate",
                json={"from_language": "sq", "to_language": "en", "input_text": text},
                headers={"accept": "application/json", "Content-Type": "application/json"},
                timeout=5
            )
            response.raise_for_status()
            translated = response.json().get("translate", "")
            print(f"Translation response: {translated}")
            return translated
        except Exception as e:
            print(f"Translation error (attempt {attempt + 1}): {e}")
            if attempt == 1:
                return f"Përkthimi dështoi: {str(e)}. Provoni përsëri ose përdorni anglisht."
    return f"Përkthimi dështoi. Provoni përsëri ose përdorni anglisht."

@spaces.GPU
def generate_image(prompt: str, seed: int = 42, aspect_ratio: str = "16:9", randomize_seed: bool = False):
    """Generate image from English prompt."""
    if pipe is None:
        raise gr.Error("Pipeline nuk u ngarkua.")

    prompt_clean = prompt.strip()
    if not prompt_clean:
        return None, seed, "Gabim: Përshkrimi është bosh."

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

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

    prompt_final = prompt_clean + ", ultra realistic, sharp, 8k resolution"

    print(f"🎯 Prompt for generation: {prompt_final}")

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

# UI Layout
def create_demo():
    with gr.Blocks() as app:
        gr.HTML("""
        <style>
        body::before {
            content: "";
            display: block;
            height: 320px;
            background-color: var(--body-background-fill);
        }
        button[aria-label="Fullscreen"], button[aria-label="Fullscreen"]:hover {
            display: none !important;
            visibility: hidden !important;
            opacity: 0 !important;
            pointer-events: none !important;
        }
        button[aria-label="Share"], button[aria-label="Share"]:hover {
            display: none !important;
        }
        button[aria-label="Download"] {
            transform: scale(3);
            transform-origin: top right;
            margin: 0 !important;
            padding: 6px !important;
        }
        </style>
        """)

        gr.Markdown("")
        gr.Markdown("")

        with gr.Row():
            with gr.Column(scale=2):
                output_image = gr.Image(label="Imazhi i Gjeneruar", interactive=False)
            with gr.Column(scale=1):
                prompt_albanian = gr.Textbox(label="Përshkrimi në Shqip", placeholder="Shkruani përshkrimin në shqip këtu", lines=3)
                prompt_english = gr.Textbox(
                    label="Përshkrimi në Anglisht",
                    placeholder="Përkthimi do të shfaqet këtu",
                    interactive=False,
                    lines=3
                )
                aspect_ratio = gr.Radio(label="Raporti i Imazhit", choices=list(ASPECT_RATIOS.keys()), value="1:1")
                randomize_seed = gr.Checkbox(label="Përdor numër të rastësishëm", value=False)
                generate_btn = gr.Button("Gjenero")
                latency = gr.Textbox(label="Koha", interactive=False)

        prompt_albanian.change(
            fn=translate_albanian_to_english,
            inputs=[prompt_albanian],
            outputs=[prompt_english]
        )

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

    return app

if __name__ == "__main__":
    print(f"Gradio version: {gr.__version__}")
    app = create_demo()
    app.launch()