import gradio as gr
import shutil
from pathlib import Path
import os
from huggingface_hub import list_models

# Optional: Spaces API might not exist
try:
    from huggingface_hub import list_spaces
except ImportError:
    list_spaces = None

# ===============================
# Configuration
# ===============================
# Password is stored securely as a Hugging Face secret, not in code
UPLOAD_PASSWORD = os.environ.get("UPLOAD_PASSWORD")  # must be set in Space secrets

WORKFLOW_DIR = Path("/data/workflows")
WORKFLOW_IMG_DIR = Path("/data/workflows_images")
LORA_DIR = Path("/data/loras")
LORA_IMG_DIR = Path("/data/loras_images")

for d in [WORKFLOW_DIR, WORKFLOW_IMG_DIR, LORA_DIR, LORA_IMG_DIR]:
    d.mkdir(parents=True, exist_ok=True)

ALLOWED_FILE_EXTS = {".json", ".yaml", ".yml", ".png", ".safetensors", ".ckpt"}
ALLOWED_IMG_EXTS = {".png", ".jpg", ".jpeg", ".webp"}

# ===============================
# Upload handlers
# ===============================
def upload_workflow(file, image=None, password=""):
    if UPLOAD_PASSWORD is None:
        return "❌ Upload password not set in Space secrets."
    if password != UPLOAD_PASSWORD:
        return "❌ Unauthorized"
    if file is None:
        return "No file uploaded."
    src = Path(file.name)
    ext = src.suffix.lower()
    if ext not in ALLOWED_FILE_EXTS:
        return f"❌ Unsupported file type: `{ext}`"
    dest = WORKFLOW_DIR / src.name
    shutil.copy(src, dest)
    if image:
        img_ext = Path(image.name).suffix.lower()
        if img_ext in ALLOWED_IMG_EXTS:
            shutil.copy(image.name, WORKFLOW_IMG_DIR / (src.stem + img_ext))
    return f"✅ Saved workflow: `{dest.name}`"


def upload_lora(file, image=None, password=""):
    if UPLOAD_PASSWORD is None:
        return "❌ Upload password not set in Space secrets."
    if password != UPLOAD_PASSWORD:
        return "❌ Unauthorized"
    if file is None:
        return "No file uploaded."
    src = Path(file.name)
    ext = src.suffix.lower()
    if ext not in ALLOWED_FILE_EXTS:
        return f"❌ Unsupported file type: `{ext}`"
    dest = LORA_DIR / src.name
    shutil.copy(src, dest)
    if image:
        img_ext = Path(image.name).suffix.lower()
        if img_ext in ALLOWED_IMG_EXTS:
            shutil.copy(image.name, LORA_IMG_DIR / (src.stem + img_ext))
    return f"✅ Saved LoRA: `{dest.name}`"

# ===============================
# Listing functions
# ===============================
def list_uploaded_workflows():
    files = sorted(WORKFLOW_DIR.glob("*"))
    if not files:
        return "No uploaded workflows yet."
    md = ""
    for f in files:
        img_path = WORKFLOW_IMG_DIR / f"{f.stem}.png"
        if img_path.exists():
            md += f"### {f.name}\n![ref image]({img_path})\n---\n"
        else:
            md += f"### {f.name}\n(no image)\n---\n"
    return md


def list_uploaded_loras():
    files = sorted(LORA_DIR.glob("*"))
    if not files:
        return "No uploaded LoRAs yet."
    md = ""
    for f in files:
        img_path = LORA_IMG_DIR / f"{f.stem}.png"
        if img_path.exists():
            md += f"### {f.name}\n![ref image]({img_path})\n---\n"
        else:
            md += f"### {f.name}\n(no image)\n---\n"
    return md

# ===============================
# Hub browsing functions
# ===============================
def browse_loras(limit=20):
    models = list(list_models(
        tags=["lora"],
        sort="downloads",
        direction=-1,
        limit=limit,
    ))
    if not models:
        return "No LoRAs found."
    return "\n---\n".join(
        f"### {m.modelId}\n⬇️ {m.downloads or 0} | ⭐ {m.likes or 0}\nhttps://huggingface.co/{m.modelId}"
        for m in models
    )


def browse_workflows(limit=20):
    uploaded_files = sorted(WORKFLOW_DIR.glob("*"))
    uploaded_md = ""
    if uploaded_files:
        uploaded_md = "## 🧑‍💻 Uploaded Workflows\n\n" + "\n---\n".join(
            f"### {f.name}\nLocal upload"
            for f in uploaded_files
        )
    if list_spaces is None:
        remote_md = "Spaces API unavailable."
    else:
        spaces = list(list_spaces(
            sort="likes",
            direction=-1,
            limit=limit,
        ))
        workflows = [
            s for s in spaces
            if any(tag in (s.tags or []) for tag in ["comfyui", "workflow", "pipeline", "diffusers"])
        ]
        if workflows:
            remote_md = "## 🌍 Community Workflows\n\n" + "\n---\n".join(
                f"### {s.id}\n⭐ {s.likes or 0}\nhttps://huggingface.co/spaces/{s.id}"
                for s in workflows[:limit]
            )
        else:
            remote_md = "No community workflows found."
    return uploaded_md + "\n\n" + remote_md

# ===============================
# Gradio UI
# ===============================
with gr.Blocks() as demo:
    gr.Markdown("# 🤗 Workflow & LoRA Browser")

    with gr.Tabs():
        # Workflows Tab
        with gr.Tab("⚙️ Workflows"):
            wf_out = gr.Markdown()
            gr.Button("Browse Workflows").click(
                browse_workflows,
                outputs=wf_out,
            )

        # LoRAs Tab
        with gr.Tab("🧩 LoRAs"):
            lora_out = gr.Markdown()
            gr.Button("Browse LoRAs").click(
                browse_loras,
                outputs=lora_out,
            )

        # Upload Workflow Tab (always visible, password-protected)
        with gr.Tab("⬆️ Upload Workflow"):
            gr.Markdown("Upload workflow files with optional reference image (.json/.yaml/.png). Owner only (password required).")
            wf_file = gr.File(label="Workflow file")
            wf_image = gr.File(label="Reference image (optional)")
            wf_password = gr.Textbox(label="Password", type="password")
            wf_status = gr.Markdown()
            gr.Button("Upload").click(
                upload_workflow,
                inputs=[wf_file, wf_image, wf_password],
                outputs=wf_status,
            )
            wf_list = gr.Markdown()
            gr.Button("Refresh Uploaded Workflows").click(
                list_uploaded_workflows,
                outputs=wf_list,
            )

        # Upload LoRA Tab (always visible, password-protected)
        with gr.Tab("⬆️ Upload LoRA"):
            gr.Markdown("Upload LoRA files with optional reference image (.safetensors/.ckpt/.png). Owner only (password required).")
            lora_file = gr.File(label="LoRA file")
            lora_image = gr.File(label="Reference image (optional)")
            lora_password = gr.Textbox(label="Password", type="password")
            lora_status = gr.Markdown()
            gr.Button("Upload").click(
                upload_lora,
                inputs=[lora_file, lora_image, lora_password],
                outputs=lora_status,
            )
            lora_list = gr.Markdown()
            gr.Button("Refresh Uploaded LoRAs").click(
                list_uploaded_loras,
                outputs=lora_list,
            )

demo.launch()