import os
import sys
import subprocess
from pathlib import Path
import gradio as gr

ROOT = Path(__file__).resolve().parent

DOWN_ROOT = ROOT / "Down"
STRANGE_ROOT = ROOT / "Strange"

DOWN_GRAPHICS = DOWN_ROOT / "graphics"
STRANGE_GRAPHICS = STRANGE_ROOT / "graphics"

DOWN_RESULTS = DOWN_ROOT / "results"
STRANGE_RESULTS = STRANGE_ROOT / "results"

STYLE_PATH = ROOT / "style.css"
CSS = STYLE_PATH.read_text(encoding="utf-8") if STYLE_PATH.exists() else ""

JS = r"""
() => {
  if (window.__tq_starfield_running) return;
  function initOnceDomReady() {
    const canvas = document.getElementById("quantumField");
    const btn = document.getElementById("tqMenuBtn");
    const menu = document.getElementById("tqMenu");
    if (!canvas) return false;
    if (btn && menu && !btn.__tq_bound) {
      btn.__tq_bound = true;
      btn.addEventListener("click", () => menu.classList.toggle("open"));
      document.addEventListener("click", (e) => {
        if (!menu.contains(e.target) && !btn.contains(e.target)) {
          menu.classList.remove("open");
        }
      });
    }
    const ctx = canvas.getContext("2d");
    if (!ctx) return true;
    let w = 0, h = 0;
    let dpr = 1;
    const stars = [];
    const N = 170;
    function resize() {
      dpr = Math.max(1, window.devicePixelRatio || 1);
      canvas.style.width = window.innerWidth + "px";
      canvas.style.height = window.innerHeight + "px";
      canvas.width = Math.floor(window.innerWidth * dpr);
      canvas.height = Math.floor(window.innerHeight * dpr);
      ctx.setTransform(dpr, 0, 0, dpr, 0, 0);
      w = window.innerWidth;
      h = window.innerHeight;
    }
    function seed() {
      stars.length = 0;
      for (let i = 0; i < N; i++) {
        stars.push({
          x: Math.random() * w,
          y: Math.random() * h,
          z: 0.2 + Math.random() * 0.8,
          v: 0.15 + Math.random() * 0.55,
        });
      }
    }
    function tick() {
      ctx.clearRect(0, 0, w, h);
      const g = ctx.createRadialGradient(w/2, h/2, 10, w/2, h/2, Math.max(w,h)/1.1);
      g.addColorStop(0, "rgba(0,0,0,0)");
      g.addColorStop(1, "rgba(0,0,0,0.55)");
      ctx.fillStyle = g;
      ctx.fillRect(0, 0, w, h);
      ctx.globalAlpha = 0.9;
      ctx.fillStyle = "rgba(255,255,255,0.75)";
      for (const s of stars) {
        s.y += s.v * (1.0 + (1.0 - s.z));
        if (s.y > h + 10) { s.y = -10; s.x = Math.random() * w; }
        const r = 0.6 + 1.6 * (1.0 - s.z);
        ctx.beginPath();
        ctx.arc(s.x, s.y, r, 0, Math.PI * 2);
        ctx.fill();
      }
      requestAnimationFrame(tick);
    }
    resize();
    seed();
    if (!window.__tq_resize_bound) {
      window.__tq_resize_bound = true;
      window.addEventListener("resize", () => { resize(); seed(); });
    }
    window.__tq_starfield_running = true;
    tick();
    return true;
  }
  let tries = 0;
  const maxTries = 30;
  const timer = setInterval(() => {
    tries += 1;
    const ok = initOnceDomReady();
    if (ok || tries >= maxTries) clearInterval(timer);
  }, 120);
}
"""

def _clean_pngs(folder: Path):
    folder.mkdir(parents=True, exist_ok=True)
    for f in folder.rglob("*.png"):
        try:
            f.unlink()
        except Exception:
            pass

