import json
import os
from functools import lru_cache

import gradio as gr
from huggingface_hub import hf_hub_download, hf_hub_url


DATASET_REPO_ID = 'marcuskwan/relight-five-model-comparison-assets'
MODEL_COLUMNS = [
    ("qwen", "Qwen"),
    ("lightx2v", "LightX2V"),
    ("step1x", "Step1X"),
    ("ic_light", "IC-Light"),
    ("firered", "FireRed"),
]


@lru_cache(maxsize=1)
def load_manifest():
    path = hf_hub_download(DATASET_REPO_ID, "manifest.jsonl", repo_type="dataset")
    rows = []
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            if line.strip():
                rows.append(json.loads(line))
    return rows


def asset_url(path):
    return hf_hub_url(DATASET_REPO_ID, filename=path, repo_type="dataset")


def image_cell(label, url):
    return f"""
    <div class='image-cell'>
      <div class='label'>{label}</div>
      <a href='{url}' target='_blank'><img src='{url}'></a>
    </div>
    """


def render(sample_id):
    rows = load_manifest()
    row = next((r for r in rows if r["sample_id"] == sample_id), rows[0])
    cells = [
        image_cell("Reference", asset_url(row["reference_image"])),
        image_cell("Original gen", asset_url(row["original_gen"])),
    ]
    for key, label in MODEL_COLUMNS:
        cells.append(image_cell(label, asset_url(row["outputs"][key])))
    prompt = row["prompt"].replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
    return f"""
    <style>
      .wrap {{ font-family: system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; }}
      .prompt {{ padding: 12px; border: 1px solid #ddd; background: #fafafa; margin-bottom: 14px; white-space: pre-wrap; line-height: 1.45; }}
      .grid {{ display: grid; grid-template-columns: repeat(auto-fit, minmax(180px, 1fr)); gap: 12px; align-items: start; }}
      .image-cell {{ border: 1px solid #ddd; border-radius: 8px; padding: 8px; background: white; }}
      .label {{ font-weight: 700; margin-bottom: 6px; }}
      img {{ width: 100%; max-height: 360px; object-fit: contain; display: block; }}
    </style>
    <div class='wrap'>
      <h3>Sample {row["sample_id"]}</h3>
      <div class='prompt'>{prompt}</div>
      <div class='grid'>{''.join(cells)}</div>
    </div>
    """


with gr.Blocks(title="Relight Five-Model Comparison") as demo:
    gr.Markdown("# Relight Five-Model Comparison")
    rows = load_manifest()
    sample_ids = [r["sample_id"] for r in rows]
    selector = gr.Dropdown(sample_ids, value=sample_ids[0], label="Sample")
    view = gr.HTML(render(sample_ids[0]))
    selector.change(render, inputs=selector, outputs=view)


if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))