import os
import re
import gradio as ui
from huggingface_hub import HfApi, hf_hub_url

REPO_ID = "Crocody/mymuse"
REPO_TYPE = "dataset"
HF_TOKEN = os.getenv("HF_TOKEN") 

api = HfApi(token=HF_TOKEN)

def natural_sort_key(s):
    return [int(text) if text.isdigit() else text.lower() for text in re.split(r'(\d+)', s)]

def get_lora_samples(sort_desc=False):
    try:
        files = api.list_repo_files(repo_id=REPO_ID, repo_type=REPO_TYPE)
        samples = [f for f in files if f.startswith("LoRA_Samples/") and f.lower().endswith(('.png', '.jpg', '.jpeg'))]
        samples.sort(key=lambda x: natural_sort_key(os.path.basename(x)), reverse=sort_desc)
        return samples
    except Exception as e:
        print(f"파일 로드 실패: {e}")
        return []

def update_gallery_data(sort_desc=False):
    paths = get_lora_samples(sort_desc=sort_desc)
    gallery_val = [(hf_hub_url(repo_id=REPO_ID, repo_type=REPO_TYPE, filename=p), os.path.splitext(os.path.basename(p))[0]) for p in paths]
    return gallery_val, paths

def generate_download_links(image_path):
    if not image_path:
        return "<div style='color: #94a3b8; text-align: center; margin-top: 140px; font-size: 15px;'>왼쪽 목록에서 로라 카드를 클릭하면<br>버전별 상태 정보 및 다운로드 창이 활성화됩니다.</div>"
    
    filename = os.path.splitext(os.path.basename(image_path))[0]
    img_url = hf_hub_url(repo_id=REPO_ID, repo_type=REPO_TYPE, filename=image_path)
    target_folders = ["Qwen2509", "Qwen2512", "Zit"]
    
    # 1. 상단 이미지 배치 및 2. 하단 텍스트 내용 유기적 결합 (HTML 레이아웃)
    html_output = f"""
    <div style='display: flex; flex-direction: column; align-items: center; gap: 20px;'>
        <div style='width: 100%; max-width: 280px; aspect-ratio: 3/4; overflow: hidden; border-radius: 8px; box-shadow: 0 2px 8px rgba(0,0,0,0.1);'>
            <img src='{img_url}' style='width: 100%; height: 100%; object-fit: cover;' alt='{filename} preview'/>
        </div>
        
        <div style='width: 100%;'>
            <h3 style='margin: 10px 0 5px 0; font-size: 18px; color: #1e293b; display: flex; align-items: center; gap: 6px;'>📦 {filename} LoRA 다운로드 링크</h3>
            <table class='download-table'>
                <thead>
                    <tr><th>버전 / 폴더</th><th>다운로드 링크</th></tr>
                </thead>
                <tbody>
    """
    
    all_files = api.list_repo_files(repo_id=REPO_ID, repo_type=REPO_TYPE)
    
    for folder in target_folders:
        target_file = f"{folder}/{filename}.safetensors"
        match = [f for f in all_files if f.lower() == target_file.lower()]
        
        if match:
            download_url = hf_hub_url(repo_id=REPO_ID, repo_type=REPO_TYPE, filename=match[0])
            html_output += f"<tr><td><b>{folder}</b></td><td><a href='{download_url}' target='_blank' class='btn-download'>다운로드 ({filename}.safetensors)</a></td></tr>"
        else:
            html_output += f"<tr><td><b>{folder}</b></td><td style='color: #ef4444;'>파일 없음</td></tr>"
            
    html_output += """
                </tbody>
            </table>
        </div>
    </div>
    """
    return html_output

# --- 스타일시트 ---
custom_css = """
body, .gradio-container { background-color: #f4f6f9 !important; font-family: sans-serif !important; }
#lora-gallery { background: transparent !important; border: none !important; box-shadow: none !important; }
#lora-gallery .gallery-item { background: #ffffff !important; border-radius: 12px !important; border: 1px solid #e2e8f0 !important; padding: 14px !important; }
#info-panel { background: #ffffff !important; border-radius: 16px !important; border: 1px solid #e2e8f0 !important; padding: 24px !important; min-height: 500px; }
.download-table { width: 100%; border-collapse: separate; border-spacing: 0; margin-top: 15px; border-radius: 8px; overflow: hidden; border: 1px solid #e2e8f0; }
.download-table th { background-color: #f8fafc; color: #64748b; padding: 12px 16px; font-size: 14px; border-bottom: 1px solid #e2e8f0; }
.download-table td { padding: 14px 16px; font-size: 14px; border-bottom: 1px solid #e2e8f0; color: #334155; }
.btn-download { color: #2563eb !important; font-weight: 600; text-decoration: none; }
.btn-download:hover { text-decoration: underline; }
.nav-btn { font-weight: 600 !important; background: #ffffff !important; border: 1px solid #e2e8f0 !important; border-radius: 8px !important; }
"""

with ui.Blocks(css=custom_css) as demo:
    with ui.Row():
        ui.Markdown("# ✨ mymuse LoRA Management Center")
        with ui.Row(elem_id="action-bar"):
            btn_asc = ui.Button("⬇️ 오름차순", elem_classes="nav-btn", scale=1)
            btn_desc = ui.Button("⬆️ 내림차순", elem_classes="nav-btn", scale=1)
            btn_refresh = ui.Button("🔄 새로고침", elem_classes="nav-btn", scale=1)

    with ui.Row():
        # 왼쪽 그리드 영역
        with ui.Column(scale=2):
            initial_paths = get_lora_samples(sort_desc=False)
            initial_gallery = [(hf_hub_url(repo_id=REPO_ID, repo_type=REPO_TYPE, filename=p), os.path.splitext(os.path.basename(p))[0]) for p in initial_paths]
            
            gallery = ui.Gallery(
                value=initial_gallery,
                columns=4, 
                object_fit="cover",
                height=750,
                allow_preview=False,
                show_label=False,
                elem_id="lora-gallery"
            )
            paths_holder = ui.State(value=initial_paths)

        # 오른쪽 정보 출력 영역 (상단 이미지 + 하단 정보 구조 반영)
        with ui.Column(scale=1):
            with ui.HTML(elem_id="info-panel") as output_area:
                output_area.value = "<div style='color: #94a3b8; text-align: center; margin-top: 180px; font-size: 15px;'>왼쪽 목록에서 로라 카드를 클릭하면<br>버전별 상태 정보 및 다운로드 창이 활성화됩니다.</div>"

    # --- 이벤트 핸들링 ---
    btn_asc.click(
        fn=lambda: update_gallery_data(sort_desc=False),
        inputs=[],
        outputs=[gallery, paths_holder]
    )
    
    btn_desc.click(
        fn=lambda: update_gallery_data(sort_desc=True),
        inputs=[],
        outputs=[gallery, paths_holder]
    )
    
    def refresh_all():
        gallery_val, paths = update_gallery_data(sort_desc=False)
        reset_html = "<div style='color: #94a3b8; text-align: center; margin-top: 180px; font-size: 15px;'>왼쪽 목록에서 로라 카드를 클릭하면<br>버전별 상태 정보 및 다운로드 창이 활성화됩니다.</div>"
        return gallery_val, paths, reset_html

    btn_refresh.click(
        fn=refresh_all,
        inputs=[],
        outputs=[gallery, paths_holder, output_area]
    )

    def on_select(evt: ui.SelectData, paths):
        return generate_download_links(paths[evt.index])

    gallery.select(fn=on_select, inputs=[paths_holder], outputs=[output_area])

demo.launch()