#!/usr/bin/env bash
set -euo pipefail

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
LOG_DIR="$ROOT_DIR/.logs"
LOG_FILE="$LOG_DIR/jupyterlab.log"

usage() {
  cat <<'EOF'
Usage: bash ./scripts/app.sh <command> [options]

Consumer-facing control surface for QuantumLearning.

Commands:
  bootstrap [--skip-browser] [--run-validation]
      Create or refresh the repo-local environment.

  start [--port PORT] [--open] [--foreground] [--restart]
      Start the repo-local Jupyter server.
      Default mode is background service mode.

  stop
      Stop the repo-local Jupyter server for this project.

  reset-state
      Remove repo-local Jupyter runtime and saved UI state files.

  restart [--port PORT] [--open] [--foreground]
      Restart the repo-local Jupyter server.

  status
      Print the current project and Jupyter status.

  url
      Print the current Jupyter URL for this project.

  open
      Open the current Jupyter URL in the default browser.

  logs [-f|--follow]
      Show the Jupyter log file. Use -f to follow live output.

  validate [--quick|--standard|--full]
      Run the validation suite.

Examples:
  bash ./scripts/app.sh bootstrap
  bash ./scripts/app.sh start --open
  bash ./scripts/app.sh restart
  bash ./scripts/app.sh status
  bash ./scripts/app.sh stop
  bash ./scripts/app.sh reset-state
  bash ./scripts/app.sh validate --quick
EOF
}

die() {
  echo "$*" >&2
  exit 1
}

ensure_dirs() {
  mkdir -p \
    "$ROOT_DIR/.cache/matplotlib" \
    "$ROOT_DIR/.ipython" \
    "$ROOT_DIR/.jupyter_config" \
    "$ROOT_DIR/.jupyter_data" \
    "$ROOT_DIR/.jupyter_runtime" \
    "$LOG_DIR"
}

ensure_venv() {
  if [[ ! -x "$ROOT_DIR/.venv/bin/python" ]]; then
    die "Missing repo-local .venv. Run: bash ./scripts/app.sh bootstrap"
  fi
}

set_jupyter_env() {
  export PATH="$ROOT_DIR/.venv/bin:/usr/bin:/bin:/usr/sbin:/sbin"
  export JUPYTER_CONFIG_DIR="$ROOT_DIR/.jupyter_config"
  export JUPYTER_DATA_DIR="$ROOT_DIR/.jupyter_data"
  export JUPYTER_RUNTIME_DIR="$ROOT_DIR/.jupyter_runtime"
  export IPYTHONDIR="$ROOT_DIR/.ipython"
  export MPLCONFIGDIR="$ROOT_DIR/.cache/matplotlib"
}

ensure_kernel() {
  local kernel_dir="$ROOT_DIR/.venv/share/jupyter/kernels/quantum-learning"
  if [[ ! -f "$kernel_dir/kernel.json" ]]; then
    "$ROOT_DIR/.venv/bin/python" -m ipykernel install \
      --sys-prefix \
      --name quantum-learning \
      --display-name "QuantumLearning (.venv)" \
      --env IPYTHONDIR "$IPYTHONDIR" \
      --env MPLCONFIGDIR "$MPLCONFIGDIR"
  fi
}

server_json_for_root() {
  "$ROOT_DIR/.venv/bin/python" - <<'PY' "$ROOT_DIR"
from __future__ import annotations

import json
import sys
from pathlib import Path

root = Path(sys.argv[1]).resolve()
runtime_dir = root / ".jupyter_runtime"
matches = []
for path in runtime_dir.glob("jpserver-*.json"):
    try:
        payload = json.loads(path.read_text())
    except Exception:
        continue
    if Path(payload.get("root_dir", "")).resolve() == root:
        matches.append((path.stat().st_mtime, path))

if matches:
    print(max(matches)[1])
PY
}

