import os
import re
from datetime import datetime

import gradio as gr
from huggingface_hub import HfApi, hf_hub_download

# =========================================================
# CONFIG
# =========================================================
REPO_ID = os.getenv("REPO_ID")
REPO_TYPE = "dataset"
HF_TOKEN = os.getenv("HF_TOKEN")

api = HfApi(token=HF_TOKEN)

if not REPO_ID:
    raise RuntimeError("REPO_ID environment variable is not set")
if not HF_TOKEN:
    raise RuntimeError("HF_TOKEN environment variable is not set")

LOCAL_CACHE_DIR = "/tmp/story_showcase"
os.makedirs(LOCAL_CACHE_DIR, exist_ok=True)


# =========================================================
# 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:
    m = re.match(r"job_(\d+)_(.+)$", folder_name)
    if not m:
        return folder_name.replace("_", " ")
    return m.group(2).replace("_", " ").strip()


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 list_repo_files():
    files = []
    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)

    for story in stories.values():
        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)
            ]
        )

    return sorted(stories.values(), key=lambda x: x["timestamp"], reverse=True)


def download_repo_file(path: str):
    if not path:
        return None
    try:
        return hf_hub_download(
            repo_id=REPO_ID,
            filename=path,
            repo_type=REPO_TYPE,
            token=HF_TOKEN,
            local_dir=LOCAL_CACHE_DIR,
        )
    except Exception as e:
        print(f"[DOWNLOAD ERROR] {path}: {e}")
        return None


def download_markdown(path: str) -> str:
    if not path:
        return "_No story markdown found._"
    try:
        local_path = download_repo_file(path)
        if not local_path:
            return "_Failed to load markdown._"
        with open(local_path, "r", encoding="utf-8") as f:
            return f.read()
    except Exception as e:
        return f"_Failed to load markdown: {e}_"


# =========================================================
# BUILD UI DATA
# =========================================================
def build_story_outputs():
    stories = collect_stories()

    if not stories:
        return [
            gr.update(value="<div class='page-wrap'><div class='empty-state'><h2>No stories found</h2><p>Could not find any <code>job_*</code> folders in the repo.</p></div></div>")
        ]

    hero_html = 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-actions">
          <div class="hero-stat">
            <div class="stat-num">{len(stories)}</div>
            <div class="stat-label">Stories</div>
          </div>
        </div>
      </div>
    </div>
    """

    outputs = [gr.update(value=hero_html)]

    max_cards = 20
    for idx in range(max_cards):
        if idx < len(stories):
            story = stories[idx]
            story_md = download_markdown(story["markdown_path"])
            video_path = download_repo_file(story["video_path"]) if story["video_path"] else None
            image_paths = [download_repo_file(p) for p in story["image_paths"]]
            image_paths = [p for p in image_paths if p]

            title_md = f"## {story['title']}"
            sub_md = story["folder"]
            meta_html = (
                f"<span class='meta-pill'>{format_ts(story['timestamp'])}</span>"
                f"<span class='meta-pill'>{len(story['image_paths'])} images</span>"
                f"<span class='meta-pill'>{'Theme' if story['csv_path'] else 'No Theme'}</span>"
            )

            outputs.extend([
                gr.update(visible=True),          # card container
                gr.update(value=title_md),        # title
                gr.update(value=sub_md),          # subtitle
                gr.update(value=meta_html),       # meta
                gr.update(value=video_path, visible=bool(video_path)),   # video
                gr.update(value=story_md),        # story markdown
                gr.update(value=image_paths, visible=bool(image_paths)), # gallery
            ])
        else:
            outputs.extend([
                gr.update(visible=False),
                gr.update(value=""),
                gr.update(value=""),
                gr.update(value=""),
                gr.update(value=None, visible=False),
                gr.update(value=""),
                gr.update(value=[], visible=False),
            ])

    return outputs


# =========================================================
# PAGINATION CONFIG
# =========================================================
# =========================================================
# PAGINATION CONFIG
# =========================================================
STORIES_PER_PAGE = 20


# =========================================================
# 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);
}

html, body {
  margin: 0 !important;
  padding: 0 !important;
  width: 100% !important;
  overflow-x: hidden !important;
}

.gradio-container {
  max-width: 100% !important;
  width: 100% !important;
  margin: 0 !important;
  padding-left: 16px !important;
  padding-right: 16px !important;
}

.gradio-container > .main,
.gradio-container .main {
  max-width: 100% !important;
  width: 100% !important;
  margin: 0 !important;
  padding: 0 !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 {
  width: 100%;
  max-width: 100%;
  padding: 6px 4px 30px 4px;
  box-sizing: border-box;
}

.hero-shell {
  display: grid;
  grid-template-columns: 1fr auto;
  gap: 14px;
  align-items: stretch;
  width: 100%;
  margin-bottom: 20px;
}

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

.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-actions {
  display: flex;
  align-items: center;
  gap: 14px;
}

.hero-stat {
  min-width: 170px;
  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;
}

.hero-refresh,
.hero-refresh button,
.hero-refresh > button {
  min-width: 60px !important;
  width: 60px !important;
  height: 60px !important;
  border-radius: 20px !important;
  padding: 0 !important;
  font-size: 24px !important;
}

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

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

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

.story-sub .prose,
.story-sub p {
  margin: 0 !important;
  color: var(--muted) !important;
  font-size: 13px !important;
  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;
  margin-right: 8px;
  margin-bottom: 8px;
}

.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 .prose,
.story-markdown .prose p,
.story-markdown .prose li,
.story-markdown .prose h1,
.story-markdown .prose h2,
.story-markdown .prose h3,
.story-markdown .prose h4,
.story-markdown .prose h5,
.story-markdown .prose h6,
.story-markdown .prose blockquote {
  color: var(--text) !important;
}

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

.story-gallery {
  min-height: 80px;
}

.story-gallery .grid-wrap,
.story-gallery .grid-container,
.story-gallery .gallery {
  border-radius: 16px !important;
}

.pagination-wrap {
  display: flex;
  justify-content: center;
  align-items: center;
  gap: 12px;
  margin: 10px 0 40px 0;
  flex-wrap: wrap;
  width: 100%;
}

.page-indicator {
  background: white;
  border: 1px solid var(--border);
  border-radius: 16px;
  padding: 10px 16px;
  box-shadow: var(--shadow);
  color: var(--text);
  font-weight: 600;
}

.page-btn button,
.page-btn > button {
  border-radius: 16px !important;
  min-width: 120px !important;
  height: 44px !important;
}

.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) {
  .hero-shell {
    grid-template-columns: 1fr;
  }

  .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;
  }

  .hero-refresh,
  .hero-refresh button,
  .hero-refresh > button {
    width: 100% !important;
  }
}
"""


