app.py

import os
import glob
import uuid
from typing import Tuple, Optional
from io import BytesIO

import gradio as gr
from PIL import Image, ImageFile

from google import genai
from google.genai import types

import logging

logger = logging.getLogger(__name__)

# -----------------------------------------
# Configuration
# -----------------------------------------

TEXTURE_FOLDERS = {
    "Laminált padló": "laminate",
    "Járólap": "tile"
}

ImageFile.LOAD_TRUNCATED_IMAGES = True

IMAGE_MODEL_ID = os.getenv("AI_IMAGE_MODEL_ID", "image-model-default")
print("IMAGE_MODEL_ID:", IMAGE_MODEL_ID)

# -----------------------------------------
# Structured Prompt for A | B Full Room Duplication
# -----------------------------------------
# A prompt most már biztonságosan, környezeti változóból töltődik be.
PROMPT_TEMPLATE = os.getenv("FLOOR_PROMPT_TEMPLATE")

if not PROMPT_TEMPLATE:
    logger.warning("⚠️ FLOOR_PROMPT_TEMPLATE secret is missing! Please set it in Hugging Face Space Settings.")
    # Fallback to prevent immediate crash, though API will likely reject an empty prompt.
    PROMPT_TEMPLATE = "{{ \"error\": \"Secret prompt missing.\" }}"

# Gallery layout
GALLERY_COLUMNS = 5
MAX_IMAGE_DIMENSION = 1280

# -----------------------------------------
# Custom CSS
# -----------------------------------------
CUSTOM_CSS = """
<style>
.gradio-gallery .selected, 
.gradio-gallery .gallery-item.selected,
.gradio-gallery :focus-within {
    border-color: #FF8C00 !important; 
    border-width: 4px !important;
    transform: scale(0.96);            
    transition: all 0.2s ease-in-out;
    outline: 4px solid #FF8C00 !important; 
}
</style>
"""

# -----------------------------------------
# Texture loading
# -----------------------------------------

def load_textures_from_folder(folder_path: str, max_size: Tuple[int, int] = (256, 256)):
    valid_names = []
    gallery_items = []
    path_map = {}

    if not os.path.isdir(folder_path):
        os.makedirs(folder_path, exist_ok=True)
        return valid_names, gallery_items, path_map

    exts = ("*.png", "*.jpg", "*.jpeg", "*.webp")
    found_files = []
    for ext in exts:
        found_files.extend(glob.glob(os.path.join(folder_path, ext)))
    
    found_files = sorted(found_files)

    for path in found_files:
        if not os.path.isfile(path):
            continue
            
        name = os.path.basename(path)
        display_name = os.path.splitext(name)[0]
        
        try:
            img = Image.open(path).convert("RGB")
            img.thumbnail(max_size, Image.LANCZOS)
            
            valid_names.append(display_name)
            gallery_items.append((img, display_name))
            path_map[display_name] = os.path.abspath(path)
            
        except Exception as e:
            continue

    return valid_names, gallery_items, path_map

CATEGORY_DATA = {}
ALL_TEXTURE_PATHS = {}

for category, folder in TEXTURE_FOLDERS.items():
    names, items, paths = load_textures_from_folder(folder)
    CATEGORY_DATA[category] = {
        "names": names,
        "items": items
    }
    ALL_TEXTURE_PATHS.update(paths)


# -----------------------------------------
# Image client
# -----------------------------------------

def get_image_client() -> genai.Client:
    return genai.Client()

try:
    image_client = get_image_client()
    IMAGE_CLIENT_INIT_ERROR = None
except Exception as e:
    image_client = None
    IMAGE_CLIENT_INIT_ERROR = str(e)


# -----------------------------------------
# Helper Functions
# -----------------------------------------

def smart_resize(img: Image.Image, max_dim: int) -> Image.Image:
    width, height = img.size
    if width <= max_dim and height <= max_dim:
        return img
    
    if width > height:
        new_width = max_dim
        new_height = int(height * (max_dim / width))
    else:
        new_height = max_dim
        new_width = int(width * (max_dim / height))
        
    return img.resize((new_width, new_height), Image.LANCZOS)


