import os
import re
import html
from datetime import datetime

import gradio as gr
from huggingface_hub import HfApi, hf_hub_download, hf_hub_url

# =========================================================
# CONFIG
# =========================================================
REPO_ID = os.getenv("REPO_ID")
REPO_TYPE = "dataset"
HF_TOKEN = os.getenv("HF_TOKEN")  # optional if repo is public

api = HfApi(token=HF_TOKEN)


# =========================================================
# HELPERS
# =========================================================
def format_ts(ts: int) -> str:
    try:
        return datetime.fromtimestamp(ts).strftime("%d %b %Y, %I:%M %p")
    except Exception:
        return "Unknown date"


def natural_story_title(folder_name: str) -> str:
    # job_1775798146_Sunbeam_s_Secret -> Sunbeam s Secret
    m = re.match(r"job_(\d+)_(.+)$", folder_name)
    if not m:
        return folder_name.replace("_", " ")
    raw = m.group(2).replace("_", " ").strip()
    return raw


def extract_job_ts(folder_name: str) -> int:
    m = re.match(r"job_(\d+)_", folder_name)
    return int(m.group(1)) if m else 0


def build_resolve_url(path: str) -> str:
    return hf_hub_url(
        repo_id=REPO_ID,
        filename=path,
        repo_type=REPO_TYPE,
    )


def list_repo_files():
    """
    Return all file paths from the dataset repo.
    """
    files = []

    # list_repo_tree is the recommended API for tree listing on HF Hub.
    tree = api.list_repo_tree(
        repo_id=REPO_ID,
        repo_type=REPO_TYPE,
        recursive=True,
        expand=True,
    )

    for item in tree:
        path = getattr(item, "path", None)
        if path and not path.endswith("/"):
            files.append(path)

    return files


def collect_stories():
    
    repo_files = list_repo_files()
    stories = {}

    for path in repo_files:
        parts = path.split("/")
        if len(parts) < 2:
            continue

        folder = parts[0]
        if not folder.startswith("job_"):
            continue

        if folder not in stories:
            stories[folder] = {
                "folder": folder,
                "title": natural_story_title(folder),
                "timestamp": extract_job_ts(folder),
                "markdown_path": None,
                "video_path": None,
                "csv_path": None,
                "image_paths": [],
            }

        lower = path.lower()

        if lower.endswith(".md"):
            stories[folder]["markdown_path"] = path
        elif lower.endswith("full_video.mp4"):
            stories[folder]["video_path"] = path
        elif lower.endswith("dataset.csv"):
            stories[folder]["csv_path"] = path
        elif "/images/" in lower and lower.endswith((".png", ".jpg", ".jpeg", ".webp")):
            stories[folder]["image_paths"].append(path)

    # Sort images inside each story
    for _, story in stories.items():
        story["image_paths"] = sorted(
            story["image_paths"],
            key=lambda p: [
                int(x) if x.isdigit() else x.lower()
                for x in re.split(r"(\d+)", p)
            ]
        )

    # Newest first
    sorted_stories = sorted(
        stories.values(),
        key=lambda x: x["timestamp"],
        reverse=True,
    )

    return sorted_stories


def download_markdown(path: str) -> str:
    if not path:
        return "_No story markdown found._"

    try:
        local_path = hf_hub_download(
            repo_id=REPO_ID,
            filename=path,
            repo_type=REPO_TYPE,
            token=HF_TOKEN,
        )
        with open(local_path, "r", encoding="utf-8") as f:
            return f.read()
    except Exception as e:
        return f"_Failed to load markdown: {e}_"