# =========================================================
# OUTPUT BUILDER
# =========================================================
def build_story_outputs(page=1):
    stories = collect_stories()
    total_stories = len(stories)
    total_pages = max(1, (total_stories + STORIES_PER_PAGE - 1) // STORIES_PER_PAGE)

    try:
        page = int(page)
    except Exception:
        page = 1

    page = max(1, min(page, total_pages))

    start_idx = (page - 1) * STORIES_PER_PAGE
    end_idx = start_idx + STORIES_PER_PAGE
    page_stories = stories[start_idx:end_idx]

    if not stories:
        hero_html = """
        <div class="hero">
          <div>
            <div class="eyebrow">OhamLab Story Showcase</div>
            <h1>OhamLab Enabled</h1>
            <p>No stories found in the dataset repo.</p>
          </div>
          <div class="hero-actions">
            <div class="hero-stat">
              <div class="stat-num">0</div>
              <div class="stat-label">Stories Rendered</div>
            </div>
          </div>
        </div>
        """

        outputs = [
            gr.update(value=hero_html),

            gr.update(value="Page 1 of 1"),
            gr.update(interactive=False),
            gr.update(interactive=False),
        ]

        for _ in range(STORIES_PER_PAGE):
            outputs.extend([
                gr.update(visible=False),
                gr.update(value=""),
                gr.update(value=""),
                gr.update(value=""),
                gr.update(value=None, visible=False),
                gr.update(value=""),
                gr.update(value=[], visible=False),
            ])

        return outputs

    hero_html = f"""
    <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-actions">
        <div class="hero-stat">
          <div class="stat-num">{total_stories}</div>
          <div class="stat-label">Stories Rendered</div>
        </div>
      </div>
    </div>
    """

    outputs = [
        gr.update(value=hero_html),

        # Bottom pagination only
        gr.update(value=f"Page {page} of {total_pages}"),
        gr.update(interactive=page > 1),
        gr.update(interactive=page < total_pages),
    ]

    for idx in range(STORIES_PER_PAGE):
        if idx < len(page_stories):
            story = page_stories[idx]
            story_md = download_markdown(story["markdown_path"])
            video_path = download_repo_file(story["video_path"]) if story["video_path"] else None
            image_paths = [download_repo_file(p) for p in story["image_paths"]]
            image_paths = [p for p in image_paths if p]

            title_md = f"## {story['title']}"
            sub_md = story["folder"]
            meta_html = (
                f"<span class='meta-pill'>{format_ts(story['timestamp'])}</span>"
                f"<span class='meta-pill'>{len(story['image_paths'])} images</span>"
                f"<span class='meta-pill'>{'Theme' if story['csv_path'] else 'No Theme'}</span>"
            )

            outputs.extend([
                gr.update(visible=True),
                gr.update(value=title_md),
                gr.update(value=sub_md),
                gr.update(value=meta_html),
                gr.update(value=video_path, visible=bool(video_path)),
                gr.update(value=story_md),
                gr.update(value=image_paths, visible=bool(image_paths)),
            ])
        else:
            outputs.extend([
                gr.update(visible=False),
                gr.update(value=""),
                gr.update(value=""),
                gr.update(value=""),
                gr.update(value=None, visible=False),
                gr.update(value=""),
                gr.update(value=[], visible=False),
            ])

    return outputs


def load_first_page():
    return [gr.update(value=1)] + build_story_outputs(1)


def refresh_to_first_page():
    return [gr.update(value=1)] + build_story_outputs(1)


def go_prev_page(page):
    try:
        page = int(page)
    except Exception:
        page = 1
    new_page = max(1, page - 1)
    return [gr.update(value=new_page)] + build_story_outputs(new_page)


def go_next_page(page):
    stories = collect_stories()
    total_pages = max(1, (len(stories) + STORIES_PER_PAGE - 1) // STORIES_PER_PAGE)
    try:
        page = int(page)
    except Exception:
        page = 1
    new_page = min(total_pages, page + 1)
    return [gr.update(value=new_page)] + build_story_outputs(new_page)


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

    with gr.Column(elem_classes=["page-wrap"]):
        with gr.Row(elem_classes=["hero-shell"]):
            hero_html = gr.HTML()
            refresh_btn = gr.Button("↻", variant="secondary", elem_classes=["hero-refresh"], scale=0)

        with gr.Column(elem_classes=["stories-list"]):
            card_components = []

            for _ in range(STORIES_PER_PAGE):
                with gr.Column(visible=False, elem_classes=["story-card"]) as card:
                    with gr.Row(elem_classes=["story-header"]):
                        with gr.Column(scale=6):
                            title_md = gr.Markdown()
                            sub_md = gr.Markdown(elem_classes=["story-sub"])
                        with gr.Column(scale=4):
                            meta_html = gr.HTML(elem_classes=["story-meta"])

                    video_comp = gr.Video(label=None, visible=False, elem_classes=["story-video-wrap"])

                    with gr.Row(elem_classes=["story-content"]):
                        with gr.Column(elem_classes=["story-text"]):
                            gr.Markdown("### Story")
                            story_md = gr.Markdown(elem_classes=["story-markdown"])

                        with gr.Column(elem_classes=["story-visuals"]):
                            gr.Markdown("### Images")
                            gallery_comp = gr.Gallery(
                                label=None,
                                columns=4,
                                height="auto",
                                preview=True,
                                visible=False,
                                elem_classes=["story-gallery"],
                            )

                card_components.extend([
                    card,
                    title_md,
                    sub_md,
                    meta_html,
                    video_comp,
                    story_md,
                    gallery_comp,
                ])

        # BOTTOM pagination only
        with gr.Row(elem_classes=["pagination-wrap"]):
            prev_btn_bottom = gr.Button("← Previous", variant="secondary", elem_classes=["page-btn"], scale=0)
            page_info_bottom = gr.Markdown("Page 1 of 1", elem_classes=["page-indicator"])
            next_btn_bottom = gr.Button("Next →", variant="secondary", elem_classes=["page-btn"], scale=0)

    all_outputs = [
        hero_html,
        page_info_bottom,
        prev_btn_bottom,
        next_btn_bottom,
    ] + card_components

    demo.load(
        fn=load_first_page,
        inputs=[],
        outputs=[page_state] + all_outputs,
    )

    refresh_btn.click(
        fn=refresh_to_first_page,
        inputs=[],
        outputs=[page_state] + all_outputs,
    )

    prev_btn_bottom.click(
        fn=go_prev_page,
        inputs=[page_state],
        outputs=[page_state] + all_outputs,
    )

    next_btn_bottom.click(
        fn=go_next_page,
        inputs=[page_state],
        outputs=[page_state] + all_outputs,
    )

demo.launch(share=True)