#!/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.

  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 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/COURSE_BLUEPRINT.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
}

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
  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."
    return 0
  fi

  pid="$(server_field "$json_path" "pid")"
  if ! kill -0 "$pid" >/dev/null 2>&1; then
    local removed_paths
    removed_paths="$(cleanup_stale_runtime_files 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
    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
      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

  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" "$@"
}

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
    ;;
  restart)
    restart_command "$@"
    ;;
  status)
    exec bash "$ROOT_DIR/scripts/project_status.sh"
    ;;
  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