#!/bin/bash
# Install Unreal MCP and connect it to VS Code (via Continue MCP) on Arch Linux
# - Installs deps: git, jq, uv, python
# - Clones https://github.com/chongdashu/unreal-mcp
# - Creates a launcher: ~/.local/bin/unreal-mcp-server
# - Configures VS Code Continue MCP: ~/.continue/config.json
# - Optional: copies UnrealMCP plugin into a specified .uproject's Plugins/

set -euo pipefail

SCRIPT_NAME="$(basename "$0")"

# Source common library for shared functions
SCRIPT_DIR="$(dirname "$(readlink -f "$0")")"
# shellcheck source=../../lib/common.sh
source "$SCRIPT_DIR/../../lib/common.sh"

# ---------- User/paths ----------
set_actual_user_vars

INSTALL_ROOT_DEFAULT="$USER_HOME/.local/share/unreal-mcp"
INSTALL_ROOT="$INSTALL_ROOT_DEFAULT"
REPO_URL="https://github.com/chongdashu/unreal-mcp.git"
REPO_DIR="" # will be set after INSTALL_ROOT known

PROJECT_UPROJECT=""     # optional: path to .uproject
RESOLVED_PROJECT_DIR="" # directory containing the resolved .uproject
CONFIGURE_CONTINUE=true
CONFIGURE_VSCODE_USER=true
FORCE_UPDATE=false

log() { echo "[$(date '+%Y-%m-%d %H:%M:%S')] $*"; }
fail() {
  echo "[ERROR] $*" >&2
  exit 1
}

usage() {
  cat << EOF
Usage: $SCRIPT_NAME [options]

Options:
  --install-dir DIR        Install root for repo (default: $INSTALL_ROOT_DEFAULT)
  --project PATH           Path to your Unreal project (.uproject file) or a directory containing one
                           Copies UnrealMCP plugin into this Unreal project
  --no-continue            Skip configuring VS Code Continue MCP
  --no-vscode              Skip adding MCP server to VS Code user profile via --add-mcp
  --force-update           If repo exists, fetch and reset to origin/main
  -h, --help               Show this help

Examples:
  $SCRIPT_NAME --project ~/UnrealProjects/MyGame/MyGame.uproject
  $SCRIPT_NAME --install-dir "$USER_HOME/dev/unreal-mcp"
EOF
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --install-dir)
      shift
      [[ $# -gt 0 ]] || fail "--install-dir requires a value"
      INSTALL_ROOT="$1"
      ;;
    --project)
      shift
      [[ $# -gt 0 ]] || fail "--project requires a path to .uproject"
      PROJECT_UPROJECT="$1"
      ;;
    --no-continue)
      CONFIGURE_CONTINUE=false
      ;;
    --no-vscode)
      CONFIGURE_VSCODE_USER=false
      ;;
    --force-update)
      FORCE_UPDATE=true
      ;;
    -h | --help)
      usage
      exit 0
      ;;
    *)
      fail "Unknown option: $1"
      ;;
  esac
  shift
done

REPO_DIR="$INSTALL_ROOT/unreal-mcp"

# ---------- Dependencies ----------
require_cmd() { command -v "$1" > /dev/null 2>&1; }