server_json_for_pid() {
  local pid="$1"
  "$ROOT_DIR/.venv/bin/python" - <<'PY' "$ROOT_DIR" "$pid"
from __future__ import annotations

import json
import sys
from pathlib import Path

root = Path(sys.argv[1]).resolve()
target_pid = int(sys.argv[2])
runtime_dir = root / ".jupyter_runtime"
matches = []
for path in runtime_dir.glob("jpserver-*.json"):
    try:
        payload = json.loads(path.read_text())
    except Exception:
        continue
    if Path(payload.get("root_dir", "")).resolve() == root and int(payload.get("pid", -1)) == target_pid:
        matches.append((path.stat().st_mtime, path))

if matches:
    print(max(matches)[1])
PY
}

server_field() {
  local json_path="$1"
  local field_name="$2"
  "$ROOT_DIR/.venv/bin/python" - <<'PY' "$json_path" "$field_name"
from __future__ import annotations

import json
import sys
from pathlib import Path

path = Path(sys.argv[1])
field_name = sys.argv[2]
payload = json.loads(path.read_text())
value = payload[field_name]
print(value)
PY
}

server_url_from_json() {
  local json_path="$1"
  "$ROOT_DIR/.venv/bin/python" - <<'PY' "$json_path"
from __future__ import annotations

import json
import sys
from pathlib import Path

path = Path(sys.argv[1])
payload = json.loads(path.read_text())
print(f"{payload['url']}lab/tree/notebooks/START_HERE.ipynb?token={payload['token']}")
PY
}

cleanup_stale_runtime_files() {
  "$ROOT_DIR/.venv/bin/python" - <<'PY' "$ROOT_DIR" "${1:-0}"
from __future__ import annotations

import json
import os
import sys
from pathlib import Path

root = Path(sys.argv[1]).resolve()
verbose = sys.argv[2] == "1"
runtime_dir = root / ".jupyter_runtime"
removed = []

for path in sorted(runtime_dir.glob("jpserver-*.json")):
    try:
        payload = json.loads(path.read_text())
    except Exception:
        continue

    if Path(payload.get("root_dir", "")).resolve() != root:
        continue

    pid = int(payload.get("pid", -1))
    alive = False
    if pid > 0:
        try:
            os.kill(pid, 0)
        except ProcessLookupError:
            alive = False
        except PermissionError:
            alive = True
        else:
            alive = True

    if alive:
        continue

    for candidate in (path, path.with_name(f"{path.stem}-open.html")):
        if candidate.exists():
            candidate.unlink()
            removed.append(candidate)

if verbose:
    for path in removed:
        print(path)
PY
}

purge_runtime_artifacts() {
  "$ROOT_DIR/.venv/bin/python" - <<'PY' "$ROOT_DIR" "${1:-0}"
from __future__ import annotations

import sys
from pathlib import Path

root = Path(sys.argv[1]).resolve()
verbose = sys.argv[2] == "1"
runtime_dir = root / ".jupyter_runtime"
patterns = (
    "jpserver-*.json",
    "jpserver-*-open.html",
    "nbserver-*.json",
    "nbserver-*-open.html",
    "kernel-*.json",
)
removed = []

for pattern in patterns:
    for path in sorted(runtime_dir.glob(pattern)):
        if path.exists():
            path.unlink()
            removed.append(path)

if verbose:
    for path in removed:
        print(path)
PY
}

purge_saved_ui_state() {
  "$ROOT_DIR/.venv/bin/python" - <<'PY' "$ROOT_DIR" "${1:-0}"
from __future__ import annotations

import sys
from pathlib import Path

root = Path(sys.argv[1]).resolve()
verbose = sys.argv[2] == "1"
data_dir = root / ".jupyter_data"
paths = [
    data_dir / "nbsignatures.db",
    data_dir / "notebook_secret",
]
workspace_dir = data_dir / "lab" / "workspaces"
removed = []

for path in paths:
    if path.exists():
        path.unlink()
        removed.append(path)

for path in sorted(workspace_dir.glob("*.jupyterlab-workspace")):
    if path.exists():
        path.unlink()
        removed.append(path)

if verbose:
    for path in removed:
        print(path)
PY
}