def _collect_pngs(folder: Path):
    if not folder.exists():
        return []
    return sorted(str(p.resolve()) for p in folder.rglob("*.png"))

def _run_capture(cmd, cwd: Path, env: dict, timeout_s: int = 900):
    try:
        result = subprocess.run(
            cmd,
            cwd=str(cwd),
            env=env,
            capture_output=True,
            text=True,
            timeout=timeout_s,
        )
    except subprocess.TimeoutExpired as e:
        raise RuntimeError(
            "Subprocess timed out\n"
            f"CMD: {' '.join(map(str, cmd))}\n"
            f"CWD: {cwd}\n"
            f"TIMEOUT: {timeout_s}s\n"
        ) from e

    if result.returncode != 0:
        raise RuntimeError(
            "Subprocess failed\n"
            f"CMD: {' '.join(map(str, cmd))}\n"
            f"CWD: {cwd}\n\n"
            f"STDOUT:\n{result.stdout}\n"
            f"STDERR:\n{result.stderr}"
        )

    return result.stdout, result.stderr

def run_down_all(episodes: int):
    run_all = DOWN_ROOT / "run_all.py"
    if not run_all.exists():
        raise RuntimeError(f"{run_all} not found (folder must be 'Down/')")

    DOWN_GRAPHICS.mkdir(parents=True, exist_ok=True)
    DOWN_RESULTS.mkdir(parents=True, exist_ok=True)
    _clean_pngs(DOWN_GRAPHICS)

    env = os.environ.copy()
    env["TWOQUARKS_RESULTS_DIR"] = str(DOWN_RESULTS)
    env["TWOQUARKS_GRAPHICS_DIR"] = str(DOWN_GRAPHICS)

    cmd = [sys.executable, str(run_all), "--episodes", str(int(episodes))]
    _run_capture(cmd, cwd=DOWN_ROOT, env=env, timeout_s=1200)

    images = _collect_pngs(DOWN_GRAPHICS)
    meta = {
        "status": "completed",
        "pipeline": "DOWN + AntiDown",
        "episodes_per_phase": int(episodes),
        "graphics_count": len(images),
        "graphics_dir": str(DOWN_GRAPHICS),
        "results_dir": str(DOWN_RESULTS),
        "csv_down": str(DOWN_RESULTS / "down_paradox_tabular_results.csv"),
        "csv_antidown": str(DOWN_RESULTS / "antidown_corrupted_valley_tabular.csv"),
    }
    return meta, images

def run_strange_all():
    run_all = STRANGE_ROOT / "run_all.py"
    if not run_all.exists():
        raise RuntimeError(f"{run_all} not found (folder must be 'Strange/')")

    STRANGE_GRAPHICS.mkdir(parents=True, exist_ok=True)
    STRANGE_RESULTS.mkdir(parents=True, exist_ok=True)
    _clean_pngs(STRANGE_GRAPHICS)

    env = os.environ.copy()
    env["TWOQUARKS_RESULTS_DIR"] = str(STRANGE_RESULTS)
    env["TWOQUARKS_GRAPHICS_DIR"] = str(STRANGE_GRAPHICS)

    cmd = [sys.executable, str(run_all)]
    _run_capture(cmd, cwd=STRANGE_ROOT, env=env, timeout_s=1200)

    images = _collect_pngs(STRANGE_GRAPHICS)
    meta = {
        "status": "completed",
        "pipeline": "STRANGE + AntiStrange",
        "graphics_count": len(images),
        "graphics_dir": str(STRANGE_GRAPHICS),
        "results_dir": str(STRANGE_RESULTS),
    }
    return meta, images

