import os

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"
os.environ["MKL_THREADING_LAYER"] = "GNU"
os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"
os.environ["TF_USE_LEGACY_KERAS"] = "1"
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"

import json
import time
from datetime import datetime
from typing import List, Optional, Tuple

import gradio as gr
from fpdf import FPDF
from pipeline import NeuroSightPipeline

pipeline = NeuroSightPipeline()

# ─────────────────────────────────────────────
# CSS  (cross-browser: Safari 14+, Chrome, Firefox, Edge)
# ─────────────────────────────────────────────
CSS = """
@import url('https://fonts.googleapis.com/css2?family=Space+Grotesk:wght@300;400;500;600;700&family=JetBrains+Mono:wght@400;500&display=swap');

/* ── Reset ── */
*, *::before, *::after {
  -webkit-box-sizing: border-box;
  box-sizing: border-box;
  margin: 0;
  padding: 0;
}

/* ── Design tokens ── */
:root {
  --bg:         #0A0C0F;
  --surface:    #111318;
  --surface-2:  #181C23;
  --border:     rgba(255,255,255,0.07);
  --border-lit: rgba(0,194,255,0.35);
  --accent:     #00C2FF;
  --accent-dim: rgba(0,194,255,0.12);
  --purple:     #7B61FF;
  --green:      #00E5A0;
  --green-dim:  rgba(0,229,160,0.12);
  --red:        #FF5A5A;
  --text:       #E8ECF0;
  --text-2:     #8C9BAB;
  --text-3:     #4A5568;
  --radius:     14px;
  --radius-sm:  8px;
  --font:       'Space Grotesk', -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
  --mono:       'JetBrains Mono', 'SF Mono', 'Menlo', 'Consolas', monospace;
}

/* ── Gradio shell overrides ── */
body, .gradio-container, #root {
  background: var(--bg) !important;
  font-family: var(--font) !important;
  color: var(--text) !important;
  min-height: 100vh;
}
.gradio-container {
  max-width: 1200px !important;
  margin: 0 auto !important;
  padding: 0 24px 80px !important;
}
footer { display: none !important; }
.svelte-1ipelgc, .wrap { background: transparent !important; }

/* ── Header ── */
#ns-header {
  text-align: center;
  padding: 60px 0 48px;
  position: relative;
}
/* Safari: pseudo-element gradient needs solid fallback color */
#ns-header::after {
  content: '';
  display: block;
  margin: 32px auto 0;
  width: 120px;
  height: 1px;
  /* Solid fallback for very old Safari */
  background-color: var(--accent);
  background-image: -webkit-linear-gradient(left, transparent, var(--accent), transparent);
  background-image: linear-gradient(90deg, transparent, var(--accent), transparent);
}

.ns-wordmark {
  font-size: 13px;
  font-weight: 600;
  letter-spacing: 0.22em;
  text-transform: uppercase;
  color: var(--accent);
  margin-bottom: 16px;
  display: -webkit-flex;
  display: flex;
  -webkit-align-items: center;
  align-items: center;
  -webkit-justify-content: center;
  justify-content: center;
  /* Safari: use column-gap instead of gap shorthand for older versions */
  -webkit-column-gap: 10px;
  column-gap: 10px;
}
.ns-wordmark::before, .ns-wordmark::after {
  content: '';
  display: inline-block;
  width: 28px;
  height: 1px;
  background: var(--accent);
  opacity: 0.5;
}

.ns-title {
  /* Fallback color for browsers that don't support background-clip: text */
  color: #E8ECF0;
  font-size: 48px; /* Static fallback for clamp() */
  font-size: clamp(38px, 6vw, 58px);
  font-weight: 700;
  letter-spacing: -0.02em;
  line-height: 1.0;
  /* Gradient text: vendor-prefixed first, then standard */
  background: -webkit-linear-gradient(45deg, #E8ECF0 30%, var(--accent) 100%);
  background: linear-gradient(135deg, #E8ECF0 30%, var(--accent) 100%);
  -webkit-background-clip: text;
  background-clip: text;
  /* text-fill-color must come AFTER background-clip for Safari to parse correctly */
  -webkit-text-fill-color: transparent;
}

.ns-subtitle {
  margin-top: 12px;
  font-size: 15px;
  color: var(--text-2);
  font-weight: 400;
  letter-spacing: 0.01em;
}

/* ── Upload zone ──
   Safari: gr.Row renders a div; we must target it with explicit display:grid
   because Gradio's own flex layout overrides can fight the CSS.
   We use a wrapper class approach. */
.upload-zone-wrapper {
  display: -webkit-box;
  display: -ms-flexbox;
  display: -webkit-flex;
  display: flex;
  -webkit-flex-wrap: wrap;
  flex-wrap: wrap;
  gap: 16px;
  row-gap: 16px;
  -webkit-column-gap: 16px;
  column-gap: 16px;
  margin-bottom: 20px;
}
.upload-zone-wrapper > * {
  -webkit-flex: 1 1 220px;
  flex: 1 1 220px;
  min-width: 0;
}

/* Also target Gradio's generated row ID for grid layout */
#upload-zone {
  display: -webkit-flex !important;
  display: flex !important;
  -webkit-flex-wrap: wrap !important;
  flex-wrap: wrap !important;
  gap: 16px !important;
  row-gap: 16px !important;
  -webkit-column-gap: 16px !important;
  column-gap: 16px !important;
  margin-bottom: 20px;
}
#upload-zone > * {
  -webkit-flex: 1 1 220px !important;
  flex: 1 1 220px !important;
  min-width: 0 !important;
}

.upload-card {
  background: var(--surface);
  border: 1px solid var(--border);
  border-radius: var(--radius);
  padding: 20px;
  /* Safari: explicit transition properties, no shorthand ambiguity */
  -webkit-transition: border-color 0.25s ease, box-shadow 0.25s ease;
  transition: border-color 0.25s ease, box-shadow 0.25s ease;
}
.upload-card:hover {
  border-color: rgba(0,194,255,0.2);
  -webkit-box-shadow: 0 0 0 1px rgba(0,194,255,0.08), 0 8px 32px rgba(0,0,0,0.4);
  box-shadow: 0 0 0 1px rgba(0,194,255,0.08), 0 8px 32px rgba(0,0,0,0.4);
}
.upload-card-label {
  display: -webkit-flex;
  display: flex;
  -webkit-align-items: center;
  align-items: center;
  -webkit-column-gap: 8px;
  column-gap: 8px;
  font-size: 12px;
  font-weight: 600;
  letter-spacing: 0.08em;
  text-transform: uppercase;
  color: var(--text-2);
  margin-bottom: 14px;
}
.upload-card-label .badge {
  display: -webkit-inline-flex;
  display: inline-flex;
  -webkit-align-items: center;
  align-items: center;
  -webkit-justify-content: center;
  justify-content: center;
  width: 22px;
  height: 22px;
  border-radius: 6px;
  background: var(--accent-dim);
  color: var(--accent);
  font-size: 10px;
  font-weight: 700;
  font-family: var(--mono);
  -webkit-flex-shrink: 0;
  flex-shrink: 0;
}

/* Gradio file widget overrides */
.upload-card .wrap { background: transparent !important; }
.upload-card .upload-container,
.upload-card [data-testid="fileupload"],
.upload-card .upload {
  background: var(--surface-2) !important;
  border: 1.5px dashed rgba(255,255,255,0.1) !important;
  border-radius: var(--radius-sm) !important;
  -webkit-transition: border-color 0.2s ease !important;
  transition: border-color 0.2s ease !important;
}
.upload-card .upload-container:hover,
.upload-card [data-testid="fileupload"]:hover {
  border-color: rgba(0,194,255,0.4) !important;
}
.upload-card .file-preview {
  background: var(--surface-2) !important;
  border-radius: var(--radius-sm) !important;
}

/* ── Scenario panel ── */
#scenario-panel {
  background: var(--surface);
  border: 1px solid var(--border);
  border-radius: var(--radius);
  padding: 20px 24px;
  margin-bottom: 20px;
  display: -webkit-flex;
  display: flex;
  -webkit-align-items: center;
  align-items: center;
  -webkit-flex-wrap: wrap;
  flex-wrap: wrap;
  gap: 24px;
  row-gap: 16px;
  -webkit-column-gap: 24px;
  column-gap: 24px;
}
.scenario-inputs-status {
  display: -webkit-flex;
  display: flex;
  -webkit-flex-wrap: wrap;
  flex-wrap: wrap;
  gap: 12px;
  row-gap: 8px;
  -webkit-column-gap: 12px;
  column-gap: 12px;
}
.input-chip {
  display: -webkit-inline-flex;
  display: inline-flex;
  -webkit-align-items: center;
  align-items: center;
  -webkit-column-gap: 6px;
  column-gap: 6px;
  padding: 5px 12px;
  border-radius: 100px;
  font-size: 12px;
  font-weight: 500;
  font-family: var(--mono);
  -webkit-transition: all 0.2s ease;
  transition: all 0.2s ease;
}
.input-chip.active {
  background: var(--green-dim);
  color: var(--green);
  border: 1px solid rgba(0,229,160,0.25);
}
.input-chip.inactive {
  background: rgba(255,255,255,0.04);
  color: var(--text-3);
  border: 1px solid var(--border);
}
.scenario-divider {
  width: 1px;
  height: 40px;
  background: var(--border);
  -webkit-flex-shrink: 0;
  flex-shrink: 0;
}
.scenario-info {
  -webkit-flex: 1;
  flex: 1;
  min-width: 200px;
}
.scenario-label {
  font-size: 10px;
  font-weight: 600;
  letter-spacing: 0.1em;
  text-transform: uppercase;
  color: var(--text-3);
  margin-bottom: 4px;
}
.scenario-name {
  font-size: 14px;
  font-weight: 600;
  color: var(--text);
}
.scenario-name.active { color: var(--accent); }
.module-chain {
  display: -webkit-flex;
  display: flex;
  -webkit-align-items: center;
  align-items: center;
  -webkit-flex-wrap: wrap;
  flex-wrap: wrap;
  gap: 4px;
  row-gap: 4px;
  -webkit-column-gap: 4px;
  column-gap: 4px;
  margin-top: 6px;
}
.module-node-sm {
  display: -webkit-inline-flex;
  display: inline-flex;
  -webkit-align-items: center;
  align-items: center;
  -webkit-justify-content: center;
  justify-content: center;
  padding: 2px 8px;
  border-radius: 4px;
  font-size: 11px;
  font-weight: 600;
  font-family: var(--mono);
  background: var(--accent-dim);
  color: var(--accent);
  border: 1px solid rgba(0,194,255,0.2);
}
.module-arrow-sm {
  color: var(--text-3);
  font-size: 10px;
}

/* ── Run button ── */
#run-btn button {
  width: 100% !important;
  height: 52px !important;
  /* Vendor-prefixed gradient for Safari < 15 */
  background: -webkit-linear-gradient(315deg, var(--accent) 0%, var(--purple) 100%) !important;
  background: linear-gradient(135deg, var(--accent) 0%, var(--purple) 100%) !important;
  border: none !important;
  border-radius: var(--radius) !important;
  color: #fff !important;
  font-family: var(--font) !important;
  font-size: 14px !important;
  font-weight: 600 !important;
  letter-spacing: 0.04em !important;
  cursor: pointer !important;
  -webkit-transition: opacity 0.2s ease, -webkit-transform 0.15s ease, -webkit-box-shadow 0.2s ease !important;
  transition: opacity 0.2s ease, transform 0.15s ease, box-shadow 0.2s ease !important;
  -webkit-box-shadow: 0 0 32px rgba(0,194,255,0.2) !important;
  box-shadow: 0 0 32px rgba(0,194,255,0.2) !important;
}
#run-btn button:hover {
  opacity: 0.92 !important;
  -webkit-transform: translateY(-1px) !important;
  transform: translateY(-1px) !important;
  -webkit-box-shadow: 0 0 48px rgba(0,194,255,0.3) !important;
  box-shadow: 0 0 48px rgba(0,194,255,0.3) !important;
}
#run-btn button:active {
  -webkit-transform: translateY(0) !important;
  transform: translateY(0) !important;
}

/* ── NeuroBio Agent link button ── */
#neurobio-agent-btn {
  display: -webkit-flex;
  display: flex;
  -webkit-align-items: center;
  align-items: center;
  -webkit-justify-content: center;
  justify-content: center;
  width: 100%;
  height: 52px;
  margin-top: 4px;
  background: var(--surface);
  border: 1.5px solid var(--border-lit);
  border-radius: var(--radius);
  color: var(--accent);
  font-family: var(--font);
  font-size: 14px;
  font-weight: 600;
  letter-spacing: 0.04em;
  text-decoration: none;
  cursor: pointer;
  -webkit-transition: background 0.2s ease, -webkit-box-shadow 0.2s ease, -webkit-transform 0.15s ease;
  transition: background 0.2s ease, box-shadow 0.2s ease, transform 0.15s ease;
}
#neurobio-agent-btn:hover {
  background: var(--accent-dim);
  -webkit-box-shadow: 0 0 32px rgba(0,194,255,0.18);
  box-shadow: 0 0 32px rgba(0,194,255,0.18);
  -webkit-transform: translateY(-1px);
  transform: translateY(-1px);
}
#neurobio-agent-btn:active {
  -webkit-transform: translateY(0);
  transform: translateY(0);
}

/* ── Animations ── */
@-webkit-keyframes ns-pulse {
  0%, 100% { opacity: 1; }
  50% { opacity: 0.4; }
}
@keyframes ns-pulse {
  0%, 100% { opacity: 1; }
  50% { opacity: 0.4; }
}
@-webkit-keyframes ns-tl-pulse {
  0%, 100% {
    -webkit-box-shadow: 0 0 0 4px rgba(0,194,255,0.08), 0 0 24px rgba(0,194,255,0.25);
    box-shadow: 0 0 0 4px rgba(0,194,255,0.08), 0 0 24px rgba(0,194,255,0.25);
  }
  50% {
    -webkit-box-shadow: 0 0 0 8px rgba(0,194,255,0.04), 0 0 40px rgba(0,194,255,0.35);
    box-shadow: 0 0 0 8px rgba(0,194,255,0.04), 0 0 40px rgba(0,194,255,0.35);
  }
}
@keyframes ns-tl-pulse {
  0%, 100% { box-shadow: 0 0 0 4px rgba(0,194,255,0.08), 0 0 24px rgba(0,194,255,0.25); }
  50% { box-shadow: 0 0 0 8px rgba(0,194,255,0.04), 0 0 40px rgba(0,194,255,0.35); }
}
@-webkit-keyframes ns-flow {
  from { left: -100%; }
  to   { left: 200%; }
}
@keyframes ns-flow {
  from { left: -100%; }
  to   { left: 200%; }
}

/* ── Execution monitor ── */
#exec-monitor {
  background: var(--surface);
  border: 1px solid var(--border);
  border-radius: var(--radius);
  margin-bottom: 20px;
  /* Safari: overflow+border-radius subpixel fix */
  overflow: hidden;
  -webkit-transform: translateZ(0);
  transform: translateZ(0);
}
.exec-monitor-header {
  padding: 16px 24px;
  border-bottom: 1px solid var(--border);
  display: -webkit-flex;
  display: flex;
  -webkit-align-items: center;
  align-items: center;
  -webkit-justify-content: space-between;
  justify-content: space-between;
}
.exec-monitor-title {
  font-size: 12px;
  font-weight: 600;
  letter-spacing: 0.08em;
  text-transform: uppercase;
  color: var(--text-2);
}
.exec-status-dot {
  width: 8px;
  height: 8px;
  border-radius: 50%;
  background: var(--text-3);
  -webkit-flex-shrink: 0;
  flex-shrink: 0;
}
.exec-status-dot.running {
  background: var(--accent);
  -webkit-box-shadow: 0 0 8px var(--accent);
  box-shadow: 0 0 8px var(--accent);
  -webkit-animation: ns-pulse 1.4s ease-in-out infinite;
  animation: ns-pulse 1.4s ease-in-out infinite;
}
.exec-status-dot.done {
  background: var(--green);
  -webkit-box-shadow: 0 0 8px var(--green);
  box-shadow: 0 0 8px var(--green);
}

/* Stage rows */
.stage-list { padding: 12px 0; }
.stage-row {
  display: -webkit-flex;
  display: flex;
  -webkit-align-items: center;
  align-items: center;
  -webkit-column-gap: 14px;
  column-gap: 14px;
  padding: 10px 24px;
  -webkit-transition: background 0.2s ease;
  transition: background 0.2s ease;
}
.stage-row:hover { background: rgba(255,255,255,0.02); }
.stage-icon {
  width: 28px;
  height: 28px;
  border-radius: 8px;
  display: -webkit-flex;
  display: flex;
  -webkit-align-items: center;
  align-items: center;
  -webkit-justify-content: center;
  justify-content: center;
  font-size: 13px;
  -webkit-flex-shrink: 0;
  flex-shrink: 0;
  font-weight: 700;
  line-height: 1;
}
.stage-icon.waiting {
  background: rgba(255,255,255,0.05);
  color: var(--text-3);
  border: 1px solid rgba(255,255,255,0.08);
}
.stage-icon.running {
  background: var(--accent-dim);
  color: var(--accent);
  border: 1px solid rgba(0,194,255,0.3);
  -webkit-animation: ns-pulse 1.2s ease-in-out infinite;
  animation: ns-pulse 1.2s ease-in-out infinite;
}
.stage-icon.done {
  background: var(--green-dim);
  color: var(--green);
  border: 1px solid rgba(0,229,160,0.25);
}
.stage-icon.error {
  background: rgba(255,90,90,0.12);
  color: var(--red);
  border: 1px solid rgba(255,90,90,0.25);
}
.stage-text { -webkit-flex: 1; flex: 1; min-width: 0; }
.stage-name {
  font-size: 13px;
  font-weight: 500;
  color: var(--text);
}
.stage-name.waiting { color: var(--text-3); }
.stage-name.running { color: var(--accent); }
.stage-name.done { color: var(--text-2); }
.stage-tag {
  display: inline-block;
  padding: 1px 6px;
  border-radius: 4px;
  font-size: 10px;
  font-weight: 700;
  font-family: var(--mono);
  background: rgba(255,255,255,0.06);
  color: var(--text-3);
  margin-left: 6px;
  vertical-align: middle;
}
.stage-tag.running { background: var(--accent-dim); color: var(--accent); }
.stage-tag.done    { background: var(--green-dim);  color: var(--green); }

/* ── Pipeline timeline ── */
#pipeline-timeline {
  background: var(--surface);
  border: 1px solid var(--border);
  border-radius: var(--radius);
  padding: 24px;
  margin-bottom: 20px;
  display: -webkit-flex;
  display: flex;
  -webkit-align-items: center;
  align-items: center;
  -webkit-justify-content: center;
  justify-content: center;
  overflow-x: auto;
  -webkit-overflow-scrolling: touch;
}
.tl-node {
  display: -webkit-flex;
  display: flex;
  -webkit-flex-direction: column;
  flex-direction: column;
  -webkit-align-items: center;
  align-items: center;
  row-gap: 8px;
  -webkit-flex-shrink: 0;
  flex-shrink: 0;
}
.tl-node-circle {
  width: 52px;
  height: 52px;
  border-radius: 14px;
  display: -webkit-flex;
  display: flex;
  -webkit-align-items: center;
  align-items: center;
  -webkit-justify-content: center;
  justify-content: center;
  font-family: var(--mono);
  font-size: 13px;
  font-weight: 700;
  -webkit-transition: all 0.4s ease;
  transition: all 0.4s ease;
  /* Safari: promote to own layer to prevent z-index stacking artifacts */
  -webkit-transform: translateZ(0);
  transform: translateZ(0);
  will-change: box-shadow;
}
.tl-node-circle.inactive {
  background: rgba(255,255,255,0.04);
  border: 1.5px solid rgba(255,255,255,0.08);
  color: var(--text-3);
}
.tl-node-circle.active {
  background: var(--accent-dim);
  border: 1.5px solid rgba(0,194,255,0.4);
  color: var(--accent);
  -webkit-box-shadow: 0 0 24px rgba(0,194,255,0.15);
  box-shadow: 0 0 24px rgba(0,194,255,0.15);
}
.tl-node-circle.done {
  background: var(--green-dim);
  border: 1.5px solid rgba(0,229,160,0.4);
  color: var(--green);
  -webkit-box-shadow: 0 0 20px rgba(0,229,160,0.12);
  box-shadow: 0 0 20px rgba(0,229,160,0.12);
}
.tl-node-circle.running {
  background: var(--accent-dim);
  border: 1.5px solid var(--accent);
  color: var(--accent);
  -webkit-animation: ns-tl-pulse 1.2s ease-in-out infinite;
  animation: ns-tl-pulse 1.2s ease-in-out infinite;
}
.tl-node-label {
  font-size: 10px;
  font-weight: 500;
  color: var(--text-3);
  text-align: center;
  letter-spacing: 0.04em;
  text-transform: uppercase;
  -webkit-transition: color 0.3s ease;
  transition: color 0.3s ease;
  /* Prevent label from wrapping oddly on narrow mobile */
  white-space: nowrap;
}
.tl-node-label.active, .tl-node-label.running { color: var(--accent); }
.tl-node-label.done { color: var(--green); }

/* Connector bar between nodes */
.tl-connector {
  width: 40px;
  height: 2px;
  background: var(--border);
  margin: 0 4px;
  margin-bottom: 22px;
  position: relative;
  overflow: hidden;
  -webkit-transition: background 0.4s ease;
  transition: background 0.4s ease;
  -webkit-flex-shrink: 0;
  flex-shrink: 0;
}
.tl-connector.lit { background: rgba(0,229,160,0.35); }
/* Safari: ::after animations on non-positioned parents can be buggy.
   We use a child <span> instead of ::after for the flow animation,
   injected via Python when state=running. See timeline_html(). */
.tl-flow-anim {
  position: absolute;
  top: 0;
  left: -100%;
  width: 100%;
  height: 100%;
  background: -webkit-linear-gradient(left, transparent, var(--accent), transparent);
  background: linear-gradient(90deg, transparent, var(--accent), transparent);
  -webkit-animation: ns-flow 1.2s linear infinite;
  animation: ns-flow 1.2s linear infinite;
  -webkit-animation-fill-mode: none;
  animation-fill-mode: none;
}

/* ── Results panel ── */
#results-panel {
  background: var(--surface);
  border: 1px solid var(--border);
  border-radius: var(--radius);
  overflow: hidden;
  -webkit-transform: translateZ(0);
  transform: translateZ(0);
  margin-bottom: 20px;
}
.results-header {
  padding: 20px 24px;
  border-bottom: 1px solid var(--border);
  display: -webkit-flex;
  display: flex;
  -webkit-align-items: center;
  align-items: center;
  -webkit-column-gap: 12px;
  column-gap: 12px;
}
.results-header-title {
  font-size: 13px;
  font-weight: 600;
  letter-spacing: 0.06em;
  text-transform: uppercase;
  color: var(--text-2);
}

/* Results metrics grid */
.results-grid {
  display: -webkit-flex;
  display: flex;
  border-bottom: 1px solid var(--border);
}
.result-metric {
  -webkit-flex: 1;
  flex: 1;
  background: var(--surface);
  padding: 20px 24px;
  border-right: 1px solid var(--border);
}
.result-metric:last-child { border-right: none; }
.result-metric-label {
  font-size: 10px;
  font-weight: 600;
  letter-spacing: 0.1em;
  text-transform: uppercase;
  color: var(--text-3);
  margin-bottom: 6px;
}
.result-metric-value {
  font-size: 20px;
  font-weight: 700;
  font-family: var(--mono);
  color: var(--text);
  /* Prevent long IDs from overflowing */
  word-break: break-all;
  overflow-wrap: anywhere;
}
.result-metric-value.accent { color: var(--accent); }
.result-metric-value.green  { color: var(--green); }

.confidence-bar {
  margin-top: 8px;
  height: 3px;
  background: rgba(255,255,255,0.07);
  border-radius: 2px;
  overflow: hidden;
}
.confidence-fill {
  height: 100%;
  border-radius: 2px;
  background: -webkit-linear-gradient(left, var(--accent), var(--green));
  background: linear-gradient(90deg, var(--accent), var(--green));
  -webkit-transition: width 0.8s ease;
  transition: width 0.8s ease;
}

/* Downloads */
.downloads-section {
  padding: 20px 24px;
  border-bottom: 1px solid var(--border);
}
.downloads-label {
  font-size: 11px;
  font-weight: 600;
  letter-spacing: 0.1em;
  text-transform: uppercase;
  color: var(--text-3);
  margin-bottom: 12px;
}
.download-cards {
  display: -webkit-flex;
  display: flex;
  -webkit-flex-wrap: wrap;
  flex-wrap: wrap;
  gap: 12px;
  row-gap: 8px;
  -webkit-column-gap: 12px;
  column-gap: 12px;
}
.download-card {
  display: -webkit-flex;
  display: flex;
  -webkit-align-items: center;
  align-items: center;
  -webkit-column-gap: 10px;
  column-gap: 10px;
  padding: 10px 14px;
  background: var(--surface-2);
  border: 1px solid var(--border);
  border-radius: var(--radius-sm);
  cursor: default;
  -webkit-transition: border-color 0.2s ease, -webkit-box-shadow 0.2s ease;
  transition: border-color 0.2s ease, box-shadow 0.2s ease;
}
.download-card:hover {
  border-color: rgba(0,194,255,0.3);
  -webkit-box-shadow: 0 4px 16px rgba(0,0,0,0.3);
  box-shadow: 0 4px 16px rgba(0,0,0,0.3);
}
.download-icon { font-size: 18px; line-height: 1; }
.download-name { font-size: 12px; font-weight: 600; color: var(--text); }
.download-meta { font-size: 10px; color: var(--text-3); font-family: var(--mono); margin-top: 2px; }

/* Literature */
.literature-section { padding: 20px 24px; }
.lit-label {
  font-size: 11px;
  font-weight: 600;
  letter-spacing: 0.1em;
  text-transform: uppercase;
  color: var(--text-3);
  margin-bottom: 12px;
}
.lit-body {
  font-size: 13px;
  color: var(--text-2);
  line-height: 1.7;
  font-family: var(--mono);
  white-space: pre-wrap;
  max-height: 240px;
  overflow-y: auto;
  -webkit-overflow-scrolling: touch;
}

/* File output styling */
#pdf-out-real, #json-out-real {
  margin-top: 0;
}

/* Scrollbar */
::-webkit-scrollbar { width: 6px; height: 6px; }
::-webkit-scrollbar-track { background: transparent; }
::-webkit-scrollbar-thumb { background: var(--border); border-radius: 3px; }
::-webkit-scrollbar-thumb:hover { background: rgba(0,194,255,0.3); }

/* ── Mobile ── */
@media (max-width: 768px) {
  .results-grid {
    -webkit-flex-direction: column;
    flex-direction: column;
  }
  .result-metric {
    border-right: none;
    border-bottom: 1px solid var(--border);
  }
  .result-metric:last-child { border-bottom: none; }
  #scenario-panel {
    -webkit-flex-direction: column;
    flex-direction: column;
    -webkit-align-items: flex-start;
    align-items: flex-start;
  }
  .scenario-divider { display: none; }
  .tl-connector { width: 24px; }
  .tl-node-circle { width: 44px; height: 44px; font-size: 11px; }
}

/* ── Reduce motion ── */
@media (prefers-reduced-motion: reduce) {
  *, *::before, *::after {
    -webkit-animation-duration: 0.01ms !important;
    animation-duration: 0.01ms !important;
    -webkit-transition-duration: 0.01ms !important;
    transition-duration: 0.01ms !important;
  }
}
"""