def build_floor_prompt(
    user_extra_prompt: str,
    flooring_type: str,
    width: int,
    height: int
) -> str:
    flooring_type_en = "Laminate" if flooring_type == "Laminált padló" else "Tile"
    material_desc = "laminate floor planks" if flooring_type == "Laminált padló" else "floor tiles"
    safe_user_prompt = user_extra_prompt if user_extra_prompt else "None"

    return PROMPT_TEMPLATE.format(
        width=width,
        height=height,
        flooring_type=flooring_type_en,
        material_desc=material_desc,
        user_extra_prompt=safe_user_prompt,
    ).strip()


# -----------------------------------------
# Main Logic
# -----------------------------------------

def generate_floor_preview(
    room_image: Optional[Image.Image],
    flooring_type: str,
    selected_texture_name: Optional[str],
    user_extra_prompt: str,
) -> Tuple[Optional[Image.Image], str]:
    
    if IMAGE_CLIENT_INIT_ERROR is not None or image_client is None:
        return None, "❌ A képgeneráló szolgáltatás inicializálása sikertelen."

    if room_image is None:
        return None, "⚠️ Kérlek, először tölts fel egy fotót a szobáról."

    # 1. Bemeneti kép méretezése
    room_image = smart_resize(room_image.convert("RGB"), MAX_IMAGE_DIMENSION)
    target_w, target_h = room_image.size
    
    # 2. KIMENETI MÉRET KISZÁMÍTÁSA (Dupla szélesség!)
    double_w = target_w * 2

    if not selected_texture_name:
        return None, "⚠️ Kérlek, válassz egy textúrát."

    if selected_texture_name not in ALL_TEXTURE_PATHS:
        return None, f"⚠️ A '{selected_texture_name}' textúra nem található. Kérlek, frissítsd az oldalt."

    texture_path = ALL_TEXTURE_PATHS[selected_texture_name]
    try:
        texture_image = Image.open(texture_path).convert("RGB")
    except Exception as e:
        logger.exception("Failed to load texture image")
        return None, "❌ Nem sikerült betölteni a textúra képet."

    # Cache buster a prompt végére
    safe_user_prompt = user_extra_prompt if user_extra_prompt else ""
    prompt_with_cache_buster = f"{safe_user_prompt} [System ID: {uuid.uuid4()}]".strip()

    # EGYETLEN prompt generálása, a DUPLA SZÉLESSÉGET (double_w) adjuk át neki!
    prompt = build_floor_prompt(prompt_with_cache_buster, flooring_type, double_w, target_h)

    try:
        # EGYETLEN API hívás
        response = image_client.models.generate_content(
            model=IMAGE_MODEL_ID,
            contents=[prompt, room_image, texture_image],
            config=types.GenerateContentConfig(
                response_modalities=["IMAGE"],
            ),
        )

        edited_image: Optional[Image.Image] = None
        parts = getattr(response, "parts", None)
        if parts is None and getattr(response, "candidates", None):
            candidate0 = response.candidates[0]
            parts = getattr(candidate0, "content", candidate0).parts

        if not parts:
            return None, "⚠️ A szolgáltatás válaszolt, de nem küldött vissza képet."

        for part in parts:
            img_method = getattr(part, "as_image", None)
            if callable(img_method):
                pil_img = img_method()
                if isinstance(pil_img, Image.Image):
                    edited_image = pil_img.convert("RGB")
                    break
            
            inline = getattr(part, "inline_data", None)
            data = getattr(inline, "data", None) if inline is not None else None
            if data:
                try:
                    edited_image = Image.open(BytesIO(data)).convert("RGB")
                    break
                except Exception:
                    continue

        if edited_image is None:
            return None, "⚠️ A kép dekódolása sikertelen."

        # KULCSFONTOSSÁGÚ: Itt most a DUPLA szélességet ellenőrizzük, 
        # nehogy a Python visszanyomja a képet szimpla méretűre!
        if edited_image.size != (double_w, target_h):
            edited_image = edited_image.resize((double_w, target_h), Image.LANCZOS)

        status = (
            "✅ Dupla padló előnézet elkészült!\n"
            f"- Típus: **{flooring_type}**\n"
            f"- Textúra: **{selected_texture_name}**\n"
            f"Eredmény: **A teljes szoba kétszer renderelve (Balra: Párhuzamos | Jobbra: Merőleges)**"
        )
        return edited_image, status

    except Exception:
        logger.exception("Image service API error")
        return None, "❌ Valami hiba történt a képszolgáltatással. Kérlek, próbáld újra."


