import shlex
import subprocess
import os

subprocess.run(shlex.split("pip install pip==24.0"), check=True)
subprocess.run(
    shlex.split(
        "pip install package/onnxruntime_gpu-1.17.0-cp310-cp310-manylinux_2_28_x86_64.whl --force-reinstall --no-deps"
    ), check=True
)
subprocess.run(
    shlex.split(
        "pip install package/nvdiffrast-0.3.1.torch-cp310-cp310-linux_x86_64.whl --force-reinstall --no-deps"
    ), check=True
)

# 모델 체크포인트 다운로드 및 torch 설정
if __name__ == "__main__":
    from huggingface_hub import snapshot_download

    snapshot_download("public-data/Unique3D", repo_type="model", local_dir="./ckpt")

    import os
    import sys
    sys.path.append(os.curdir)
    import torch
    torch.set_float32_matmul_precision('medium')
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.set_grad_enabled(False)

import fire
import gradio as gr
from gradio_app.gradio_3dgen import create_ui as create_3d_ui
from gradio_app.all_models import model_zoo

# ===============================
# Text-to-IMAGE 관련 API 함수 정의
# ===============================
def text_to_image(height, width, steps, scales, prompt, seed):
    """
    주어진 파라미터를 이용해 외부 API의 /process_and_save_image 엔드포인트를 호출하여 이미지를 생성한다.
    """
    from gradio_client import Client
    client = Client(os.getenv("CLIENT_API"))  # 기본값 설정
    result = client.predict(
        height,
        width,
        steps,
        scales,
        prompt,
        seed,
        api_name="/process_and_save_image"
    )
    if isinstance(result, dict):
        return result.get("url", None)
    else:
        return result

def update_random_seed():
    """
    외부 API의 /update_random_seed 엔드포인트를 호출하여 새로운 랜덤 시드 값을 가져온다.
    """
    from gradio_client import Client
    client = Client(os.getenv("CLIENT_API"))  # 기본값 설정
    return client.predict(api_name="/update_random_seed")


_TITLE = '''✨ 3D LLAMA Studio'''
_DESCRIPTION = '''
### Welcome to 3D Llama Studio - Your Advanced 3D Generation Platform

This platform offers two powerful features:
1. **Text/Image to 3D**: Generate detailed 3D models from text descriptions or reference images
2. **Text to Styled Image**: Create artistic images that can be used for 3D generation

*Note: Both English and Korean prompts are supported (영어와 한글 프롬프트 모두 지원됩니다)*
'''

# CSS 스타일 밝은 테마로 수정
custom_css = """
.gradio-container {
    background-color: #ffffff;
    color: #333333;
}
.tabs {
    background-color: #f8f9fa;
    border-radius: 10px;
    padding: 10px;
    margin: 10px 0;
    box-shadow: 0 2px 4px rgba(0,0,0,0.1);
}
.input-box {
    background-color: #ffffff;
    border: 1px solid #e0e0e0;
    border-radius: 8px;
    padding: 15px;
    margin: 10px 0;
    box-shadow: 0 1px 3px rgba(0,0,0,0.05);
}
.button-primary {
    background-color: #4a90e2 !important;
    border: none !important;
    color: white !important;
    transition: all 0.3s ease;
}
.button-primary:hover {
    background-color: #357abd !important;
    transform: translateY(-1px);
}
.button-secondary {
    background-color: #f0f0f0 !important;
    border: 1px solid #e0e0e0 !important;
    color: #333333 !important;
    transition: all 0.3s ease;
}
.button-secondary:hover {
    background-color: #e0e0e0 !important;
}
.main-title {
    color: #2c3e50;
    font-weight: bold;
    margin-bottom: 20px;
}
.slider-label {
    color: #2c3e50;
    font-weight: 500;
}
.textbox-input {
    border: 1px solid #e0e0e0 !important;
    background-color: #ffffff !important;
}
"""

# Gradio 테마 설정 수정
def launch():
    model_zoo.init_models()
    
    with gr.Blocks(
        title=_TITLE,
        css=custom_css,
        theme=gr.themes.Soft(
            primary_hue="blue",
            secondary_hue="slate",
            neutral_hue="slate",
            font=["Inter", "Arial", "sans-serif"]
        )
    ) as demo:

        with gr.Row():
            gr.Markdown('# ' + _TITLE, elem_classes="main-title")
        gr.Markdown(_DESCRIPTION)
        
        with gr.Tabs() as tabs:
            with gr.Tab("🎨 Text to Styled Image", elem_classes="tab"):
                with gr.Group(elem_classes="input-box"):
                    gr.Markdown("### Image Generation Settings")
                    with gr.Row():
                        with gr.Column():
                            height_slider = gr.Slider(
                                label="Image Height",
                                minimum=256,
                                maximum=2048,
                                step=64,
                                value=1024,
                                info="Select image height (pixels)"
                            )
                            width_slider = gr.Slider(
                                label="Image Width",
                                minimum=256,
                                maximum=2048,
                                step=64,
                                value=1024,
                                info="Select image width (pixels)"
                            )
                        with gr.Column():
                            steps_slider = gr.Slider(
                                label="Generation Steps",
                                minimum=1,
                                maximum=100,
                                step=1,
                                value=8,
                                info="More steps = higher quality but slower"
                            )
                            scales_slider = gr.Slider(
                                label="Guidance Scale",
                                minimum=1.0,
                                maximum=10.0,
                                step=0.1,
                                value=3.5,
                                info="How closely to follow the prompt"
                            )
                    
                    prompt_text = gr.Textbox(
                        label="Image Description",
                        placeholder="Enter your prompt here (English or Korean)",
                        lines=3,
                        elem_classes="input-box"
                    )
                    
                    with gr.Row():
                        seed_number = gr.Number(
                            label="Seed (Empty = Random)",
                            value=None,
                            elem_classes="input-box"
                        )
                        update_seed_button = gr.Button(
                            "🎲 Random Seed",
                            elem_classes="button-secondary"
                        )
                    
                    generate_button = gr.Button(
                        "🚀 Generate Image",
                        elem_classes="button-primary"
                    )
                
                with gr.Group(elem_classes="input-box"):
                    gr.Markdown("### Generated Result")
                    image_output = gr.Image(label="Output Image")
                
                update_seed_button.click(
                    fn=update_random_seed,
                    inputs=[],
                    outputs=seed_number
                )
                
                generate_button.click(
                    fn=text_to_image,
                    inputs=[height_slider, width_slider, steps_slider, scales_slider, prompt_text, seed_number],
                    outputs=image_output
                )

            with gr.Tab("🎯 Image to 3D", elem_classes="tab"):
                create_3d_ui("wkl")
                
    demo.queue().launch(share=True)

if __name__ == '__main__':
    fire.Fire(launch)