# ─────────────────────────────────────────────
# HTML helpers
# ─────────────────────────────────────────────

def header_html():
    return """
<div id="ns-header">
  <div class="ns-wordmark">NeuroSight</div>
  <div class="ns-title">Longitudinal Intelligence</div>
  <div class="ns-subtitle">Neuro-Oncology AI&nbsp;&nbsp;·&nbsp;&nbsp;MRI Segmentation&nbsp;&nbsp;·&nbsp;&nbsp;PINN Modeling&nbsp;&nbsp;·&nbsp;&nbsp;cfDNA Analysis</div>
</div>
"""

def scenario_panel_html(baseline=False, followup=False, cfdna=False):
    b_class = "active" if baseline else "inactive"
    f_class = "active" if followup else "inactive"
    c_class = "active" if cfdna else "inactive"
    b_icon = "&#10003;" if baseline else "&#10007;"
    f_icon = "&#10003;" if followup else "&#10007;"
    c_icon = "&#10003;" if cfdna else "&#10007;"

    if baseline and followup and cfdna:
        scenario = "Multimodal Longitudinal Analysis"
        modules = ["M1", "M2", "M3", "M4", "M5"]
    elif baseline and followup:
        scenario = "Longitudinal MRI Analysis"
        modules = ["M1", "M2", "M3", "M4"]
    elif baseline and cfdna:
        scenario = "MRI + Liquid Biopsy"
        modules = ["M1", "M2", "M4", "M5"]
    elif baseline:
        scenario = "Single MRI Analysis"
        modules = ["M1", "M2", "M4"]
    elif cfdna:
        scenario = "Liquid Biopsy Analysis"
        modules = ["M5"]
    else:
        scenario = None
        modules = []

    scenario_name_class = "active" if scenario else ""
    scenario_text = scenario if scenario else "Upload files above to detect scenario"

    chain_parts = []
    for i, m in enumerate(modules):
        chain_parts.append(f'<span class="module-node-sm">{m}</span>')
        if i < len(modules) - 1:
            chain_parts.append('<span class="module-arrow-sm">&#8594;</span>')
    chain_html = "".join(chain_parts)

    return f"""
<div id="scenario-panel">
  <div class="scenario-inputs-status">
    <span class="input-chip {b_class}">{b_icon}&nbsp;Baseline MRI</span>
    <span class="input-chip {f_class}">{f_icon}&nbsp;Follow-up MRI</span>
    <span class="input-chip {c_class}">{c_icon}&nbsp;cfDNA Biomarkers</span>
  </div>
  <div class="scenario-divider"></div>
  <div class="scenario-info">
    <div class="scenario-label">Detected Scenario</div>
    <div class="scenario-name {scenario_name_class}">{scenario_text}</div>
    <div class="module-chain">{chain_html}</div>
  </div>
</div>
"""