print_ui_state_summary() {
  local workspace_count="0"
  if compgen -G "$ROOT_DIR/.jupyter_data/lab/workspaces/*.jupyterlab-workspace" >/dev/null; then
    workspace_count="$(find "$ROOT_DIR/.jupyter_data/lab/workspaces" -maxdepth 1 -name '*.jupyterlab-workspace' | wc -l | tr -d ' ')"
  fi

  local trust_db="absent"
  local notebook_secret="absent"
  [[ -f "$ROOT_DIR/.jupyter_data/nbsignatures.db" ]] && trust_db="present"
  [[ -f "$ROOT_DIR/.jupyter_data/notebook_secret" ]] && notebook_secret="present"

  if [[ "$workspace_count" == "0" && "$trust_db" == "absent" && "$notebook_secret" == "absent" ]]; then
    echo "Jupyter UI state: clean"
    return 0
  fi

  echo "Jupyter UI state: workspaces=$workspace_count, trust-db=$trust_db, notebook-secret=$notebook_secret"
  echo "  Clean slate: bash ./scripts/app.sh reset-state"
}

have_running_server() {
  local json_path
  cleanup_stale_runtime_files >/dev/null
  json_path="$(server_json_for_root || true)"
  if [[ -z "$json_path" ]]; then
    return 1
  fi

  local pid
  pid="$(server_field "$json_path" "pid")"
  kill -0 "$pid" >/dev/null 2>&1
}

print_running_server() {
  local json_path url pid
  cleanup_stale_runtime_files >/dev/null
  json_path="$(server_json_for_root || true)"
  if [[ -z "$json_path" ]]; then
    die "No repo-local Jupyter runtime file found. Run: bash ./scripts/app.sh start"
  fi
  url="$(server_url_from_json "$json_path")"
  pid="$(server_field "$json_path" "pid")"
  echo "Jupyter URL: $url"
  echo "PID: $pid"
  echo "Runtime file: $json_path"
  echo "Log file: $LOG_FILE"
}

stop_server() {
  ensure_venv
  ensure_dirs

  local json_path pid removed_paths
  cleanup_stale_runtime_files >/dev/null
  json_path="$(server_json_for_root || true)"
  if [[ -z "$json_path" ]]; then
    echo "No repo-local Jupyter runtime file found."
    removed_paths="$(purge_runtime_artifacts 1 || true)"
    if [[ -n "$removed_paths" ]]; then
      echo "Removed stale runtime files:"
      echo "$removed_paths"
    fi
    return 0
  fi

  pid="$(server_field "$json_path" "pid")"
  if ! kill -0 "$pid" >/dev/null 2>&1; then
    removed_paths="$(cleanup_stale_runtime_files 1 || true)"
    local extra_removed
    extra_removed="$(purge_runtime_artifacts 1 || true)"
    echo "Jupyter runtime file exists, but the process is not running."
    echo "Runtime file: $json_path"
    if [[ -n "$removed_paths" ]]; then
      echo "Removed stale runtime files:"
      echo "$removed_paths"
    fi
    if [[ -n "$extra_removed" ]]; then
      echo "Removed leftover runtime artifacts:"
      echo "$extra_removed"
    fi
    return 0
  fi

  kill -TERM "$pid"
  for _ in $(seq 1 30); do
    if ! kill -0 "$pid" >/dev/null 2>&1; then
      cleanup_stale_runtime_files >/dev/null
      purge_runtime_artifacts >/dev/null
      echo "Stopped Jupyter (pid $pid)."
      return 0
    fi
    sleep 1
  done

  die "Jupyter did not stop cleanly after 30 seconds. Inspect: $LOG_FILE"
}