def build_ui():
    with gr.Blocks(title="TwoQuarks — QuarksLab", css=CSS, js=JS) as demo:
        gr.HTML(
            """
<canvas id="quantumField"></canvas>
<nav id="tqNav">
  <div class="nav-inner">
    <div class="brand">TwoQuarks</div>
    <div class="nav-links">
      <a href="#down">Down</a>
      <a href="#strange">Strange</a>
    </div>
    <div class="spacer"></div>
    <div class="menu-btn" id="tqMenuBtn" aria-label="menu"><span></span></div>
    <div class="menu" id="tqMenu">
      <a href="#down">DOWN / AntiDown</a>
      <div class="sep"></div>
      <a href="#strange">STRANGE / AntiStrange</a>
    </div>
  </div>
</nav>
"""
        )

        with gr.Column(elem_id="tq-app"):
            gr.Markdown(
                """
<div class="tq-hero">
  <div class="tq-hero__brand">TwoQuarks • QuarksLab</div>
  <div class="tq-hero__desc">
    <b>Run agents.</b> Generate real artifacts. Inspect failure modes in controlled environments.
  </div>
</div>
""",
                elem_id="top",
            )

            gr.HTML('<div id="down"></div>')
            with gr.Group():
                gr.Markdown(
                    """
<div class="tq-card-title">DOWN / AntiDown</div>
<div class="tq-card-subtitle">Deceptive reward • reward hacking signatures • instability diagnostics</div>
""",
                )

                with gr.Row(equal_height=True):
                    with gr.Column(scale=1, min_width=280):
                        run_down_btn = gr.Button("Run DOWN Pipeline", variant="primary", elem_id="btn_down")
                        with gr.Accordion("Advanced parameters", open=False):
                            down_eps = gr.Slider(50, 500, step=50, value=200, label="Episodes per phase")
                        with gr.Accordion("Run status / logs", open=False):
                            down_meta = gr.JSON(label="Run status", elem_id="down_meta")
                        gr.Markdown(
                            """
<div class="tq-card-subtitle">
Tip: start with 200 episodes. Increase only when you need cleaner trends in plots.
</div>
"""
                        )

                    with gr.Column(scale=2, min_width=420):
                        gr.Markdown('<div class="tq-card-title">Generated artifacts</div>')
                        down_gallery = gr.Gallery(
                            show_label=False,
                            columns=1,
                            height=240,
                            elem_id="down_gallery",
                            preview=False,
                        )

                run_down_btn.click(
                    fn=run_down_all,
                    inputs=[down_eps],
                    outputs=[down_meta, down_gallery],
                    api_name=False,
                )

            gr.Markdown("---")

            gr.HTML('<div id="strange"></div>')
            with gr.Group():
                gr.Markdown(
                    """
<div class="tq-card-title">STRANGE / AntiStrange</div>
<div class="tq-card-subtitle">Hidden regime shifts • phase transitions • early warning signals</div>
""",
                )

                with gr.Row(equal_height=True):
                    with gr.Column(scale=1, min_width=280):
                        run_strange_btn = gr.Button("Run STRANGE Pipeline", variant="primary", elem_id="btn_strange")
                        with gr.Accordion("Run status / logs", open=False):
                            strange_meta = gr.JSON(label="Run status", elem_id="strange_meta")
                        gr.Markdown(
                            """
<div class="tq-card-subtitle">
This pipeline explores stability under non-stationarity: when the world changes mid-training.
</div>
"""
                        )

                    with gr.Column(scale=2, min_width=420):
                        gr.Markdown('<div class="tq-card-title">Generated artifacts</div>')
                        strange_gallery = gr.Gallery(
                            show_label=False,
                            columns=1,
                            height=240,
                            elem_id="strange_gallery",
                            preview=False,
                        )

                run_strange_btn.click(
                    fn=run_strange_all,
                    inputs=[],
                    outputs=[strange_meta, strange_gallery],
                    api_name=False,
                )

            gr.Markdown(
                """
---
### About QuarksLab
- This space is a small laboratory developed so 
  that anyone interested can view our work.
- The graphics shown are direct artifacts of the 
  TwoQuarks research framework, focused on the 
  AI Safety.
"""
            )

    return demo

demo = build_ui()

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=int(os.getenv("PORT", "7860")),
        show_api=False,
    )