ALL_STAGES = {
    "M1": ("MRI Preprocessing &amp; Segmentation", "M1"),
    "M2": ("PINN Tumor Growth Modeling",            "M2"),
    "M3": ("Longitudinal Progression Analysis",     "M3"),
    "M4": ("Literature Retrieval &amp; Report Synthesis", "M4"),
    "M5": ("cfDNA Biomarker Classification",        "M5"),
}

def exec_monitor_html(stages, stage_states, dot_class=""):
    dot_cls = dot_class or "idle"
    rows = ""
    for mid in stages:
        label, tag = ALL_STAGES[mid]
        state = stage_states.get(mid, "waiting")
        icon_map = {"waiting": "&middot;", "running": "&#9673;", "done": "&#10003;", "error": "&#10007;"}
        icon = icon_map.get(state, "&middot;")
        rows += f"""
<div class="stage-row">
  <div class="stage-icon {state}">{icon}</div>
  <div class="stage-text">
    <span class="stage-name {state}">{label}</span>
    <span class="stage-tag {state}">{tag}</span>
  </div>
</div>"""

    return f"""
<div id="exec-monitor">
  <div class="exec-monitor-header">
    <span class="exec-monitor-title">Execution Monitor</span>
    <span class="exec-status-dot {dot_cls}"></span>
  </div>
  <div class="stage-list">{rows}</div>
</div>
"""

