"""
Hugging Face Spaces entrypoint — Gradio UI for Toy-IR compiler RL demo.
"""

from __future__ import annotations

import json
import os
import sys
import textwrap
from pathlib import Path

import gradio as gr

_ROOT = Path(__file__).resolve().parent
if str(_ROOT) not in sys.path:
    sys.path.insert(0, str(_ROOT))

from runtime_core import (
    CompilerOptimizationEnv,
    Deliverable2_Formatter,
    MOCK_PASSES,
    MockEngine,
    SAMPLE_PROGRAM,
)
from train import run_toy_training


DEFAULT_IR = json.dumps(SAMPLE_PROGRAM, indent=2)


def translate_ir(ir_json: str) -> tuple[str, str]:
    try:
        data = json.loads(ir_json.strip() or "[]")
    except json.JSONDecodeError as e:
        return "", f"Invalid JSON: {e}"
    if not isinstance(data, list):
        return "", "JSON root must be a list of instruction objects."
    d2 = Deliverable2_Formatter.translate_state(data)
    env = CompilerOptimizationEnv(MockEngine(), MOCK_PASSES, max_steps=10)
    env.reset(data)
    internal = env.state()
    return d2, internal


def parse_llm_output(llm_text: str) -> str:
    try:
        arr = Deliverable2_Formatter.extract_action_array(llm_text)
        return json.dumps(arr, indent=2)
    except ValueError as e:
        return str(e)


def parse_action_line(line: str) -> list[str]:
    line = line.strip()
    if not line:
        return []
    try:
        got = Deliverable2_Formatter.extract_action_array(line)
        return [str(x).strip() for x in got]
    except ValueError:
        parts = [p.strip().strip("\"'") for p in line.split(",") if p.strip()]
        return [p.lower() for p in parts]


def run_episode(ir_json: str, actions_multiline: str, max_steps: int) -> str:
    try:
        program = json.loads(ir_json.strip() or "[]")
    except json.JSONDecodeError as e:
        return f"Invalid program JSON: {e}"
    if not isinstance(program, list):
        return "Program must be a JSON list."

    lines = [ln for ln in actions_multiline.splitlines() if ln.strip()]
    actions: list[str] = []
    for ln in lines:
        actions.extend(parse_action_line(ln))
    if not actions:
        return "No actions parsed. Enter JSON arrays or comma-separated pass names."

    engine = MockEngine()
    env = CompilerOptimizationEnv(engine, MOCK_PASSES, max_steps=int(max_steps))
    obs0 = env.reset(program)
    log = [
        f"Baseline cycles: {env.previous_cycles}",
        f"Initial observation (env):\n{obs0}",
        "",
        f"Deliverable 2 pseudo-assembly:\n{Deliverable2_Formatter.translate_state(program)}",
        "",
        "--- steps ---",
    ]
    for i, act in enumerate(actions, start=1):
        r = env.step(act)
        log.append(
            f"{i}. {act!r} → reward={r.reward:+.3f} done={r.done} cycles={env.previous_cycles}"
        )
        if r.info:
            slim = {k: v for k, v in r.info.items() if k in ("delta_pct", "error", "no_op", "reason", "terminal_bonus")}
            if slim:
                log.append(f"   info: {slim}")
        if r.done:
            break
    log.append("")
    log.append("Episode summary:")
    log.append(json.dumps(env._episode_summary(), indent=2))
    return "\n".join(log)


def build_demo() -> gr.Blocks:
    with gr.Blocks(title="Compiler Tetris — Toy-IR RL Demo") as demo:
        gr.Markdown(
            textwrap.dedent(
                """
                # Compiler optimization (Toy-IR) — interactive demo

                This Space runs **CPU-only** demos: **Deliverable 2** state translation and action parsing,
                plus the **CompilerOptimizationEnv** loop. A **toy REINFORCE** tab trains a tiny
                stateless policy over the mock passes (see `train.py`). Full **GRPO + LLM + Unsloth**
                still belongs on Colab or a GPU machine.
                """
            ).strip()
        )

        with gr.Tabs():
            with gr.Tab("Translate IR (D2)"):
                gr.Markdown("Paste Toy-IR as a JSON **list** of instruction dicts (`op`, `args`, `dest`, …).")
                ir_in = gr.Textbox(label="IR JSON", value=DEFAULT_IR, lines=12, max_lines=24)
                btn_t = gr.Button("Translate", variant="primary")
                out_d2 = gr.Textbox(label="Deliverable2 pseudo-assembly (LLM-facing)", lines=10)
                out_env = gr.Textbox(label="Env internal pseudo-asm (with type hints)", lines=10)
                btn_t.click(translate_ir, [ir_in], [out_d2, out_env])

            with gr.Tab("Parse LLM output (D2)"):
                gr.Markdown(
                    "Paste model output. A JSON array is preferred; bracket extraction handles light noise."
                )
                llm_in = gr.Textbox(
                    label="LLM output",
                    value='Here is the plan: ["constant_folding", "dead_code_elimination"]',
                    lines=4,
                )
                out_parse = gr.Textbox(label="Parsed array or error", lines=6)
                gr.Button("Parse", variant="primary").click(parse_llm_output, [llm_in], [out_parse])

            with gr.Tab("Env episode (mock engine)"):
                gr.Markdown(
                    "One JSON program + actions: each line can be a JSON array or comma-separated names."
                )
                ir_ep = gr.Textbox(label="Program JSON", value=DEFAULT_IR, lines=10)
                acts = gr.Textbox(
                    label="Actions (one JSON array or comma-list per line)",
                    value='["constant_folding", "dead_code_elimination"]\nloop_unrolling',
                    lines=5,
                )
                ms = gr.Slider(1, 20, value=10, step=1, label="max_steps")
                out_ep = gr.Textbox(label="Log", lines=20)
                gr.Button("Run episode", variant="primary").click(run_episode, [ir_ep, acts, ms], [out_ep])

            with gr.Tab("Toy training (REINFORCE)"):
                gr.Markdown(
                    textwrap.dedent(
                        """
                        Trains a **stateless** categorical policy over the three mock passes
                        (`constant_folding`, `dead_code_elimination`, `loop_unrolling`) using
                        **REINFORCE** in pure Python. Same logic as: `python train.py --episodes 50`.

                        This is a CPU smoke run, not an LLM. For real GRPO, use your project notebooks
                        on a GPU.
                        """
                    ).strip()
                )
                tr_ep = gr.Slider(5, 200, value=50, step=1, label="episodes")
                tr_ms = gr.Slider(2, 20, value=8, step=1, label="max_steps per episode")
                tr_seed = gr.Number(value=0, label="random seed", precision=0)
                tr_lr = gr.Slider(0.01, 0.5, value=0.15, step=0.01, label="learning rate")
                out_tr = gr.Textbox(label="Training log", lines=18)
                gr.Button("Run training", variant="primary").click(
                    run_toy_training, [tr_ep, tr_ms, tr_seed, tr_lr], [out_tr]
                )

        gr.Markdown(
            "Source notebooks in the parent repo: `compiler_optimization_grpo.ipynb`, "
            "`role2_deliverable3_training_loop (2) (1).ipynb`, `compiler_tetris (1).ipynb`, `metahack1 (1).ipynb`."
        )

    return demo


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