start_server() {
  local port="$1"
  local foreground="$2"
  local open_browser="$3"
  local restart="$4"

  ensure_venv
  ensure_dirs
  set_jupyter_env
  ensure_kernel
  cleanup_stale_runtime_files >/dev/null

  if have_running_server; then
    if [[ "$restart" == "1" ]]; then
      stop_server
    else
      echo "Jupyter is already running for this project."
      print_running_server
      if [[ "$open_browser" == "1" ]]; then
        local existing_url
        existing_url="$(server_url_from_json "$(server_json_for_root)")"
        open "$existing_url" >/dev/null 2>&1 || true
      fi
      return 0
    fi
  fi

  purge_runtime_artifacts >/dev/null

  if [[ "$foreground" == "1" ]]; then
    exec "$ROOT_DIR/.venv/bin/jupyter" lab \
      --no-browser \
      --ip=127.0.0.1 \
      --port="$port"
  fi

  : > "$LOG_FILE"
  nohup "$ROOT_DIR/.venv/bin/jupyter" lab \
    --no-browser \
    --ip=127.0.0.1 \
    --port="$port" \
    < /dev/null \
    >>"$LOG_FILE" 2>&1 &
  local launch_pid=$!

  local json_path=""
  for _ in $(seq 1 60); do
    if ! kill -0 "$launch_pid" >/dev/null 2>&1; then
      echo "Jupyter exited during startup. Recent log output:" >&2
      tail -n 40 "$LOG_FILE" >&2 || true
      exit 1
    fi

    json_path="$(server_json_for_pid "$launch_pid" || true)"
    if [[ -n "$json_path" ]]; then
      local url
      url="$(server_url_from_json "$json_path")"
      echo "Jupyter started."
      echo "URL: $url"
      echo "PID: $launch_pid"
      echo "Log: $LOG_FILE"
      if [[ "$open_browser" == "1" ]]; then
        open "$url" >/dev/null 2>&1 || true
      fi
      return 0
    fi
    sleep 1
  done

  die "Timed out waiting for Jupyter startup. Inspect: $LOG_FILE"
}

show_logs() {
  ensure_dirs
  if [[ ! -f "$LOG_FILE" ]]; then
    echo "No Jupyter log file exists yet."
    return 0
  fi

  if [[ "${1:-}" == "--follow" || "${1:-}" == "-f" ]]; then
    tail -f "$LOG_FILE"
  else
    tail -n 80 "$LOG_FILE"
  fi
}

validate_command() {
  bash "$ROOT_DIR/scripts/run_validation.sh" "$@"
}

reset_state_command() {
  ensure_venv
  ensure_dirs

  if have_running_server; then
    stop_server >/dev/null
  else
    cleanup_stale_runtime_files >/dev/null
    purge_runtime_artifacts >/dev/null
  fi

  local removed_runtime removed_ui
  removed_runtime="$(purge_runtime_artifacts 1 || true)"
  removed_ui="$(purge_saved_ui_state 1 || true)"

  if [[ -z "$removed_runtime" && -z "$removed_ui" ]]; then
    echo "Jupyter state was already clean."
    return 0
  fi

  if [[ -n "$removed_runtime" ]]; then
    echo "Removed runtime artifacts:"
    echo "$removed_runtime"
  fi

  if [[ -n "$removed_ui" ]]; then
    echo "Removed saved UI state:"
    echo "$removed_ui"
  fi
}