def timeline_html(stages, stage_states):
    label_map = {"M1": "MRI", "M2": "PINN", "M3": "Trends", "M4": "Literature", "M5": "cfDNA"}
    nodes = ""
    for i, mid in enumerate(stages):
        state = stage_states.get(mid, "inactive")
        nodes += f"""
<div class="tl-node">
  <div class="tl-node-circle {state}">{mid}</div>
  <div class="tl-node-label {state}">{label_map.get(mid, mid)}</div>
</div>"""
        if i < len(stages) - 1:
            prev_state = stage_states.get(mid, "inactive")
            lit_cls   = "lit"     if prev_state == "done"    else ""
            # Use inline child span instead of ::after pseudo for Safari compat
            flow_span = '<span class="tl-flow-anim"></span>' if prev_state == "running" else ""
            nodes += f'<div class="tl-connector {lit_cls}">{flow_span}</div>'

    return f'<div id="pipeline-timeline">{nodes}</div>'

def results_html(case_id, signal, confidence, report_text=""):
    conf_pct = int((confidence or 0) * 100)
    lit_section = ""
    if report_text:
        import html as html_mod
        safe = html_mod.escape(report_text[:800])
        lit_section = f"""
<div class="literature-section">
  <div class="lit-label">Literature &amp; Report Summary</div>
  <div class="lit-body">{safe}</div>
</div>"""

    return f"""
<div id="results-panel">
  <div class="results-header">
    <span style="font-size:16px;line-height:1;">&#10022;</span>
    <span class="results-header-title">Analysis Complete</span>
  </div>
  <div class="results-grid">
    <div class="result-metric">
      <div class="result-metric-label">Case ID</div>
      <div class="result-metric-value accent" style="font-size:14px;">{case_id}</div>
    </div>
    <div class="result-metric">
      <div class="result-metric-label">Dominant Signal</div>
      <div class="result-metric-value" style="font-size:15px;">{signal or "&#8212;"}</div>
    </div>
    <div class="result-metric">
      <div class="result-metric-label">Confidence</div>
      <div class="result-metric-value green">{conf_pct}%</div>
      <div class="confidence-bar">
        <div class="confidence-fill" style="width:{conf_pct}%"></div>
      </div>
    </div>
  </div>
  <div class="downloads-section">
    <div class="downloads-label">Downloads</div>
    <div class="download-cards">
      <div class="download-card">
        <span class="download-icon">&#128196;</span>
        <div>
          <div class="download-name">NeuroSight Clinical Report</div>
          <div class="download-meta">PDF &middot; Available below</div>
        </div>
      </div>
      <div class="download-card">
        <span class="download-icon">&#128202;</span>
        <div>
          <div class="download-name">Agent Payload</div>
          <div class="download-meta">JSON &middot; Available below</div>
        </div>
      </div>
    </div>
  </div>
  {lit_section}
</div>
"""