def make_story_card(story: dict) -> str:
    title = html.escape(story["title"])
    subtitle = html.escape(story["folder"])
    created = html.escape(format_ts(story["timestamp"]))

    story_md = download_markdown(story["markdown_path"])
    story_md_html = gr.Markdown().postprocess(story_md)

    video_html = ""
    if story["video_path"]:
        video_url = build_resolve_url(story["video_path"])
        video_html = f"""
        <div class="story-video-wrap">
            <video class="story-video" controls preload="metadata">
                <source src="{video_url}" type="video/mp4">
                Your browser does not support the video tag.
            </video>
        </div>
        """

    gallery_html = ""
    if story["image_paths"]:
        imgs = []
        for img_path in story["image_paths"]:
            img_url = build_resolve_url(img_path)
            img_name = html.escape(os.path.basename(img_path))
            imgs.append(
                f"""
                <a href="{img_url}" target="_blank" class="story-image-link">
                    <img src="{img_url}" alt="{img_name}" class="story-image" loading="lazy">
                </a>
                """
            )
        gallery_html = f"""
        <div class="story-gallery">
            {''.join(imgs)}
        </div>
        """

    dataset_badge = ""
    if story["csv_path"]:
        csv_url = build_resolve_url(story["csv_path"])
        dataset_badge = (
            f'<a class="meta-pill" href="{csv_url}" target="_blank">dataset.csv</a>'
        )

    return f"""
    <section class="story-card">
        <div class="story-header">
            <div>
                <h2>{title}</h2>
                <div class="story-sub">{subtitle}</div>
            </div>
            <div class="story-meta">
                <span class="meta-pill">{created}</span>
                {dataset_badge}
                <span class="meta-pill">{len(story["image_paths"])} images</span>
            </div>
        </div>

        {video_html}

        <div class="story-content">
            <div class="story-text">
                <h3>Story</h3>
                <div class="story-markdown">{story_md_html}</div>
            </div>

            <div class="story-visuals">
                <h3>Images</h3>
                {gallery_html or '<div class="empty-box">No images found.</div>'}
            </div>
        </div>
    </section>
    """


def build_showcase_html():
    stories = collect_stories()

    if not stories:
        return """
        <div class="empty-state">
            <h2>No stories found</h2>
            <p>Could not find any <code>job_*</code> folders in the repo.</p>
        </div>
        """

    cards = [make_story_card(story) for story in stories]

    return f"""
    <div class="page-wrap">
        <div class="hero">
            <div>
                <div class="eyebrow">OhamLab Story Showcase</div>
                <h1>OhamLab Enabled</h1>
                <p>Latest generated stories, images, and videos from the dataset repo.</p>
            </div>
            <div class="hero-stat">
                <div class="stat-num">{len(stories)}</div>
                <div class="stat-label">Stories</div>
            </div>
        </div>

        <div class="stories-list">
            {''.join(cards)}
        </div>
    </div>
    """


def refresh_showcase():
    return build_showcase_html(), f"Last refreshed: {datetime.now().strftime('%d %b %Y, %I:%M:%S %p')}"


