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):
    if input_image is None:
        raise gr.Error("No image uploaded!")


# Функция обработки изображения
def preprocess(input_image, do_remove_background, foreground_ratio):
    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


# Функция генерации 3D модели
def generate(image):
    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


def start_loading():
    return gr.HTML.update(visible=True)

def stop_loading():
    return gr.HTML.update(visible=False)


# Настройка темы и CSS
class CustomTheme(gr.themes.Base):
    def __init__(self):
        super().__init__()
        self.primary_hue = "#191a1e"
        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 = """
/* Скрываем нижний колонтитул */
footer {
    visibility: hidden;
    height: 0;
    margin: 0;
    padding: 0;
    overflow: hidden;
}
/* Применяем шрифты */
@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;
    background-color: #191a1e !important;
    color: #FFFFFF;
}
/* Настройки заголовков */
h1, h2, h3, h4, h5, h6 {
    font-family: 'Poppins', sans-serif;
    font-weight: 700;
    color: #FFFFFF;
}
/* Стиль для текстовых полей и кнопок */
input[type="text"], textarea {
    background-color: #191a1e !important;
    color: #FFFFFF;
    border: 1px solid #FFFFFF;
}
/* Слайдер */
.gr-slider {
    height: 40px !important; /* Увеличение высоты слайдера */
    transition: all 0.3s ease !important; /* Плавные переходы */
}
.gr-slider .slider-value {
    display: none !important; /* Прячем текст значения */
}
.gr-slider .slider::-webkit-slider-thumb {
    height: 15px !important;
    width: 15px !important;
    background-color: #5271FF !important; /* Цвет совпадает с кнопкой */
    border-radius: 50%;
}
.gr-slider .slider::-webkit-slider-runnable-track {
    background: #5271FF !important; /* Цвет заполнения слайдера */
}
/* Кнопка Generate */
.generate-button {
    background-color: #5271FF !important;
    color: #FFFFFF !important;
    border: none;
    font-weight: bold;
}
.generate-button:hover {
    background-color: #405BBF !important; /* Цвет при наведении */
}
.generate-button {
    position: relative; /* Кнопка становится родительским контейнером */
}
/* Выделяем текст для Prompt */
.prompt-text {
    font-weight: bold;
    color: #FFFFFF;
}
/* Лоадер (анимация) */
#loading-bar {
    display: none; /* Скрыт по умолчанию */
    position: relative; /* Изменено на relative для использования в кнопке */
    margin: 0 auto; /* Центрируем внутри родительского элемента */
    width: 30px; /* Уменьшен размер для кнопки */
    height: 30px;
    border: 4px solid #f3f3f3;
    border-top: 4px solid #5271FF; /* Цвет лоадера */
    border-radius: 50%;
    animation: spin 1s linear infinite;
}
/* Анимация вращения */
@keyframes spin {
    0% { transform: translate(-50%, -50%) rotate(0deg); }
    100% { transform: translate(-50%, -50%) rotate(360deg); }
}
/* Текст CheckBox в белый */
.gr-checkbox label {
    color: #FFFFFF !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",  # Текст перекрашен в белый через CSS
            value=True,
        )
        submit = gr.Button("Generate", elem_classes="generate-button")
        loading_bar = gr.HTML("<div id='loading-bar'></div>")
        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,
)