# ─────────────────────────────────────────────
# Backend helpers
# ─────────────────────────────────────────────

def generate_pdf(report_text: str, output_path: str) -> str:
    pdf = FPDF()
    pdf.add_page()
    pdf.set_font("Arial", size=11)
    if not report_text:
        report_text = "No report generated."
    for line in report_text.split("\n"):
        pdf.multi_cell(0, 8, txt=line.encode("latin-1", "replace").decode("latin-1"))
    pdf.output(output_path)
    return output_path

def format_agent_json(pipeline_result: dict) -> dict:
    payload = pipeline_result.get("payload", {})
    delta   = payload.get("delta", {})
    scan2   = payload.get("scan2", {})
    m1      = scan2.get("m1_outputs", {})
    m2      = scan2.get("m2_outputs", {})
    m3      = scan2.get("m3_outputs", {})
    return {
        "case_id":        payload.get("patient_id", "UNKNOWN"),
        "generated_at":   datetime.now().isoformat(),
        "scan_index":     4,
        "date":           datetime.now().strftime("%Y-%m-%d"),
        "dominant_signal":m3.get("progression_class", "Unknown"),
        "confidence":     m3.get("confidence", 0.0),
        "key_parameters": {
            "mu_R":       m2.get("mu_R", 0.0),
            "mu_D":       m2.get("mu_D", 0.0),
            "gamma":      m2.get("gamma", 0.0),
            "delta_mu_R": delta.get("delta_mu_R", 0.0),
            "delta_mu_D": delta.get("delta_mu_D", 0.0),
            "delta_gamma":delta.get("delta_gamma", 0.0),
        },
        "volumes": {
            "enhancing_core_cc": m1.get("enhancing_core_volume_cc", 0.0),
            "edema_cc":          m1.get("edema_volume_cc", 0.0),
            "necrotic_core_cc":  m1.get("necrotic_core_volume_cc", 0.0),
            "pinn_predicted_cc": 0.0,
            "pinn_discrepancy_cc": 0.0,
        },
        "molecular_profile":    "IDH-wildtype, MGMT unmethylated. Post-surgical baseline.",
        "treatment_at_scan":    "TMZ + RT (concurrent phase)",
        "weeks_post_RT":        0,
        "TMZ_cycles_completed": 0,
        "agent_search_context": f"true-progression, {str(m3.get('progression_class', '')).lower()}",
        "m3_classification_history": [{
            "scan":          4,
            "date":          datetime.now().strftime("%Y-%m-%d"),
            "classification":m3.get("progression_class", "Unknown"),
            "confidence":    m3.get("confidence", 0.0),
            "delta_mu_R":    delta.get("delta_mu_R", 0.0),
            "delta_mu_D":    delta.get("delta_mu_D", 0.0),
            "delta_gamma":   delta.get("delta_gamma", 0.0),
        }],
        "full_report_path": "reports/full_report.md",
    }

