import gradio as gr
import time  # Для эмуляции времени загрузки
import tempfile
import numpy as np
import torch
from PIL import Image
from tsr.system import TSR
from tsr.utils import remove_background, resize_foreground, to_gradio_3d_orientation

# Проверяем наличие GPU
device = "cuda:0" if torch.cuda.is_available() else "cpu"

# Загружаем модель
model = TSR.from_pretrained(
    "stabilityai/TripoSR",
    config_name="config.yaml",
    weight_name="model.ckpt",
)
model.renderer.set_chunk_size(131072)
model.to(device)


# Функция для проверки изображения
def check_input_image(input_image):
    try:
        if input_image is None:
            raise gr.Error("No image uploaded!")
    except Exception as e:
        print(f"Error in check_input_image: {e}")
        raise


# Функция обработки изображения
def preprocess(input_image, do_remove_background, foreground_ratio):
    try:
        def fill_background(image):
            image = np.array(image).astype(np.float32) / 255.0
            image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5
            image = Image.fromarray((image * 255.0).astype(np.uint8))
            return image

        if do_remove_background:
            image = input_image.convert("RGB")
            image = remove_background(image)
            image = resize_foreground(image, foreground_ratio)
            image = fill_background(image)
        else:
            image = input_image
            if image.mode == "RGBA":
                image = fill_background(image)
        return image
    except Exception as e:
        print(f"Error in preprocess: {e}")
        raise


# Функция генерации 3D модели
def generate(image):
    try:
        time.sleep(3)  # Эмуляция времени обработки
        scene_codes = model(image, device=device)
        mesh = model.extract_mesh(scene_codes)[0]
        mesh = to_gradio_3d_orientation(mesh)
        mesh_path2 = tempfile.NamedTemporaryFile(suffix=".glb", delete=False)
        mesh.export(mesh_path2.name)
        return mesh_path2.name
    except Exception as e:
        print(f"Error in generate: {e}")
        raise


# Функции управления лоадером
def start_loading():
    print("Start loading triggered")
    return gr.HTML.update(visible=True)

def stop_loading():
    print("Stop loading triggered")
    return gr.HTML.update(visible=False)


# Настройка темы и CSS
class CustomTheme(gr.themes.Base):
    def __init__(self):
        super().__init__()
        self.primary_hue = "#5271FF"  # Сохраняем голубой цвет для кнопок и контейнеров
        self.background_fill_primary = "#191a1e"  # Фоновый цвет страницы
        self.background_fill_secondary = "#191a1e"
        self.background_fill_tertiary = "#191a1e"
        self.text_color_primary = "#FFFFFF"
        self.text_color_secondary = "#FFFFFF"
        self.text_color_tertiary = "#FFFFFF"
        self.input_background_fill = "#191a1e"
        self.input_text_color = "#FFFFFF"


css = """
/* Задаем фоновый цвет всей страницы */
body {
    background-color: #191a1e !important; /* Тёмный фон */
    margin: 0;
    padding: 0;
}

/* Применяем шрифты */
@import url('https://fonts.googleapis.com/css2?family=Poppins:wght@400;500;700&display=swap');
body, input, button, textarea, select, .gr-button {
    font-family: 'Poppins', sans-serif;
    color: #FFFFFF;
}

/* Настройки заголовков */
h1, h2, h3, h4, h5, h6 {
    font-family: 'Poppins', sans-serif;
    font-weight: 700;
    color: #FFFFFF;
}

/* Цвет кнопок */
.generate-button {
    background-color: #5271FF !important; /* Голубой */
    color: #FFFFFF !important;
    border: none;
    font-weight: bold;
}

/* Цвет контейнеров */
.gr-block {
    background-color: #1c1c24 !important; /* Голубоватый тёмный фон */
    border: 1px solid #5271FF !important; /* Голубая рамка */
}
"""

# Интерфейс
with gr.Blocks(theme=CustomTheme(), css=css) as demo:
    with gr.Column():
        gr.Markdown("<span style='color: white; font-weight: bold;'>Upload and Process Your Image</span>")
        with gr.Row():
            input_image = gr.Image(
                label="Upload Image",
                image_mode="RGBA",
                sources="upload",
                type="pil",
                width=400,
                height=300,
            )
            processed_image = gr.Image(
                label="Processed Image",
                interactive=False,
                width=400,
                height=300,
            )
        foreground_ratio = gr.Slider(
            label="Foreground Ratio",
            minimum=0.5,
            maximum=1.0,
            value=0.85,
            step=0.05,
        )
        do_remove_background = gr.Checkbox(
            label="Remove Background",
            value=True,
        )
        submit = gr.Button("Generate", elem_classes="generate-button")
        loading_bar = gr.HTML("<div id='loading-bar'></div>", visible=False)
        output_model = gr.Model3D(
            label="Generated GLB Model",
            interactive=False,
            elem_classes="gr-model3d-container",
        )

    submit.click(
        fn=start_loading,  # Включить прогресс-бар
        inputs=[], 
        outputs=[loading_bar]
    ).then(
        fn=check_input_image,
        inputs=[input_image],
        outputs=[]
    ).then(
        fn=preprocess,
        inputs=[input_image, do_remove_background, foreground_ratio],
        outputs=[processed_image]
    ).then(
        fn=generate,
        inputs=[processed_image],
        outputs=[output_model]
    ).then(
        fn=stop_loading,  # Отключаем прогресс-бар после завершения генерации
        inputs=[],
        outputs=[loading_bar]
    )

# Запуск приложения
demo.launch(
    server_name="0.0.0.0",
    server_port=7860,
    share=True,
)