status_command() {
  cd "$ROOT_DIR"
  ensure_dirs

  local branch commit git_state
  branch="$(git branch --show-current)"
  commit="$(git rev-parse --short HEAD)"
  if [[ -n "$(git status --short)" ]]; then
    git_state="dirty"
  else
    git_state="clean"
  fi

  echo "Root: $ROOT_DIR"
  echo "Git: $branch @ $commit ($git_state)"

  if [[ -x "$ROOT_DIR/.venv/bin/python" ]]; then
    local python_version removed_runtime json_path pid url
    python_version="$("$ROOT_DIR/.venv/bin/python" -c 'import sys; print(sys.version.split()[0])')"
    echo "Venv: present (Python $python_version)"
    cleanup_stale_runtime_files >/dev/null
    if ! have_running_server; then
      removed_runtime="$(purge_runtime_artifacts 1 || true)"
    else
      removed_runtime=""
    fi
    json_path="$(server_json_for_root || true)"

    if compgen -G "$ROOT_DIR/.playwright-browsers/chromium-*" >/dev/null; then
      echo "Browser runtime: installed"
    else
      echo "Browser runtime: missing"
    fi

    if [[ -f "$LOG_FILE" ]]; then
      echo "Jupyter log: $LOG_FILE"
    else
      echo "Jupyter log: none yet"
    fi

    if [[ -n "$json_path" ]]; then
      pid="$(server_field "$json_path" "pid")"
      url="$(server_url_from_json "$json_path")"
      if kill -0 "$pid" >/dev/null 2>&1; then
        echo "Jupyter: running (pid $pid)"
        echo "  URL: $url"
        echo "  Runtime file: $json_path"
      else
        echo "Jupyter: runtime file exists but process is not running"
        echo "  Runtime file: $json_path"
      fi
    else
      echo "Jupyter: no repo-local runtime file found"
      if [[ -n "${removed_runtime:-}" ]]; then
        echo "  Purged stale runtime artifacts:"
        echo "$removed_runtime"
      fi
    fi

    print_ui_state_summary
  else
    echo "Venv: missing"
    if compgen -G "$ROOT_DIR/.playwright-browsers/chromium-*" >/dev/null; then
      echo "Browser runtime: installed"
    else
      echo "Browser runtime: missing"
    fi
    if [[ -f "$LOG_FILE" ]]; then
      echo "Jupyter log: $LOG_FILE"
    else
      echo "Jupyter log: none yet"
    fi
    echo "Jupyter: repo-local environment is missing"
    print_ui_state_summary
  fi

  echo "Consumer:"
  echo "  Start:   bash ./scripts/app.sh start --open"
  echo "  Stop:    bash ./scripts/app.sh stop"
  echo "  Restart: bash ./scripts/app.sh restart"
  echo "  Reset:   bash ./scripts/app.sh reset-state"
  echo "  Logs:    bash ./scripts/app.sh logs -f"
  echo "Validation:"
  echo "  Quick: bash ./scripts/app.sh validate --quick"
  echo "  Full:  bash ./scripts/app.sh validate --full"
}

start_command() {
  local port="8888"
  local foreground="0"
  local open_browser="0"
  local restart="0"

  while [[ $# -gt 0 ]]; do
    case "$1" in
      --port)
        [[ $# -ge 2 ]] || die "--port requires a value"
        port="$2"
        shift 2
        ;;
      --open)
        open_browser="1"
        shift
        ;;
      --foreground)
        foreground="1"
        shift
        ;;
      --restart)
        restart="1"
        shift
        ;;
      -h|--help)
        cat <<'EOF'
Usage: bash ./scripts/app.sh start [--port PORT] [--open] [--foreground] [--restart]
EOF
        return 0
        ;;
      *)
        die "Unknown start option: $1"
        ;;
    esac
  done

  [[ "$port" =~ ^[0-9]+$ ]] || die "Port must be an integer, got: $port"
  start_server "$port" "$foreground" "$open_browser" "$restart"
}

restart_command() {
  start_command --restart "$@"
}

if [[ $# -eq 0 ]]; then
  usage
  exit 1
fi

COMMAND="$1"
shift

case "$COMMAND" in
  bootstrap)
    exec bash "$ROOT_DIR/scripts/bootstrap_mac.sh" "$@"
    ;;
  start)
    start_command "$@"
    ;;
  stop)
    stop_server
    ;;
  reset-state)
    reset_state_command "$@"
    ;;
  restart)
    restart_command "$@"
    ;;
  status)
    status_command
    ;;
  _status_internal)
    status_command
    ;;
  url)
    ensure_venv
    cleanup_stale_runtime_files >/dev/null
    json_path="$(server_json_for_root || true)"
    [[ -n "$json_path" ]] || die "No repo-local Jupyter runtime file found. Run: bash ./scripts/app.sh start"
    server_url_from_json "$json_path"
    ;;
  open)
    ensure_venv
    cleanup_stale_runtime_files >/dev/null
    json_path="$(server_json_for_root || true)"
    [[ -n "$json_path" ]] || die "No repo-local Jupyter runtime file found. Run: bash ./scripts/app.sh start"
    url="$(server_url_from_json "$json_path")"
    open "$url"
    echo "Opened: $url"
    ;;
  logs)
    show_logs "${1:-}"
    ;;
  validate)
    validate_command "$@"
    ;;
  -h|--help|help)
    usage
    ;;
  *)
    usage >&2
    die "Unknown command: $COMMAND"
    ;;
esac