def get_scenario(has_b: bool, has_f: bool, has_c: bool):
    """
    Three real scenarios — M5 is an add-on, not a scenario branch.

    Inputs → scenario_key, ordered stages
    ─────────────────────────────────────────────────────────────
    Only cfDNA                     → "cfdna",        [M5]
    1 MRI (± cfDNA)                → "single",       [M1,M2,M4] (+ M5)
    Baseline + Follow-up (± cfDNA) → "longitudinal", [M1,M2,M3,M4] (+ M5)
    Follow-up without baseline     → None  (invalid)
    Nothing                        → None
    """
    if not has_b and not has_f and not has_c:
        return None, []

    if has_f and not has_b:
        return None, []          # follow-up without baseline is invalid

    if not has_b and has_c:
        return "cfdna", ["M5"]

    if has_b and not has_f:
        stages = ["M1", "M2", "M4"]
        if has_c:
            stages.append("M5")
        return "single", stages

    # has_b and has_f
    stages = ["M1", "M2", "M3", "M4"]
    if has_c:
        stages.append("M5")
    return "longitudinal", stages

# ─────────────────────────────────────────────
# Run function — generator for streaming updates
# ─────────────────────────────────────────────


def run_neurosight(baseline_zip, followup_zip, cfdna_json, progress=gr.Progress()):
    has_b = baseline_zip is not None
    has_f = followup_zip is not None
    has_c = cfdna_json   is not None

    scenario_key, stages = get_scenario(has_b, has_f, has_c)

    if scenario_key is None:
        if has_f and not has_b:
            raise gr.Error("Please upload the baseline MRI scan — a follow-up scan alone is not valid.")
        raise gr.Error("Upload at least one input file to begin analysis.")

    states = {s: "waiting" for s in stages}

    def render(dot="running", result_html="", pdf_path=None, json_path=None):
        return (
            exec_monitor_html(stages, states, dot),
            timeline_html(stages, states),
            result_html,
            pdf_path,
            json_path,
        )

    yield render("running")
    time.sleep(0.3)

    session_dir = None
    result      = None
    m5_result   = None
    errors      = []
    agent       = {}

    try:
        # ── Scenario 1: cfDNA only ──────────────────────────────────────────
        if scenario_key == "cfdna":
            states["M5"] = "running"
            yield render("running")
            result = pipeline.run_cfdna(cfdna_json_path=cfdna_json.name)
            states["M5"] = "error" if result.get("errors") else "done"
            errors      = result.get("errors", [])
            session_dir = "sessions/cfdna_run"
            os.makedirs(session_dir, exist_ok=True)
            m5 = result.get("m5", {})
            agent = {
                "case_id":         result.get("sample_id", "UNKNOWN"),
                "dominant_signal": m5.get("predicted_label", "—"),
                "confidence":      m5.get("confidence", 0.0),
            }

        # ── Scenario 2: Single MRI (± cfDNA) ───────────────────────────────
        elif scenario_key == "single":
            states["M1"] = states["M2"] = "running"
            yield render("running")
            result  = pipeline.run_single_scan(zip_path=baseline_zip.name)
            errors  = result.get("errors", [])
            session_dir = result.get("session_dir", ".")
            for mid in ("M1", "M2"):
                states[mid] = "error" if any(
                    e.get("stage", "").lower().startswith(mid.lower()) for e in errors
                ) else "done"
            yield render("running")

            states["M4"] = "running"
            yield render("running")
            states["M4"] = "done"
            yield render("running")

            if has_c:
                states["M5"] = "running"
                yield render("running")
                m5_result    = pipeline.run_cfdna(cfdna_json_path=cfdna_json.name)
                states["M5"] = "error" if m5_result.get("errors") else "done"
                errors      += m5_result.get("errors", [])
                yield render("running")

            payload = result.get("payload", {})
            scan1   = payload.get("scan1", {})
            agent = {
                "case_id":         result.get("study_id", "UNKNOWN"),
                "dominant_signal": (
                    m5_result.get("m5", {}).get("predicted_label", "—")
                    if m5_result else "N/A (no cfDNA)"
                ),
                "confidence":      scan1.get("m1_outputs", {}).get("segmentation_confidence", 0.0),
            }

        # ── Scenario 3: Longitudinal MRI (± cfDNA) ─────────────────────────
        else:
            states["M1"] = states["M2"] = "running"
            yield render("running")
            result = pipeline.run_longitudinal(
                baseline_zip=baseline_zip.name,
                followup_zip=followup_zip.name,
            )
            errors      = result.get("errors", [])
            session_dir = result.get("session_dir", ".")
            for mid in ("M1", "M2"):
                states[mid] = "error" if any(
                    mid.lower() in e.get("stage", "").lower() for e in errors
                ) else "done"
            yield render("running")

            states["M3"] = "running"
            yield render("running")
            states["M3"] = "error" if result.get("m3", {}).get("error") else "done"
            yield render("running")

            states["M4"] = "running"
            yield render("running")
            states["M4"] = "done"
            yield render("running")

            if has_c:
                states["M5"] = "running"
                yield render("running")
                m5_result    = pipeline.run_cfdna(cfdna_json_path=cfdna_json.name)
                states["M5"] = "error" if m5_result.get("errors") else "done"
                errors      += m5_result.get("errors", [])
                yield render("running")

            payload = result.get("payload", {})
            scan2   = payload.get("scan2", {})
            m3_out  = scan2.get("m3_outputs", result.get("m3", {}))
            agent = {
                "case_id":         result.get("study_id", "UNKNOWN"),
                "dominant_signal": m3_out.get("progression_class", "—"),
                "confidence":      m3_out.get("confidence", 0.0),
            }

    except gr.Error:
        raise                                       # don't double-wrap gr.Error
    except Exception as exc:
        raise gr.Error(f"Pipeline error: {exc}")

    # ── Write outputs ───────────────────────────────────────────────────────
    os.makedirs(session_dir or ".", exist_ok=True)
    case_id   = agent.get("case_id", "UNKNOWN")
    pdf_path  = os.path.join(session_dir, f"{case_id}_report.pdf")
    json_path = os.path.join(session_dir, f"{case_id}_study.json")

    report_text = result.get("report", "") if result else ""
    generate_pdf(report_text or "NeuroSight Report", pdf_path)

    try:
        agent_payload = format_agent_json(result) if result and scenario_key != "cfdna" else agent
    except Exception:
        agent_payload = agent

    with open(json_path, "w") as fh:
        json.dump(agent_payload, fh, indent=2)

    res_html = results_html(
        case_id=agent.get("case_id", "—"),
        signal=agent.get("dominant_signal", "—"),
        confidence=agent.get("confidence", 0.0),
        report_text=report_text or "",
    )

    yield render("done", res_html, pdf_path, json_path)