ensure_packages_arch() {
  # Install with pacman using sudo when needed; keep idempotent with --needed
  local pkgs=(git jq uv python rsync)
  local to_install=()
  for p in "${pkgs[@]}"; do
    if ! pacman -Qi "$p" > /dev/null 2>&1; then
      to_install+=("$p")
    fi
  done
  if [[ ${#to_install[@]} -gt 0 ]]; then
    log "Installing packages: ${to_install[*]}"
    if [[ $EUID -eq 0 ]]; then
      pacman -S --noconfirm --needed "${to_install[@]}"
    else
      sudo pacman -S --noconfirm --needed "${to_install[@]}"
    fi
  else
    log "All required packages already installed"
  fi
}

check_python_version() {
  if require_cmd python; then
    local v
    v=$(python -V 2>&1 | awk '{print $2}')
  elif require_cmd python3; then
    local v
    v=$(python3 -V 2>&1 | awk '{print $2}')
  else
    log "python not found; pacman install will provide it"
    return 0
  fi
  # Require >= 3.12 (Unreal MCP docs)
  local major minor
  major=$(echo "$v" | cut -d. -f1)
  minor=$(echo "$v" | cut -d. -f2)
  if ((major < 3 || (major == 3 && minor < 12))); then
    log "Python $v detected; installing newer python via pacman"
    if [[ $EUID -eq 0 ]]; then
      pacman -S --noconfirm --needed python
    else
      sudo pacman -S --noconfirm --needed python
    fi
  fi
}

# ---------- Git clone/update ----------
setup_repo() {
  mkdir -p "$INSTALL_ROOT"
  if [[ ! -d "$REPO_DIR/.git" ]]; then
    log "Cloning unreal-mcp into $REPO_DIR"
    if require_cmd git; then
      git clone "$REPO_URL" "$REPO_DIR"
    else
      fail "git is required but not found after install"
    fi
  else
    log "Repo exists at $REPO_DIR"
    if [[ $FORCE_UPDATE == true ]]; then
      log "Updating repo with --force-update"
      git -C "$REPO_DIR" fetch origin
      git -C "$REPO_DIR" reset --hard origin/main
      git -C "$REPO_DIR" pull --rebase --autostash
    else
      log "Pulling latest changes"
      git -C "$REPO_DIR" pull --rebase --autostash
    fi
  fi

  # Ensure ownership for the real user when script ran via sudo
  if [[ $EUID -eq 0 ]]; then
    chown -R "$ACTUAL_USER:$ACTUAL_USER" "$INSTALL_ROOT"
  fi
}

# ---------- Launcher ----------
install_launcher() {
  local bin_dir="$USER_HOME/.local/bin"
  local python_dir="$REPO_DIR/Python"
  local launcher="$bin_dir/unreal-mcp-server"
  mkdir -p "$bin_dir"
  cat > "$launcher" << EOF
#!/bin/bash
set -euo pipefail
exec uv --directory "$python_dir" run unreal_mcp_server.py "\${1:-}" < /dev/null
EOF
  chmod +x "$launcher"
  if [[ $EUID -eq 0 ]]; then chown "$ACTUAL_USER:$ACTUAL_USER" "$launcher"; fi
  log "Installed launcher: $launcher"
}

# ---------- VS Code: Continue MCP config ----------
configure_continue() {
  if [[ $CONFIGURE_CONTINUE != true ]]; then
    log "Skipping Continue config (--no-continue)"
    return 0
  fi

  local cont_dir="$USER_HOME/.continue"
  local cont_cfg="$cont_dir/config.json"
  local python_dir="$REPO_DIR/Python"
  mkdir -p "$cont_dir"

  # Base JSON when no config exists
  local tmp_file
  tmp_file="$(mktemp)"
  if [[ ! -f $cont_cfg ]]; then
    cat > "$tmp_file" << JSON
{
  "mcpServers": {
    "unrealMCP": {
      "command": "uv",
      "args": ["--directory", "$python_dir", "run", "unreal_mcp_server.py"]
    }
  }
}
JSON
    mv "$tmp_file" "$cont_cfg"
  else
    # Merge using jq: ensure .mcpServers exists, then set/overwrite unrealMCP
    if ! require_cmd jq; then
      fail "jq is required to merge ~/.continue/config.json"
    fi
    jq --arg dir "$python_dir" '
      .mcpServers = (.mcpServers // {}) |
      .mcpServers.unrealMCP = {
        command: "uv",
        args: ["--directory", $dir, "run", "unreal_mcp_server.py"]
      }
    ' "$cont_cfg" > "$tmp_file" && mv "$tmp_file" "$cont_cfg"
  fi

  if [[ $EUID -eq 0 ]]; then chown "$ACTUAL_USER:$ACTUAL_USER" "$cont_cfg"; fi
  log "Configured Continue MCP at: $cont_cfg"
}

# ---------- VS Code user MCP (native) ----------
configure_vscode_user_mcp() {
  if [[ $CONFIGURE_VSCODE_USER != true ]]; then
    log "Skipping VS Code user MCP config (--no-vscode)"
    return 0
  fi

  if ! require_cmd jq; then
    fail "jq is required to compose VS Code --add-mcp JSON and to parse profiles"
  fi

  local python_dir="$REPO_DIR/Python"
  local json
  json=$(jq -n --arg dir "$python_dir" '{name:"unrealMCP", command:"uv", args:["--directory", $dir, "run", "unreal_mcp_server.py"]}')

  # Handle multiple VS Code variants if present
  local candidates=(code code-insiders codium)
  local found_any=false
  for cli in "${candidates[@]}"; do
    if ! command -v "$cli" > /dev/null 2>&1; then
      continue
    fi
    found_any=true
    log "Registering MCP server in VS Code user profile via: $cli --add-mcp"
    if "$cli" --add-mcp "$json" > "/tmp/${cli}-add-mcp.log" 2>&1; then
      log "[$cli] user profile: unrealMCP added/updated"
    else
      sed -n '1,200p' "/tmp/${cli}-add-mcp.log" || true
      fail "[$cli] --add-mcp failed for user profile. Ensure your VS Code supports MCP or rerun with --no-vscode."
    fi

    # Detect profiles with 'unreal' (case-insensitive) and add there too
    local data_dir=""
    case "$cli" in
      code)
        data_dir="$USER_HOME/.config/Code"
        ;;
      code-insiders)
        data_dir="$USER_HOME/.config/Code - Insiders"
        ;;
      codium)
        data_dir="$USER_HOME/.config/VSCodium"
        ;;
    esac
    local profiles_json="$data_dir/User/profiles/profiles.json"
    if [[ -f $profiles_json ]]; then
      # Extract profile names matching /unreal/i
      mapfile -t unreal_profiles < <(jq -r '.profiles // [] | .[] | .name // empty | select(test("unreal"; "i"))' "$profiles_json")
      if [[ ${#unreal_profiles[@]} -gt 0 ]]; then
        log "[$cli] Found profiles with 'unreal': ${unreal_profiles[*]}"
        local name
        for name in "${unreal_profiles[@]}"; do
          log "[$cli] Adding unrealMCP to profile: $name"
          if "$cli" --profile "$name" --add-mcp "$json" > "/tmp/${cli}-add-mcp-${name// /_}.log" 2>&1; then
            log "[$cli] profile '$name': unrealMCP added/updated"
          else
            sed -n '1,200p' "/tmp/${cli}-add-mcp-${name// /_}.log" || true
            fail "[$cli] --add-mcp failed for profile '$name'."
          fi
        done
      else
        log "[$cli] No VS Code profiles with 'unreal' in name"
      fi
    else
      log "[$cli] Profiles file not found: $profiles_json (skipping profile-specific adds)"
    fi
  done

  if [[ $found_any == false ]]; then
    fail "VS Code CLI not found (code/code-insiders/codium). Install VS Code and ensure 'code' CLI is available, or run with --no-vscode to skip."
  fi
}

# ---------- Unreal Plugin copy (optional) ----------
install_plugin_into_project() {
  [[ -n $PROJECT_UPROJECT ]] || return 0
  local upath="$PROJECT_UPROJECT"
  if [[ -d $upath ]]; then
    # Resolve .uproject in the provided directory
    mapfile -t _uprojects < <(find "$upath" -maxdepth 1 -type f -name "*.uproject" 2> /dev/null || true)
    if [[ ${#_uprojects[@]} -eq 0 ]]; then
      fail "--project directory '$upath' contains no .uproject files"
    elif [[ ${#_uprojects[@]} -gt 1 ]]; then
      printf '[ERROR] Multiple .uproject files found in %s:\n' "$upath" >&2
      printf '  - %s\n' "${_uprojects[@]}" >&2
      fail "Please pass the specific .uproject path to --project"
    else
      upath="${_uprojects[0]}"
      log "Resolved .uproject: $upath"
    fi
  elif [[ -f $upath ]]; then
    true
  else
    fail "--project path does not exist: $upath"
  fi
  if [[ ${upath##*.} != "uproject" ]]; then
    fail "--project must point to a .uproject file (got: $upath)"
  fi
  local proj_dir
  proj_dir="$(cd "$(dirname "$upath")" && pwd)"
  RESOLVED_PROJECT_DIR="$proj_dir"
  local src_plugin="$REPO_DIR/MCPGameProject/Plugins/UnrealMCP"
  local dst_plugin="$proj_dir/Plugins/UnrealMCP"
  if [[ ! -d $src_plugin ]]; then
    fail "Source plugin not found at $src_plugin (did repo layout change?)"
  fi
  mkdir -p "$proj_dir/Plugins"
  log "Copying UnrealMCP plugin to project: $dst_plugin"
  rsync -a --delete "$src_plugin/" "$dst_plugin/"
  # Set ownership back to actual user if run as root
  if [[ $EUID -eq 0 ]]; then chown -R "$ACTUAL_USER:$ACTUAL_USER" "$proj_dir/Plugins"; fi
  log "Plugin installed. Enable it from Unreal Editor (Edit > Plugins) if needed."
}

# ---------- Summary ----------
print_summary() {
  local python_dir="$REPO_DIR/Python"
  local plugin_dest="N/A"
  if [[ -n $RESOLVED_PROJECT_DIR ]]; then
    plugin_dest="$RESOLVED_PROJECT_DIR/Plugins/UnrealMCP"
  fi
  cat << EOF
============================================
Unreal MCP setup complete
============================================

Repo:        $REPO_DIR
Python dir:  $python_dir
Launcher:    $USER_HOME/.local/bin/unreal-mcp-server

VS Code (Continue) MCP configured: ${CONFIGURE_CONTINUE}
  - File: $USER_HOME/.continue/config.json
  - Server ID: unrealMCP

VS Code (User profile) MCP configured: ${CONFIGURE_VSCODE_USER}
  - Command used: code --add-mcp '{"name":"unrealMCP", "command":"uv", "args":["--directory","$python_dir","run","unreal_mcp_server.py"]}'

Optional usage:
  - Run server manually: unreal-mcp-server
  - In VS Code with Continue installed, the unrealMCP server will auto-start when needed.

Unreal plugin:
  - Source: MCPGameProject/Plugins/UnrealMCP
  - If you provided --project, the plugin was copied to: $plugin_dest
  - In the Unreal Editor: Edit > Plugins > search "UnrealMCP" and enable. Restart when prompted.

Notes:
  - Ensure you have Unreal Engine 5.5+ installed.
  - The Python server listens to the Unreal plugin on TCP port 55557 by default.
  - For other MCP clients (Claude Desktop, Cursor, Windsurf), copy the JSON snippet from the repo README to their config locations.
EOF
}

main() {
  log "Installing prerequisites (Arch Linux)"
  ensure_packages_arch
  check_python_version
  setup_repo
  install_launcher
  configure_continue
  install_plugin_into_project
  configure_vscode_user_mcp
  print_summary
}

main "$@"