"""HuggingFace Space: OpenEuroLLM pretraining progress bar.

Reads ``progress.json`` (current snapshot) uploaded by ``oellm-progress``
(tracker.py). Self-contained — no config.yaml needed on the Space. Locally it
falls back to the repo's ``data/`` dir so the module imports for
`oellm-progress-dashboard`.
"""

import base64
import json
import mimetypes
from datetime import datetime, timezone, timedelta
from pathlib import Path

import gradio as gr

# On the Space, progress.json sits next to this file. Locally it lives under data/.
_here = Path(__file__).parent
DATA_DIR = _here if (_here / "progress.json").exists() else _here.parents[2] / "data"

# Brand palette: dark blue, light blue, orange accent.
_DARKBLUE = "#0a2a66"
_LIGHTBLUE = "#38bdf8"
_ORANGE = "#f97316"
_GRAY_LIGHT = "#e5e7eb"
_GRAY_MED = "#6b7280"
_GRAY_DARK = "#1f2937"
_BG = "#f9fafb"

# A logo is picked up automatically if a file named logo.<ext> sits next to this
# module (deployed to the Space) or in the local data/ dir.
_LOGO_EXTS = (".png", ".svg", ".jpg", ".jpeg", ".webp")


def _logo_data_uri() -> str | None:
    """Return the logo as a data URI, or None. Inlined to avoid Gradio static-file serving."""
    for d in (_here, DATA_DIR):
        for ext in _LOGO_EXTS:
            p = d / f"logo{ext}"
            if p.exists():
                mime = (
                    "image/svg+xml"
                    if ext == ".svg"
                    else (mimetypes.guess_type(p.name)[0] or "image/png")
                )
                b64 = base64.b64encode(p.read_bytes()).decode()
                return f"data:{mime};base64,{b64}"
    return None


def load_snapshot() -> dict:
    path = DATA_DIR / "progress.json"
    if not path.exists():
        return {}
    try:
        return json.loads(path.read_text())
    except (json.JSONDecodeError, OSError):
        return {}



def _human_tokens(n: float) -> str:
    try:
        n = float(n)
    except (TypeError, ValueError):
        return "—"
    for unit, scale in (("T", 1e12), ("B", 1e9), ("M", 1e6), ("K", 1e3)):
        if abs(n) >= scale:
            return f"{n / scale:.2f}{unit}"
    return f"{n:.0f}"



def header_html() -> str:
    """Branded dark-blue header band: logo (if present) · title · website link."""
    snap = load_snapshot()
    title = snap.get("title", "OpenEuroLLM Pretraining Progress")
    website = snap.get("website")
    logo = _logo_data_uri()

    logo_img = (
        f'<img src="{logo}" alt="logo" style="height:60px;width:auto;'
        f'flex:0 0 auto;border-radius:10px">'
        if logo
        else ""
    )
    if website:
        label = website.replace("https://", "").replace("http://", "").rstrip("/")
        link = (
            f'<a href="{website}" target="_blank" rel="noopener" '
            f'style="color:#fff;background:{_ORANGE};padding:8px 14px;border-radius:99px;'
            f'text-decoration:none;font-weight:600;white-space:nowrap">{label} ↗</a>'
        )
    else:
        link = ""

    return (
        f'<div style="display:flex;align-items:center;gap:18px;padding:20px 26px;'
        f"border-radius:16px;background:linear-gradient(100deg,{_DARKBLUE},{_LIGHTBLUE});"
        f'box-shadow:0 2px 10px rgba(10,42,102,.25)">'
        f"{logo_img}"
        f'<h1 style="flex:1;margin:0;color:#fff;font-size:1.55em;letter-spacing:.2px">{title}</h1>'
        f"{link}"
        f"</div>"
    )


def progress_html() -> str:
    snap = load_snapshot()
    if not snap:
        return f"<p style='color:{_GRAY_MED}'>No progress data yet — run <code>oellm-progress</code>.</p>"

    pct = snap.get("progress_pct") or 0.0
    pct_disp = min(max(pct, 0.0), 100.0)
    seen = _human_tokens(snap.get("tokens_seen"))
    total = _human_tokens(snap.get("tokens_total"))
    model = snap.get("model", "Model")
    updated = snap.get("updated", "")
    try:
        _cest = timezone(timedelta(hours=2))
        updated_disp = datetime.fromisoformat(updated).replace(tzinfo=timezone.utc).astimezone(_cest).strftime("%Y-%m-%d %H:%M CEST")
    except (ValueError, TypeError):
        updated_disp = updated

    start = snap.get("start_date")
    try:
        start_disp = datetime.fromisoformat(str(start)).strftime("%d %b %Y") if start else ""
    except (ValueError, TypeError):
        start_disp = str(start)
    meta = " · ".join(
        b for b in (f"Started {start_disp}" if start_disp else "", f"Updated {updated_disp}") if b
    )

    return (
        f'<div style="padding:28px;border-radius:16px;border:1px solid {_GRAY_LIGHT};'
        f'background:{_BG};max-width:920px;margin:0 auto">'
        f'<h2 style="margin:0 0 4px;color:{_DARKBLUE}">{model} Pretraining Progress</h2>'
        f'<p style="margin:0 0 18px;color:{_GRAY_MED};font-size:.9em">{meta}</p>'
        f'<div style="background:{_GRAY_LIGHT};border-radius:99px;height:38px;position:relative;overflow:hidden">'
        f'<div style="background:linear-gradient(90deg,{_DARKBLUE},{_LIGHTBLUE});width:{pct_disp:.2f}%;'
        f'height:100%;border-radius:99px;transition:width .6s"></div>'
        f'<div style="position:absolute;inset:0;display:flex;align-items:center;justify-content:center;'
        f'font-weight:700;color:{_DARKBLUE};font-size:1.05em">{pct:.2f}%</div>'
        f"</div>"
        f'<div style="display:flex;justify-content:space-between;margin-top:12px;font-size:1.15em;font-weight:600">'
        f'<span style="color:{_ORANGE}">{seen} tokens seen</span>'
        f'<span style="color:{_GRAY_DARK}">{total} target</span>'
        f"</div>"
        + "</div>"
    )


_WANDB_URL = "https://wandb.ai/openeurollm-project/baby_9b_dense/workspace?nw=nwusertvosch"


def wandb_html() -> str:
    return (
        f'<p style="text-align:center;margin:4px 0 18px;color:{_GRAY_DARK};font-size:1.05em">'
        f"Follow the training logs live on Weights &amp; Biases: "
        f'<a href="{_WANDB_URL}" target="_blank" rel="noopener" '
        f'style="color:{_LIGHTBLUE};font-weight:700;text-decoration:none">wandb</a>'
        f"</p>"
    )


with gr.Blocks(title="OpenEuroLLM Pretraining Progress") as demo:
    header = gr.HTML(value=header_html())
    bar = gr.HTML(value=progress_html())
    gr.HTML(value=wandb_html())
    refresh = gr.Button("🔄 Refresh")
    refresh.click(
        fn=lambda: (header_html(), progress_html()),
        outputs=[header, bar],
    )


def main():
    # theme on launch() (gradio 6 moved it here from the Blocks constructor).
    theme = gr.themes.Soft(
        primary_hue="blue",
        font=["ui-sans-serif", "system-ui", "sans-serif"],
    )
    demo.launch(theme=theme)


if __name__ == "__main__":
    main()