# ─────────────────────────────────────────────
# Gradio UI
# ─────────────────────────────────────────────

with gr.Blocks(css=CSS, title="NeuroSight — Neuro-Oncology Intelligence Engine") as app:

    gr.HTML(header_html())

    # Upload zone — three columns via flex (Safari-safe)
    with gr.Row(elem_id="upload-zone"):
        with gr.Column(elem_classes="upload-card"):
            gr.HTML('<div class="upload-card-label"><span class="badge">B</span>Baseline MRI</div>')
            baseline_input = gr.File(label="", file_types=[".zip"], show_label=False)

        with gr.Column(elem_classes="upload-card"):
            gr.HTML('<div class="upload-card-label"><span class="badge">F</span>Follow-up MRI</div>')
            followup_input = gr.File(label="", file_types=[".zip"], show_label=False)

        with gr.Column(elem_classes="upload-card"):
            gr.HTML('<div class="upload-card-label"><span class="badge">C</span>cfDNA Biomarkers</div>')
            cfdna_input = gr.File(label="", file_types=[".json"], show_label=False)

    scenario_display = gr.HTML(scenario_panel_html())

    run_btn = gr.Button("Run NeuroSight Analysis", elem_id="run-btn", variant="primary")

    monitor_display  = gr.HTML(exec_monitor_html([], {}, ""))
    timeline_display = gr.HTML("")
    results_display  = gr.HTML("")

    with gr.Row():
        pdf_output  = gr.File(label="NeuroSight Clinical Report (.pdf)", elem_id="pdf-out-real")
        json_output = gr.File(label="Agent Payload (.json)",             elem_id="json-out-real")

    # Reactive scenario panel — fires on any file change
    def update_scenario(b, f, c):
        return scenario_panel_html(b is not None, f is not None, c is not None)

    for inp in [baseline_input, followup_input, cfdna_input]:
        inp.change(
            fn=update_scenario,
            inputs=[baseline_input, followup_input, cfdna_input],
            outputs=[scenario_display],
        )

    run_btn.click(
        fn=run_neurosight,
        inputs=[baseline_input, followup_input, cfdna_input],
        outputs=[monitor_display, timeline_display, results_display, pdf_output, json_output],
    )

if __name__ == "__main__":
    app.launch(server_name="0.0.0.0")