# =========================================================
# CUSTOM CSS
# =========================================================
CUSTOM_CSS = """
:root {
  --bg: #f5f7fb;
  --panel: #ffffff;
  --panel-2: #f9fafc;
  --border: rgba(0,0,0,0.08);
  --text: #1a1f36;
  --muted: #6b7280;
  --accent: #2563eb;
  --accent-2: #0ea5e9;
  --shadow: 0 8px 24px rgba(0,0,0,0.06);
}

.gradio-container {
  max-width: 1500px !important;
}

body, .gradio-container {
  background:
    radial-gradient(circle at top left, rgba(124,156,255,0.16), transparent 25%),
    radial-gradient(circle at top right, rgba(144,224,239,0.10), transparent 20%),
    linear-gradient(180deg, #0a0f1d, #0e1422 45%, #0b1020);
}

.page-wrap {
  padding: 6px 4px 30px 4px;
}

.hero {
  display: flex;
  align-items: center;
  justify-content: space-between;
  gap: 24px;
  background: white;
  border: 1px solid var(--border);
  border-radius: 24px;
  padding: 28px;
  margin-bottom: 24px;
  box-shadow: var(--shadow);
}

.eyebrow {
  color: var(--accent-2);
  font-size: 12px;
  font-weight: 700;
  letter-spacing: 0.12em;
  text-transform: uppercase;
  margin-bottom: 8px;
}

.hero h1 {
  margin: 0;
  color: var(--text);
  font-size: 36px;
  line-height: 1.1;
}

.hero p {
  margin: 10px 0 0 0;
  color: var(--muted);
  font-size: 16px;
}

.hero-stat {
  min-width: 130px;
  text-align: center;
  background: rgba(255,255,255,0.04);
  border: 1px solid var(--border);
  border-radius: 20px;
  padding: 18px;
}

.stat-num {
  color: var(--text);
  font-size: 34px;
  font-weight: 800;
}

.stat-label {
  color: var(--muted);
  font-size: 13px;
}

.stories-list {
  display: flex;
  flex-direction: column;
  gap: 22px;
}

.story-card {
background: white;
  border: 1px solid var(--border);
  border-radius: 24px;
  padding: 22px;
  box-shadow: var(--shadow);
  backdrop-filter: blur(8px);
}

.story-header {
  display: flex;
  justify-content: space-between;
  align-items: flex-start;
  gap: 18px;
  margin-bottom: 18px;
}

.story-header h2 {
  margin: 0;
  color: var(--text);
  font-size: 28px;
  line-height: 1.2;
}

.story-sub {
  color: var(--muted);
  font-size: 13px;
  margin-top: 6px;
  word-break: break-word;
}

.story-meta {
  display: flex;
  flex-wrap: wrap;
  gap: 8px;
  justify-content: flex-end;
}

.meta-pill {
  display: inline-flex;
  align-items: center;
  justify-content: center;
  text-decoration: none;
  color: var(--text);
  background: rgba(124,156,255,0.12);
  border: 1px solid rgba(124,156,255,0.25);
  border-radius: 999px;
  padding: 8px 12px;
  font-size: 12px;
  font-weight: 600;
}

.story-video-wrap {
  margin-bottom: 20px;
}

.story-video {
  width: 100%;
  max-height: 580px;
  border-radius: 18px;
  background: #000;
  border: 1px solid var(--border);
}

.story-content {
  display: grid;
  grid-template-columns: 1.15fr 0.85fr;
  gap: 20px;
}

.story-text,
.story-visuals {
  background: rgba(255,255,255,0.025);
  border: 1px solid var(--border);
  border-radius: 20px;
  padding: 18px;
}

.story-text h3,
.story-visuals h3 {
  margin: 0 0 14px 0;
  color: var(--text);
  font-size: 18px;
}

.story-markdown {
  color: var(--text);
  line-height: 1.7;
  font-size: 15px;
}

.story-markdown p,
.story-markdown li,
.story-markdown h1,
.story-markdown h2,
.story-markdown h3,
.story-markdown h4,
.story-markdown h5,
.story-markdown h6,
.story-markdown blockquote {
  color: var(--text) !important;
}

.story-gallery {
  display: grid;
  grid-template-columns: repeat(auto-fill, minmax(170px, 1fr));
  gap: 12px;
}

.story-image-link {
  display: block;
  text-decoration: none;
}

.story-image {
  width: 100%;
  height: 180px;
  object-fit: cover;
  border-radius: 16px;
  border: 1px solid var(--border);
  transition: transform 0.18s ease, box-shadow 0.18s ease;
  background: rgba(255,255,255,0.02);
}

.story-image:hover {
  transform: translateY(-2px);
  box-shadow: 0 10px 22px rgba(0,0,0,0.22);
}

.empty-box,
.empty-state {
  background: rgba(255,255,255,0.025);
  border: 1px dashed var(--border);
  border-radius: 18px;
  padding: 22px;
  color: var(--muted);
}

@media (max-width: 1100px) {
  .story-content {
    grid-template-columns: 1fr;
  }

  .story-header {
    flex-direction: column;
  }

  .story-meta {
    justify-content: flex-start;
  }

  .hero {
    flex-direction: column;
    align-items: flex-start;
  }
}
"""


# =========================================================
# UI
# =========================================================
with gr.Blocks(
    title="Story Factory Showcase",
    theme=gr.themes.Soft(),
    css=CUSTOM_CSS,
) as demo:
    showcase_state = gr.State(value=None)

    with gr.Row():
        refresh_btn = gr.Button("Refresh Stories", variant="primary")
        refreshed_at = gr.Markdown("Loading...")

    showcase_html = gr.HTML()

    demo.load(
        fn=refresh_showcase,
        inputs=[],
        outputs=[showcase_html, refreshed_at],
    )

    refresh_btn.click(
        fn=refresh_showcase,
        inputs=[],
        outputs=[showcase_html, refreshed_at],
    )

demo.launch()