# -----------------------------------------
# Ergonomic UI Layout
# -----------------------------------------

with gr.Blocks(title="Padló Előnézet Tervező") as demo:
    
    gr.HTML(CUSTOM_CSS)

    selected_texture_state = gr.State(value=None)
    
    initial_category = "Laminált padló"
    initial_gallery_items = CATEGORY_DATA[initial_category]["items"]
    
    if initial_gallery_items:
        selected_texture_state.value = CATEGORY_DATA[initial_category]["names"][0]

    gr.Markdown(
        """
        # 🏠 Padló látványtervező
        **Tölts fel egy fotót a szobáról, válassz padlótípust, és nézd meg a látványterveket azonnal.**
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            
            room_image_input = gr.Image(
                label="1. Szoba fotójának feltöltése",
                type="pil",
                height=450, 
                sources=["upload", "clipboard"],
            )

            gr.Markdown("### 2. Padló beállítása")
            
            flooring_type_radio = gr.Radio(
                choices=["Laminált padló", "Járólap"],
                value=initial_category,
                label="Válassz padlótípust",
                interactive=True
            )
            
            texture_gallery = gr.Gallery(
                value=initial_gallery_items,
                label="Válassz textúrát",
                columns=GALLERY_COLUMNS, 
                height=320,              
                allow_preview=False,
                interactive=True,
                object_fit="cover",
                type="pil",
                show_label=True
            )
            
            selection_feedback = gr.Markdown(
                f"Kiválasztva: **{selected_texture_state.value if selected_texture_state.value else 'Semmi'}**"
            )
            
            with gr.Accordion("Haladó beállítások", open=False):
                extra_prompt_input = gr.Textbox(
                    label="Egyedi módosítások",
                    lines=2,
                    placeholder="pl. Legyen fényesebb, távolítsd el a szőnyeget...",
                )

            generate_button = gr.Button("✨ Látványtervek Generálása", variant="primary", size="lg")

        with gr.Column(scale=2):
            edited_output_image = gr.Image(
                label="MI Eredmény (A: Párhuzamos | B: Merőleges)", 
                interactive=False,
                height=750 
            )
            
            status_text = gr.Markdown()

    # -----------------------------------------
    # Event Wiring
    # -----------------------------------------

    def on_flooring_type_change(new_type):
        items = CATEGORY_DATA.get(new_type, {}).get("items", [])
        names = CATEGORY_DATA.get(new_type, {}).get("names", [])
        
        new_selection = names[0] if names else None
        feedback_str = f"Kiválasztva: **{new_selection}**" if new_selection else "Kiválasztva: **Semmi**"
        
        return items, new_selection, feedback_str

    flooring_type_radio.change(
        fn=on_flooring_type_change,
        inputs=[flooring_type_radio],
        outputs=[texture_gallery, selected_texture_state, selection_feedback]
    )

    def update_selection(evt: gr.SelectData, current_type):
        names = CATEGORY_DATA.get(current_type, {}).get("names", [])
        if evt.index < len(names):
            name = names[evt.index]
            return name, f"Kiválasztva: **{name}**"
        return None, "Kiválasztva: **Semmi**"

    texture_gallery.select(
        fn=update_selection,
        inputs=[flooring_type_radio], 
        outputs=[selected_texture_state, selection_feedback]
    )

    def clear_output():
        return None, "⏳ Generálás folyamatban... (Dupla széles A|B nézet készül)"

    generate_button.click(
        fn=clear_output,
        inputs=[],
        outputs=[edited_output_image, status_text],
        queue=False
    ).then(
        fn=generate_floor_preview,
        inputs=[room_image_input, flooring_type_radio, selected_texture_state, extra_prompt_input],
        outputs=[edited_output_image, status_text],
    )

if __name__ == "__main__":
    DARK_MODE_JS = """
    () => {
        document.body.classList.toggle('dark', true);
    }
    """
    demo.launch(js=DARK_MODE_JS)