"""
Nova-1-XL Dataset Generation - HF Space UI
By SmilyAI Labs
"""

import os
import json
import threading
import gradio as gr

HF_TOKEN = os.environ.get("HF_TOKEN", "")
os.environ["HF_TOKEN"] = HF_TOKEN

LOCAL_CACHE = "./nova_datagen_cache"
os.makedirs(LOCAL_CACHE, exist_ok=True)

from datagen import generate_dataset, HF_DATASET_REPO

# ── Global state ───────────────────────────────────────────────────────────────
is_running   = False
gen_thread   = None
start_time   = None


# ── Status reader ──────────────────────────────────────────────────────────────

def get_status() -> str:
    meta_path = os.path.join(LOCAL_CACHE, "metadata.json")

    if not os.path.exists(meta_path):
        if is_running:
            return (
                "## 🔄 Starting up...\n\n"
                "Producers are launching. First sample coming soon.\n\n"
                f"Check logs for HTTP 200 responses from endpoint."
            )
        return "## ⏸️ Ready\n\nClick **Start Generation** to begin."

    try:
        with open(meta_path) as f:
            d = json.load(f)

        status     = d.get("status", "unknown")
        n          = d.get("n_generated", 0)
        total_tok  = d.get("total_tokens", 0)
        target_tok = d.get("target_tokens", 815_000_000)
        pct        = d.get("pct_complete", 0)
        rate       = d.get("rate_mh", 0)
        elapsed    = d.get("elapsed_h", 0)
        eta        = (target_tok - total_tok) / max(rate * 1_000_000, 1) / 3600
        cats       = d.get("cat_counts", {})
        q_size     = d.get("queue_size", 0)
        avg_r      = d.get("avg_reasoning_len", 0)
        avg_a      = d.get("avg_response_len", 0)
        dupes      = d.get("duplicates", 0)
        fails      = d.get("failures", 0)
        too_long   = d.get("too_long", 0)

        icon = "✅" if status == "complete" else ("🔄" if is_running else "⏸️")

        # Category breakdown
        cat_lines = "\n".join(
            f"| {k} | {v:,} |"
            for k, v in sorted(cats.items(), key=lambda x: -x[1])
        ) if cats else "| (none yet) | 0 |"

        return f"""
## {icon} Nova-1-XL Generation {'Complete!' if status == 'complete' else 'In Progress...' if is_running else 'Paused'}

### Progress
| Metric | Value |
|---|---|
| **Samples Generated** | {n:,} |
| **Tokens** | {total_tok/1e6:.1f}M / {target_tok/1e6:.0f}M |
| **Complete** | {pct:.1f}% |
| **Rate** | {rate:.1f}M tok/h |
| **Elapsed** | {elapsed:.2f}h |
| **ETA** | {eta:.1f}h |
| **Queue Size** | {q_size} |

### Quality
| Metric | Value |
|---|---|
| **Avg Reasoning Length** | {avg_r:,} chars |
| **Avg Response Length** | {avg_a:,} chars |
| **Duplicates Skipped** | {dupes:,} |
| **Failures** | {fails:,} |
| **Too Long (skipped)** | {too_long:,} |

### Category Breakdown
| Category | Samples |
|---|---|
{cat_lines}

### Dataset
[View on HF Hub](https://huggingface.co/datasets/{HF_DATASET_REPO})
"""

    except Exception as e:
        return f"## ⚠️ Error reading status\n\n```{e}```"


def get_latest_sample() -> str:
    """Show the most recently generated sample."""
    samples_dir = os.path.join(LOCAL_CACHE, "samples")
    if not os.path.exists(samples_dir):
        return "No samples yet."

    files = sorted(os.listdir(samples_dir))
    if not files:
        return "No samples yet."

    latest = os.path.join(samples_dir, files[-1])
    try:
        with open(latest, encoding="utf-8") as f:
            return f.read()
    except Exception as e:
        return f"Error reading sample: {e}"


# ── Controls ───────────────────────────────────────────────────────────────────

def start_generation() -> str:
    global is_running, gen_thread, start_time

    if is_running:
        return "⚠️ Already running!"

    if gen_thread and gen_thread.is_alive():
        return "⚠️ Thread still alive!"

    is_running = True
    import time
    start_time = time.time()

    def run():
        global is_running
        try:
            generate_dataset()
        except Exception as e:
            import traceback
            print(f"Generation error: {e}")
            traceback.print_exc()
        finally:
            is_running = False

    gen_thread = threading.Thread(target=run, daemon=True, name="generator")
    gen_thread.start()

    return "🚀 Generation started! Status updates every 30 seconds."


def stop_generation() -> str:
    global is_running
    is_running = False
    return "⏹️ Stop requested. Current sample will finish before stopping."


def refresh() -> tuple:
    return get_status(), get_latest_sample()


# ── UI ─────────────────────────────────────────────────────────────────────────

with gr.Blocks(
    title="Nova-1-XL Dataset Generator",
    theme=gr.themes.Soft(),
) as demo:

    gr.Markdown("""
# 🌟 Nova-1-XL Dataset Generator
**By SmilyAI Labs** | Generates personality-anchored reasoning training data
""")

    with gr.Row():
        start_btn   = gr.Button("🚀 Start Generation", variant="primary",   scale=3)
        stop_btn    = gr.Button("⏹️ Stop",             variant="stop",      scale=1)
        refresh_btn = gr.Button("🔄 Refresh",           variant="secondary", scale=1)

    msg_box = gr.Textbox(
        label="Control Message",
        interactive=False,
        max_lines=1,
    )

    with gr.Row():
        with gr.Column(scale=2):
            status_md = gr.Markdown(get_status())
        with gr.Column(scale=3):
            sample_box = gr.Code(
                label="Latest Sample",
                language="markdown",
                lines=30,
                value=get_latest_sample(),
            )

    # ── Auto-refresh timer ─────────────────────────────────────────────────────
    try:
        timer = gr.Timer(value=30)
        timer.tick(fn=refresh, outputs=[status_md, sample_box])
    except AttributeError:
        gr.Markdown("*Auto-refresh unavailable - use Refresh button*")

    # ── Button actions ─────────────────────────────────────────────────────────
    start_btn.click(
        fn=start_generation,
        outputs=msg_box,
    )
    stop_btn.click(
        fn=stop_generation,
        outputs=msg_box,
    )
    refresh_btn.click(
        fn=refresh,
        outputs=[status_md, sample_box],
    )


if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
    )