#!/usr/bin/env bash

# Exit immediately if any command fails
set -e

# Make sure the script is running with superuser privileges
if [ "$EUID" -ne 0 ]; then
  echo "Error: This deployment script must be run as root (sudo)." >&2
  exit 1
fi

# ====================================================
# ⚙️ CONFIGURATION: Client ID, Central URL, Agent Token
# ====================================================
# Resolution priority for each value: explicit CLI arg > existing install > default.
# Reading from the existing install makes re-running with no args a safe in-place
# upgrade — the node keeps its identity. This is what fleet self-update relies on.
#   Usage: deploy_agent.sh [CLIENT_ID] [CENTRAL_URL] [AGENT_TOKEN]
EXISTING_PY="/opt/seekright-agent/client_agent_prototype.py"

CLIENT_ID="${1:-}"
CENTRAL_URL="${2:-}"
ARG_TOKEN="${3:-}"

if [ -f "$EXISTING_PY" ]; then
  if [ -z "$CLIENT_ID" ]; then
    CLIENT_ID="$(sed -n 's/^CLIENT_ID = "\(.*\)"/\1/p' "$EXISTING_PY" | head -1)"
  fi
  if [ -z "$CENTRAL_URL" ]; then
    CENTRAL_URL="$(sed -n 's/^CENTRAL_BASE = "\(.*\)"/\1/p' "$EXISTING_PY" | head -1)"
  fi
fi

# Fallbacks when nothing was provided and there is no prior install
if [ -z "$CLIENT_ID" ]; then CLIENT_ID="$(hostname)"; fi
if [ -z "$CENTRAL_URL" ]; then CENTRAL_URL="http://rmm-backend.seekright.com"; fi

# Agent token. The server injects the live token into the line below at download
# time (replacing __AGENT_TOKEN__). Priority: CLI arg > server-injected > existing install.
AGENT_TOKEN="__AGENT_TOKEN__"
if [ -n "$ARG_TOKEN" ]; then
  AGENT_TOKEN="$ARG_TOKEN"
elif [ "$AGENT_TOKEN" = "__AGENT_""TOKEN__" ]; then
  # (placeholder split in two so the server's token injection can't rewrite this
  # comparison — it must only replace the assignment above)
  if [ -f "$EXISTING_PY" ]; then
    AGENT_TOKEN="$(sed -n 's/^AGENT_TOKEN = "\(.*\)"/\1/p' "$EXISTING_PY" | head -1)"
  else
    AGENT_TOKEN=""
  fi
fi

echo "🚀 Preparing SeekRight RMM Agent installation (agent v3.2-auth)..."
echo "📍 Target Client ID: $CLIENT_ID"
echo "🌐 Central Server URL: $CENTRAL_URL"

# 1. Install System Dependencies
echo "🔄 Updating system package indexes and installing python dependencies..."

# Fix potential MySQL GPG key expiration issues before updating (common issue on Ubuntu 22.04)
apt-key del B7B3B788A8D3785C || true
wget -q -O - https://repo.mysql.com/RPM-GPG-KEY-mysql-2025 | apt-key add - || true
wget -q -O - https://repo.mysql.com/RPM-GPG-KEY-mysql-2025 | gpg --dearmor --yes -o /usr/share/keyrings/mysql-apt-config.gpg || true

apt-get update -y || true
apt-get install -y python3 python3-pip python3-psutil

# 2. Create Destination Directories
echo "📁 Creating installation workspace..."
mkdir -p /opt/seekright-agent
chmod 755 /opt/seekright-agent

# 3. Dynamically Generate Python Agent File
echo "📝 Generating client agent file..."
cat << 'EOF' > /opt/seekright-agent/client_agent_prototype.py
import urllib.request
import urllib.parse
import json
import time
import subprocess
import psutil
import os
import sys
import string
import threading

def get_cpu_temp():
    if not hasattr(psutil, "sensors_temperatures"):
        return None
    try:
        temps = psutil.sensors_temperatures()
        if not temps:
            return None
        for name in ['coretemp', 'k10temp', 'acpitz', 'cpu_thermal']:
            if name in temps and temps[name]:
                return temps[name][0].current
        for name, entries in temps.items():
            if entries:
                return entries[0].current
    except Exception:
        pass
    return None

AGENT_VERSION = "3.3-shift"

CLIENT_ID = "TEMPLATE_CLIENT_ID"
CENTRAL_BASE = "TEMPLATE_CENTRAL_URL"
AGENT_TOKEN = "TEMPLATE_AGENT_TOKEN"

# Attach the shared headers (agent token + tunnel bypass) to every outbound request
# by installing a default opener, so we don't have to touch each call site.
_opener = urllib.request.build_opener()
_opener.addheaders = [('ngrok-skip-browser-warning', 'true'), ('X-Agent-Token', AGENT_TOKEN)]
urllib.request.install_opener(_opener)

CENTRAL_API_URL = f"{CENTRAL_BASE}/api/get-command?client_id={CLIENT_ID}"
CENTRAL_HEARTBEAT_URL = f"{CENTRAL_BASE}/api/heartbeat?client_id={CLIENT_ID}"

platform_name = "nt" if os.name == "nt" else "posix"
CENTRAL_COMMANDS_URL = f"{CENTRAL_BASE}/api/get-commands?platform={platform_name}"

CENTRAL_FILE_REQUEST_URL = f"{CENTRAL_BASE}/api/get-file-request?client_id={CLIENT_ID}"
CENTRAL_FILE_STATUS_URL = f"{CENTRAL_BASE}/api/file-transfer-status?client_id={CLIENT_ID}"

CACHED_COMMANDS = {}
command_in_progress = False
file_transfer_in_progress = False
SHIFT_PATH = None          # per-site SHIFT folder, configured on the dashboard, delivered via heartbeat
search_in_progress = False

def run_command_async(command_key, safe_cmd_list, is_shutdown_or_restart):
    global command_in_progress
    try:
        result = subprocess.run(safe_cmd_list, capture_output=True, text=True)
        output = result.stdout.strip()
        error_output = result.stderr.strip() if result.stderr else ""
        
        # Post execution results to Central Server
        send_command_result_to_server(command_key, result.returncode, output, error_output)
        
        if not is_shutdown_or_restart:
            if result.returncode == 0:
                print(f"    -> Output: {output}")
                send_rocket_chat_notification(f"✅ Success: {command_key}", output, "#00FF00")
            else:
                print(f"    -> Error: {error_output}")
                send_rocket_chat_notification(f"❌ Failed: {command_key}", error_output, "#FF0000")
    except Exception as ex:
        print(f"    -> [!] Execution Error: {ex}")
        send_command_result_to_server(command_key, -1, "", str(ex))
        if not is_shutdown_or_restart:
            send_rocket_chat_notification(f"❌ Execution Error: {command_key}", str(ex), "#FF0000")
    finally:
        command_in_progress = False

def execute_command_async(command_key):
    # Resolve a command key received via heartbeat against the whitelist and run it.
    # (The heartbeat loop sets command_in_progress=True before dispatching here, so
    # every exit path must clear it unless run_command_async takes over that duty.)
    global command_in_progress
    try:
        allowed_commands = load_allowed_commands()
        if command_key not in allowed_commands:
            warning_msg = f"Server requested unknown/unwhitelisted command: '{command_key}'. Ignored."
            print(f"    -> [!] SECURITY WARNING: {warning_msg}")
            send_rocket_chat_notification("⚠️ SECURITY WARNING", warning_msg, "#FFA500")
            command_in_progress = False
            return

        print(f"    -> Executing approved command: '{command_key}'")
        safe_cmd_list = allowed_commands[command_key]
        is_shutdown_or_restart = any(x in command_key.lower() for x in ["reboot", "restart", "kill_python"])
        if is_shutdown_or_restart:
            print("    -> Shutdown/restart command detected. Sending pre-execution notification...")
            send_rocket_chat_notification(
                f"🔄 Executing: {command_key}",
                f"System/agent is performing an action: {' '.join(safe_cmd_list)}\nConnection may drop temporarily.",
                "#FFA500"
            )
            time.sleep(3)
        # run_command_async clears command_in_progress in its finally block
        run_command_async(command_key, safe_cmd_list, is_shutdown_or_restart)
    except Exception as e:
        print(f"    -> [!] Command dispatch error: {e}")
        command_in_progress = False

def load_allowed_commands():
    global CACHED_COMMANDS
    try:
        req = urllib.request.Request(CENTRAL_COMMANDS_URL)
        req.add_header('ngrok-skip-browser-warning', 'true')
        with urllib.request.urlopen(req, timeout=15) as response:
            commands = json.loads(response.read().decode())
            if commands:
                CACHED_COMMANDS = commands
            return CACHED_COMMANDS
    except Exception as e:
        print(f"    -> [!] Failed to load commands from central API: {e}")
        if CACHED_COMMANDS:
            print("    -> Returning last successfully cached whitelist.")
            return CACHED_COMMANDS
        return {"whoami": ["whoami"]}

ROCKET_CHAT_WEBHOOK = "https://text.seekright.com/hooks/6a0eb41ea88367bc6e101f0c/gvXfD8zSNRti43kEabqdE9tMCvNi9zAAoGfbao7cxcKbiW6R"

def send_rocket_chat_notification(title, message, color="#764FA5"):
    try:
        if not message:
            message = "No output provided."
        payload = {
            "alias": f"Agent: {CLIENT_ID}",
            "text": "Execution Report",
            "attachments": [{
                "title": title,
                "text": message,
                "color": color
            }]
        }
        data = json.dumps(payload).encode('utf-8')
        req = urllib.request.Request(ROCKET_CHAT_WEBHOOK, data=data, headers={'Content-Type': 'application/json'}, method='POST')
        urllib.request.urlopen(req, timeout=15)
    except Exception as e:
        print(f"    -> [!] Failed to send webhook: {e}")

def get_linux_labels():
    labels = {}
    try:
        if os.path.exists("/dev/disk/by-label"):
            for label in os.listdir("/dev/disk/by-label"):
                try:
                    full_path = os.path.join("/dev/disk/by-label", label)
                    real_dev = os.path.realpath(full_path)
                    labels[real_dev] = label
                except Exception:
                    pass
    except Exception:
        pass
    return labels

def get_windows_label(drive):
    try:
        import ctypes
        volumeNameBuffer = ctypes.create_unicode_buffer(260)
        rc = ctypes.windll.kernel32.GetVolumeInformationW(
            ctypes.c_wchar_p(drive),
            volumeNameBuffer,
            ctypes.sizeof(volumeNameBuffer),
            None, None, None, None, 0
        )
        if rc:
            return volumeNameBuffer.value
    except Exception:
        pass
    return ""

def send_heartbeat():
    try:
        gpus = []
        try:
            result = subprocess.run(
                ["nvidia-smi", "--query-gpu=name,utilization.gpu,temperature.gpu,memory.used,memory.total,power.draw,power.limit,fan.speed", "--format=csv,noheader"], 
                capture_output=True, text=True
            )
            if result.returncode == 0:
                for line in result.stdout.strip().split('\n'):
                    if line:
                        parts = [p.strip() for p in line.split(',')]
                        if len(parts) >= 8:
                            gpus.append({
                                "name": parts[0],
                                "utilization": parts[1],
                                "temp": parts[2] + "C",
                                "memory_used": parts[3],
                                "memory_total": parts[4],
                                "power_draw": parts[5],
                                "power_limit": parts[6],
                                "fan_speed": parts[7]
                            })
        except FileNotFoundError:
            pass

        linux_labels = get_linux_labels() if os.name != "nt" else {}
        disks = []
        seen_devices = set()  # Deduplicate: skip bind-mounts that reuse the same device
        for part in psutil.disk_partitions(all=False):
            # Skip loop devices and unknown/empty fstypes
            if 'loop' in part.device or 'loop' in part.fstype or part.fstype == '':
                continue
            # Skip snap & squashfs mounts
            if part.mountpoint.startswith('/snap') or part.fstype == 'squashfs':
                continue
            # Skip boot, WSL internal, and other virtual/system paths
            if part.mountpoint.startswith('/boot') or part.mountpoint.startswith('/mnt/wsl'):
                continue
            if part.mountpoint.startswith('/var/snap') or part.mountpoint.startswith('/run'):
                continue
            # Skip virtual/overlay filesystems
            if part.fstype in ('tmpfs', 'devtmpfs', 'overlay', 'aufs', 'fusectl',
                               'proc', 'sysfs', 'cgroup', 'cgroup2', 'pstore',
                               'securityfs', 'debugfs', 'tracefs', 'bpf'):
                continue
            # Deduplicate: each physical device reported only once (eliminates bind-mounts)
            if part.device in seen_devices:
                continue
            seen_devices.add(part.device)
            try:
                usage = psutil.disk_usage(part.mountpoint)
                
                # Resolve label/friendly name
                label = ""
                if os.name == "nt":
                    label = get_windows_label(part.mountpoint)
                    if not label:
                        if part.mountpoint == "C:\\":
                            label = "OS"
                        else:
                            _mp = part.mountpoint.strip('\\')
                            label = f"Local Disk ({_mp})"
                else:
                    label = linux_labels.get(part.device, "")
                    if not label:
                        if part.mountpoint == "/":
                            label = "Computer"
                        else:
                            label = os.path.basename(part.mountpoint) or "Volume"

                disks.append({
                    "mount": part.mountpoint,
                    "device": part.device,
                    "label": label,
                    "total_gb": round(usage.total / (1024**3), 2),
                    "free_gb": round(usage.free / (1024**3), 2),
                    "percent": usage.percent
                })
            except PermissionError:
                continue

        payload = {
            "cpu_percent": psutil.cpu_percent(interval=0.5),
            "cpu_temp": get_cpu_temp(),
            "memory_percent": psutil.virtual_memory().percent,
            "memory_total_gb": round(psutil.virtual_memory().total / (1024**3), 2),
            "memory_free_gb": round(psutil.virtual_memory().available / (1024**3), 2),
            "disks": disks,
            "gpus": gpus,
            "agent_version": AGENT_VERSION
        }

        data = json.dumps(payload).encode('utf-8')
        req = urllib.request.Request(CENTRAL_HEARTBEAT_URL, data=data, headers={'Content-Type': 'application/json'}, method='POST')
        req.add_header('ngrok-skip-browser-warning', 'true')
        with urllib.request.urlopen(req, timeout=30) as response:
            res = json.loads(response.read().decode())
            
            # Handle commands
            cmd = res.get("command", "none")
            if cmd != "none":
                global command_in_progress
                if not command_in_progress:
                    command_in_progress = True
                    t = threading.Thread(target=execute_command_async, args=(cmd,))
                    t.daemon = True
                    t.start()
                    
            # Handle file requests
            fname = res.get("filename", "none")
            if fname != "none":
                global file_transfer_in_progress
                if not file_transfer_in_progress:
                    file_transfer_in_progress = True
                    t = threading.Thread(target=handle_file_request_async, args=(fname,))
                    t.daemon = True
                    t.start()

            # Server-configured SHIFT folder for this node
            global SHIFT_PATH
            sp = res.get("shift_path")
            if sp:
                SHIFT_PATH = sp

            # Handle video search requests
            squery = res.get("search", "none")
            if squery and squery != "none":
                global search_in_progress
                if not search_in_progress:
                    search_in_progress = True
                    t = threading.Thread(target=run_file_search, args=(squery,))
                    t.daemon = True
                    t.start()

        print("[*] Heartbeat successful.")
    except Exception as e:
        print(f"[!] Failed heartbeat: {e}")

def send_command_result_to_server(command: str, returncode: int, stdout: str, stderr: str):
    try:
        base_url = CENTRAL_API_URL.split("/api/")[0]
        url = f"{base_url}/api/command-result?client_id={CLIENT_ID}"
        payload = {
            "command": command,
            "returncode": returncode,
            "stdout": stdout,
            "stderr": stderr
        }
        data = json.dumps(payload).encode('utf-8')
        req = urllib.request.Request(
            url, 
            data=data, 
            headers={'Content-Type': 'application/json'}, 
            method='POST'
        )
        req.add_header('ngrok-skip-browser-warning', 'true')
        with urllib.request.urlopen(req, timeout=15) as response:
            pass
        print("    -> Sent command execution output to Central Server.")
    except Exception as e:
        print(f"    -> [!] Failed to send command result to server: {e}")

def poll_server():
    global command_in_progress
    if command_in_progress:
        return
    try:
        print(f"[*] Checking for commands...")
        req = urllib.request.Request(CENTRAL_API_URL)
        req.add_header('ngrok-skip-browser-warning', 'true')
        with urllib.request.urlopen(req, timeout=15) as response:
            data = json.loads(response.read().decode())
            command_key = data.get("command", "none")
            
            if command_key == "none":
                print("    -> No pending commands.")
            else:
                allowed_commands = load_allowed_commands()
                if command_key in allowed_commands:
                    print(f"    -> Executing approved command: '{command_key}'")
                    safe_cmd_list = allowed_commands[command_key]
                    
                    is_shutdown_or_restart = any(x in command_key.lower() for x in ["reboot", "restart", "kill_python"])
                    if is_shutdown_or_restart:
                        print("    -> Shutdown/restart command detected. Sending pre-execution success notification...")
                        send_rocket_chat_notification(
                            f"🔄 Executing: {command_key}", 
                            f"System/agent is performing a scheduled action: {' '.join(safe_cmd_list)}\nConnection may drop temporarily.", 
                            "#FFA500"
                        )
                        time.sleep(3)
                        
                    command_in_progress = True
                    t = threading.Thread(target=run_command_async, args=(command_key, safe_cmd_list, is_shutdown_or_restart))
                    t.daemon = True
                    t.start()
                else:
                    warning_msg = f"Server requested unknown/malicious command: '{command_key}'. Ignored."
                    print(f"    -> [!] SECURITY WARNING: {warning_msg}")
                    send_rocket_chat_notification("⚠️ SECURITY WARNING", warning_msg, "#FFA500")
    except Exception as e:
        print(f"    -> [!] Failed to connect to central server: {e}")

def get_dir_size(path):
    if os.name != "nt":
        try:
            result = subprocess.run(["du", "-sb", path], capture_output=True, text=True, timeout=30)
            if result.returncode == 0:
                return int(result.stdout.split()[0])
        except Exception:
            pass
    total_size = 0
    try:
        for dirpath, dirnames, filenames in os.walk(path):
            for f in filenames:
                fp = os.path.join(dirpath, f)
                try:
                    total_size += os.path.getsize(fp)
                except Exception:
                    pass
    except Exception:
        pass
    return total_size

def get_takeleap_sizes(root_path, subdirs):
    sizes = {d: 0 for d in subdirs}
    total_size = 0
    subdir_abs_paths = {}
    for d in subdirs:
        subdir_abs_paths[os.path.join(root_path, d)] = d
    try:
        for dirpath, dirnames, filenames in os.walk(root_path):
            dir_size = 0
            for f in filenames:
                fp = os.path.join(dirpath, f)
                try:
                    dir_size += os.path.getsize(fp)
                except Exception:
                    pass
            total_size += dir_size
            for sub_abs, sub_name in subdir_abs_paths.items():
                if dirpath == sub_abs or dirpath.startswith(sub_abs + os.sep):
                    sizes[sub_name] += dir_size
                    break
    except Exception:
        pass
    return total_size, sizes

def format_size(size_in_bytes):
    for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
        if size_in_bytes < 1024.0:
            return f"{size_in_bytes:.2f} {unit}"
        size_in_bytes /= 1024.0
    return f"{size_in_bytes:.2f} PB"

def check_takeleap_folders():
    target_names = {"SHIFT", "Error_Videos", "UPLOAD_FOLDER"}
    parent_found = False
    files_found = False
    
    if os.name == "nt":
        drives = [f"{letter}:\\" for letter in string.ascii_uppercase if os.path.exists(f"{letter}:\\")]
    else:
        drives = ["/mnt", "/media", "/home"]
        if os.path.exists("/TAKELEAP"):
            drives.append("/TAKELEAP")
        
    for drive in drives:
        for root, dirs, files in os.walk(drive, topdown=True):
            if os.name == "nt":
                dirs[:] = [d for d in dirs if d.lower() not in {"windows", "program files", "program files (x86)", "appdata", "programdata", "$recycle.bin", "system volume information", "node_modules", ".git"}]
            else:
                if root == "/":
                    dirs[:] = [d for d in dirs if d.lower() not in {"proc", "sys", "dev", "var", "lib", "run", "boot", "snap", "node_modules", ".git"}]
                elif root == "/mnt":
                    dirs[:] = [d for d in dirs if not (len(d) == 1 and d.isalpha()) and d.lower() not in {"wslg", "wsl"}]
                
            if os.path.basename(root).upper() == "TAKELEAP":
                parent_found = True
                for sub in dirs:
                    if sub in target_names:
                        subpath = os.path.join(root, sub)
                        try:
                            for item in os.listdir(subpath):
                                if os.path.isfile(os.path.join(subpath, item)):
                                    files_found = True
                                    break
                        except Exception:
                            pass
                    if files_found:
                        break
            if files_found:
                break
        if files_found:
            break
            
    if not parent_found:
        print("Folder not found")
    else:
        print("true" if files_found else "false")

def list_takeleap_parent_folders():
    found_any = False
    
    if os.name == "nt":
        drives = [f"{letter}:\\" for letter in string.ascii_uppercase if os.path.exists(f"{letter}:\\")]
    else:
        drives = ["/mnt", "/media", "/home"]
        if os.path.exists("/TAKELEAP"):
            drives.append("/TAKELEAP")
        
    for drive in drives:
        for root, dirs, files in os.walk(drive, topdown=True):
            if os.name == "nt":
                dirs[:] = [d for d in dirs if d.lower() not in {"windows", "program files", "program files (x86)", "appdata", "programdata", "$recycle.bin", "system volume information", "node_modules", ".git"}]
            else:
                if root == "/":
                    dirs[:] = [d for d in dirs if d.lower() not in {"proc", "sys", "dev", "var", "lib", "run", "boot", "snap", "node_modules", ".git"}]
                elif root == "/mnt":
                    dirs[:] = [d for d in dirs if not (len(d) == 1 and d.isalpha()) and d.lower() not in {"wslg", "wsl"}]
                
            if os.path.basename(root).upper() == "TAKELEAP":
                found_any = True
                print(f"\n[+] Found TAKELEAP at: {root}")
                try:
                    items = os.listdir(root)
                    if items:
                        print("    Contents:")
                        for idx, item in enumerate(items, 1):
                            is_dir = "Folder" if os.path.isdir(os.path.join(root, item)) else "File"
                            print(f"      {idx}. {item} ({is_dir})")
                    else:
                        print("    (Empty)")
                except Exception as e:
                    print(f"    (Error listing directory: {e})")
                    
    if not found_any:
        print("No TAKELEAP folders found on this system.")

def list_takeleap_parent_size():
    found_any = False
    
    if os.name == "nt":
        drives = [f"{letter}:\\" for letter in string.ascii_uppercase if os.path.exists(f"{letter}:\\")]
    else:
        drives = ["/mnt", "/media", "/home"]
        if os.path.exists("/TAKELEAP"):
            drives.append("/TAKELEAP")
        
    for drive in drives:
        for root, dirs, files in os.walk(drive, topdown=True):
            if os.name == "nt":
                dirs[:] = [d for d in dirs if d.lower() not in {"windows", "program files", "program files (x86)", "appdata", "programdata", "$recycle.bin", "system volume information", "node_modules", ".git"}]
            else:
                if root == "/":
                    dirs[:] = [d for d in dirs if d.lower() not in {"proc", "sys", "dev", "var", "lib", "run", "boot", "snap", "node_modules", ".git"}]
                elif root == "/mnt":
                    dirs[:] = [d for d in dirs if not (len(d) == 1 and d.isalpha()) and d.lower() not in {"wslg", "wsl"}]
                
            if os.path.basename(root).upper() == "TAKELEAP":
                found_any = True
                subdirs = list(dirs)
                total_size, subdir_sizes = get_takeleap_sizes(root, subdirs)
                print(f"\n[+] Found TAKELEAP at: {root} (Total Size: {format_size(total_size)})")
                if subdirs:
                    print("    Subfolders present:")
                    for d in subdirs:
                        sub_size = subdir_sizes.get(d, 0)
                        print(f"      - {d} (Size: {format_size(sub_size)})")
                else:
                    print("    (No subfolders found inside)")
                    
    if not found_any:
        print("No TAKELEAP folders found on this system.")

def list_takeleap_subfolder_size(subfolder_name):
    found_any = False
    
    if os.name == "nt":
        drives = [f"{letter}:\\" for letter in string.ascii_uppercase if os.path.exists(f"{letter}:\\")]
    else:
        drives = ["/mnt", "/media", "/home"]
        if os.path.exists("/TAKELEAP"):
            drives.append("/TAKELEAP")
        
    for drive in drives:
        for root, dirs, files in os.walk(drive, topdown=True):
            if os.name == "nt":
                dirs[:] = [d for d in dirs if d.lower() not in {"windows", "program files", "program files (x86)", "appdata", "programdata", "$recycle.bin", "system volume information", "node_modules", ".git"}]
            else:
                if root == "/":
                    dirs[:] = [d for d in dirs if d.lower() not in {"proc", "sys", "dev", "var", "lib", "run", "boot", "snap", "node_modules", ".git"}]
                elif root == "/mnt":
                    dirs[:] = [d for d in dirs if not (len(d) == 1 and d.isalpha()) and d.lower() not in {"wslg", "wsl"}]
                
            if os.path.basename(root).upper() == "TAKELEAP":
                for sub in dirs:
                    if sub.lower() == subfolder_name.lower():
                        subpath = os.path.join(root, sub)
                        try:
                            found_any = True
                            sub_size = get_dir_size(subpath)
                            print(f"\n[+] Folder: {subpath} (Total Size: {format_size(sub_size)})")
                        except Exception as e:
                            print(f"  [!] Error accessing folder {subpath}: {e}")
                            
    if not found_any:
        print(f"No TAKELEAP/{subfolder_name} folders found on this system.")

def list_takeleap_subfolder_files(subfolder_name):
    found_any = False
    
    if os.name == "nt":
        drives = [f"{letter}:\\" for letter in string.ascii_uppercase if os.path.exists(f"{letter}:\\")]
    else:
        drives = ["/mnt", "/media", "/home"]
        if os.path.exists("/TAKELEAP"):
            drives.append("/TAKELEAP")
        
    for drive in drives:
        for root, dirs, files in os.walk(drive, topdown=True):
            if os.name == "nt":
                dirs[:] = [d for d in dirs if d.lower() not in {"windows", "program files", "program files (x86)", "appdata", "programdata", "$recycle.bin", "system volume information", "node_modules", ".git"}]
            else:
                if root == "/":
                    dirs[:] = [d for d in dirs if d.lower() not in {"proc", "sys", "dev", "var", "lib", "run", "boot", "snap", "node_modules", ".git"}]
                elif root == "/mnt":
                    dirs[:] = [d for d in dirs if not (len(d) == 1 and d.isalpha()) and d.lower() not in {"wslg", "wsl"}]
                
            if os.path.basename(root).upper() == "TAKELEAP":
                for sub in dirs:
                    if sub.lower() == subfolder_name.lower():
                        subpath = os.path.join(root, sub)
                        try:
                            found_any = True
                            items = [item for item in os.listdir(subpath) if os.path.isfile(os.path.join(subpath, item))]
                            print(f"\n[+] Folder Contents: {subpath}")
                            if items:
                                for idx, item in enumerate(items, 1):
                                    print(f"  {idx}. {item}")
                            else:
                                print("  (Empty)")
                        except Exception as e:
                            print(f"  [!] Error accessing folder {subpath}: {e}")
                            
    if not found_any:
        print(f"No TAKELEAP/{subfolder_name} folders found on this system.")

def list_takeleap_subfolder(subfolder_name):
    list_takeleap_subfolder_files(subfolder_name)

def find_shift_roots():
    # The dashboard-configured path wins; without one, fall back to any
    # */TAKELEAP/SHIFT discovered by the legacy drive scan.
    if SHIFT_PATH and os.path.isdir(SHIFT_PATH):
        return [SHIFT_PATH]
    roots = []
    if os.name == "nt":
        drives = [f"{letter}:\\" for letter in string.ascii_uppercase if os.path.exists(f"{letter}:\\")]
    else:
        drives = ["/mnt", "/media", "/home"]
        if os.path.exists("/TAKELEAP"):
            drives.append("/TAKELEAP")
    for drive in drives:
        for root, dirs, files in os.walk(drive, topdown=True):
            if os.name == "nt":
                dirs[:] = [d for d in dirs if d.lower() not in {"windows", "program files", "program files (x86)", "appdata", "programdata", "$recycle.bin", "system volume information", "node_modules", ".git"}]
            else:
                if root == "/":
                    dirs[:] = [d for d in dirs if d.lower() not in {"proc", "sys", "dev", "var", "lib", "run", "boot", "snap", "node_modules", ".git"}]
                elif root == "/mnt":
                    dirs[:] = [d for d in dirs if not (len(d) == 1 and d.isalpha()) and d.lower() not in {"wslg", "wsl"}]
            if os.path.basename(root).upper() == "TAKELEAP":
                for sub in dirs:
                    if sub.upper() == "SHIFT":
                        roots.append(os.path.join(root, sub))
                dirs[:] = []  # don't descend into an already-claimed subtree
    return roots

def run_file_search(query):
    # Case-insensitive filename substring search in the SHIFT folder(s);
    # results are posted back for the dashboard to render.
    global search_in_progress
    MAX_RESULTS = 200
    try:
        target = query.lower()
        roots = find_shift_roots()
        results = []
        for base in roots:
            for froot, fdirs, ffiles in os.walk(base):
                for f in ffiles:
                    if target in f.lower():
                        fp = os.path.join(froot, f)
                        entry = {"name": f, "path": fp}
                        try:
                            st = os.stat(fp)
                            entry["size_mb"] = round(st.st_size / (1024 * 1024), 2)
                            entry["modified"] = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(st.st_mtime))
                        except Exception:
                            pass
                        results.append(entry)
                        if len(results) >= MAX_RESULTS:
                            break
                if len(results) >= MAX_RESULTS:
                    break
            if len(results) >= MAX_RESULTS:
                break
        payload = {"query": query, "results": results, "searched_path": ", ".join(roots)}
        if not roots:
            payload["error"] = "No SHIFT folder configured or found on this system"
        data = json.dumps(payload).encode('utf-8')
        req = urllib.request.Request(f"{CENTRAL_BASE}/api/search-results?client_id={CLIENT_ID}", data=data, headers={'Content-Type': 'application/json'}, method='POST')
        req.add_header('ngrok-skip-browser-warning', 'true')
        urllib.request.urlopen(req, timeout=30)
        print(f"    -> Search '{query}': posted {len(results)} match(es).")
    except Exception as e:
        print(f"    -> [!] File search failed: {e}")
    finally:
        search_in_progress = False

# ============================================================
# Remote File Transfer: dashboard requests a file by name,
# agent finds it under any TAKELEAP folder and uploads it.
# ============================================================

def send_file_transfer_status(filename, status, message="", progress_percent=None):
    try:
        payload = {"filename": filename, "status": status, "message": message}
        if progress_percent is not None:
            payload["progress_percent"] = progress_percent
        data = json.dumps(payload).encode('utf-8')
        req = urllib.request.Request(
            CENTRAL_FILE_STATUS_URL,
            data=data,
            headers={'Content-Type': 'application/json'},
            method='POST'
        )
        req.add_header('ngrok-skip-browser-warning', 'true')
        urllib.request.urlopen(req, timeout=30)
    except Exception as e:
        print(f"    -> [!] Failed to send file transfer status: {e}")

def find_file_in_takeleap(filename):
    target = filename.lower()

    # Fast-path optimization: Extract YYYYMMDD from filename
    # e.g., 20260612112725_000000.MP4 -> SHIFT/2026/June/12
    fast_path_rel = None
    try:
        if len(filename) >= 8 and filename[:8].isdigit():
            import calendar
            y = filename[0:4]
            m = int(filename[4:6])
            d = filename[6:8]
            if 1 <= m <= 12:
                month_name = calendar.month_name[m]
                fast_path_rel = os.path.join("SHIFT", y, month_name, d)
    except Exception:
        pass

    # Configured SHIFT path takes priority — SR/SHIFT sites have no TAKELEAP
    # folder anywhere, so the legacy scan below can never find their files.
    if SHIFT_PATH and os.path.isdir(SHIFT_PATH):
        if fast_path_rel:
            # SHIFT_PATH ends in .../SHIFT and fast_path_rel starts with SHIFT/,
            # so join the date subpath onto SHIFT_PATH's parent.
            specific_dir = os.path.join(os.path.dirname(SHIFT_PATH), fast_path_rel)
            if os.path.exists(specific_dir):
                for froot, fdirs, ffiles in os.walk(specific_dir):
                    for f in ffiles:
                        if f.lower() == target:
                            return os.path.join(froot, f)
        for froot, fdirs, ffiles in os.walk(SHIFT_PATH):
            for f in ffiles:
                if f.lower() == target:
                    return os.path.join(froot, f)

    if os.name == "nt":
        drives = [f"{letter}:\\" for letter in string.ascii_uppercase if os.path.exists(f"{letter}:\\")]
    else:
        drives = ["/mnt", "/media", "/home"]
        if os.path.exists("/TAKELEAP"):
            drives.append("/TAKELEAP")

    for drive in drives:
        for root, dirs, files in os.walk(drive, topdown=True):
            if os.name == "nt":
                dirs[:] = [d for d in dirs if d.lower() not in {"windows", "program files", "program files (x86)", "appdata", "programdata", "$recycle.bin", "system volume information", "node_modules", ".git"}]
            else:
                if root == "/":
                    dirs[:] = [d for d in dirs if d.lower() not in {"proc", "sys", "dev", "var", "lib", "run", "boot", "snap", "node_modules", ".git"}]
                elif root == "/mnt":
                    dirs[:] = [d for d in dirs if not (len(d) == 1 and d.isalpha()) and d.lower() not in {"wslg", "wsl"}]

            if os.path.basename(root).upper() == "TAKELEAP":
                # FAST PATH: Check the specific date folder first
                if fast_path_rel:
                    specific_dir = os.path.join(root, fast_path_rel)
                    if os.path.exists(specific_dir):
                        for froot, fdirs, ffiles in os.walk(specific_dir):
                            for f in ffiles:
                                if f.lower() == target:
                                    return os.path.join(froot, f)

                # Exhaustively search this TAKELEAP subtree for the requested file
                for froot, fdirs, ffiles in os.walk(root):
                    for f in ffiles:
                        if f.lower() == target:
                            return os.path.join(froot, f)
                dirs[:] = []  # subtree fully searched, don't descend again
    return None

def upload_file_to_server(filename, filepath):
    import time
    size = os.path.getsize(filepath)
    quoted = urllib.parse.quote(filename)
    CHUNK_SIZE = 512 * 1024  # 512 KB: small enough to finish well under the tunnel gateway timeout

    if size == 0:
        total_chunks = 1
    else:
        total_chunks = (size + CHUNK_SIZE - 1) // CHUNK_SIZE

    def fetch_received():
        # Returns the set of chunk indexes the server already holds, or None if unreachable
        url = f"{CENTRAL_BASE}/api/received-chunks?client_id={CLIENT_ID}&filename={quoted}&chunk_size={CHUNK_SIZE}&file_size={size}"
        try:
            req = urllib.request.Request(url)
            req.add_header('ngrok-skip-browser-warning', 'true')
            with urllib.request.urlopen(req, timeout=30) as r:
                return set(json.loads(r.read().decode()).get("received", []))
        except Exception as e:
            print(f"    -> [!] Could not fetch resume state ({e}).")
            return None

    def upload_one(f, chunk_index):
        # Returns True on success, False if the server reports the transfer cancelled
        f.seek(chunk_index * CHUNK_SIZE)
        chunk_data = f.read(CHUNK_SIZE)
        chunk_url = f"{CENTRAL_BASE}/api/upload-chunk?client_id={CLIENT_ID}&filename={quoted}&chunk_index={chunk_index}"
        while True:
            try:
                req = urllib.request.Request(chunk_url, data=chunk_data, method='POST')
                req.add_header('Content-Type', 'application/octet-stream')
                req.add_header('ngrok-skip-browser-warning', 'true')
                with urllib.request.urlopen(req, timeout=120) as response:
                    res = json.loads(response.read().decode())
                    return res.get("status") != "cancelled"
            except Exception as e:
                print(f"    -> [!] Chunk {chunk_index} failed ({e}). Retrying in 5s...")
                time.sleep(5)

    received = fetch_received() or set()
    if received:
        print(f"    -> Resuming: server already has {len(received)}/{total_chunks} chunks.")

    uploaded = 0
    with open(filepath, 'rb') as f:
        for chunk_index in range(total_chunks):
            if chunk_index in received:
                continue
            if not upload_one(f, chunk_index):
                print(f"    -> [!] Upload cancelled by server.")
                send_file_transfer_status(filename, "error", "Upload cancelled.")
                return
            uploaded += 1
            # Brief pause between chunks so heartbeat traffic can slip through
            time.sleep(0.1)
            if uploaded % 8 == 0 or chunk_index == total_chunks - 1:
                progress = round(((len(received) + uploaded) / total_chunks) * 100, 1)
                send_file_transfer_status(filename, "uploading", f"Uploading... ({progress}%)", progress_percent=progress)

        # Finalize with self-healing: verify the server holds every chunk (re-uploading
        # any lost to a mid-transfer cancel or cleanup), then ask it to stitch the file.
        while True:
            on_server = fetch_received()
            if on_server is None:
                time.sleep(5)
                continue
            missing = sorted(set(range(total_chunks)) - on_server)
            if missing:
                print(f"    -> [!] Server is missing {len(missing)} chunk(s). Re-uploading...")
                for chunk_index in missing:
                    if not upload_one(f, chunk_index):
                        print(f"    -> [!] Upload cancelled by server.")
                        send_file_transfer_status(filename, "error", "Upload cancelled.")
                        return
                    time.sleep(0.1)
                continue
            complete_url = f"{CENTRAL_BASE}/api/upload-complete?client_id={CLIENT_ID}&filename={quoted}&total_chunks={total_chunks}"
            try:
                req = urllib.request.Request(complete_url, method='POST', data=b'')
                req.add_header('ngrok-skip-browser-warning', 'true')
                with urllib.request.urlopen(req, timeout=180) as response:
                    return json.loads(response.read().decode())
            except Exception as e:
                print(f"    -> [!] Finalize failed ({e}). Re-verifying chunks in 5s...")
                time.sleep(5)

def handle_file_request_async(filename):
    global file_transfer_in_progress
    try:
        print(f"    -> File request received: '{filename}'. Searching TAKELEAP folders...")
        send_file_transfer_status(filename, "searching", "Searching TAKELEAP folders on client...")
        filepath = find_file_in_takeleap(filename)

        if not filepath:
            print(f"    -> [!] File '{filename}' not found in any TAKELEAP folder.")
            send_file_transfer_status(filename, "not_found", "File was not found in any TAKELEAP folder on this client.")
            send_rocket_chat_notification(f"❌ File Not Found: {filename}", f"Requested file was not found on client {CLIENT_ID}.", "#FF0000")
            return

        size_mb = round(os.path.getsize(filepath) / (1024 * 1024), 2)
        print(f"    -> Found at {filepath} ({size_mb} MB). Uploading...")
        send_file_transfer_status(filename, "uploading", f"Found at {filepath} ({size_mb} MB). Uploading to server...")
        upload_file_to_server(filename, filepath)
        # Server marks the transfer 'ready' once the upload completes
        print(f"    -> Upload complete: {filename}")
        send_rocket_chat_notification(f"✅ File Uploaded: {filename}", f"Uploaded {size_mb} MB from {filepath} to the central server.", "#00FF00")
    except Exception as ex:
        print(f"    -> [!] File transfer error: {ex}")
        send_file_transfer_status(filename, "error", f"Transfer failed on client: {ex}")
        send_rocket_chat_notification(f"❌ File Transfer Failed: {filename}", str(ex), "#FF0000")
    finally:
        file_transfer_in_progress = False

def poll_file_request():
    global file_transfer_in_progress
    if file_transfer_in_progress:
        return
    try:
        req = urllib.request.Request(CENTRAL_FILE_REQUEST_URL)
        req.add_header('ngrok-skip-browser-warning', 'true')
        with urllib.request.urlopen(req, timeout=15) as response:
            data = json.loads(response.read().decode())
            filename = data.get("filename", "none")
            if filename and filename != "none":
                file_transfer_in_progress = True
                t = threading.Thread(target=handle_file_request_async, args=(filename,))
                t.daemon = True
                t.start()
    except Exception as e:
        print(f"    -> [!] Failed to poll file requests: {e}")

if __name__ == "__main__":
    if len(sys.argv) > 1:
        if sys.argv[1] == "--check-takeleap":
            check_takeleap_folders()
            sys.exit(0)
        elif sys.argv[1] == "--list-takeleap":
            list_takeleap_parent_folders()
            sys.exit(0)
        elif sys.argv[1] == "--list-takeleap-size":
            list_takeleap_parent_size()
            sys.exit(0)
        elif sys.argv[1] == "--list-takeleap-sub" and len(sys.argv) > 2:
            list_takeleap_subfolder(sys.argv[2])
            sys.exit(0)
        elif sys.argv[1] == "--list-takeleap-sub-size" and len(sys.argv) > 2:
            list_takeleap_subfolder_size(sys.argv[2])
            sys.exit(0)
        elif sys.argv[1] == "--list-takeleap-sub-files" and len(sys.argv) > 2:
            list_takeleap_subfolder_files(sys.argv[2])
            sys.exit(0)
        else:
            print(f"Unknown argument or missing parameters: {sys.argv[1]}")
            sys.exit(1)
    else:
        print(f"Starting Agent v{AGENT_VERSION} for {CLIENT_ID}...")
        try:
            while True:
                send_heartbeat()
                time.sleep(10)
        except KeyboardInterrupt:
            print("\n[*] Agent stopped.")
EOF

# Inject the Client ID, Central URL, and Agent Token dynamically
sed -i "s/TEMPLATE_CLIENT_ID/$CLIENT_ID/g" /opt/seekright-agent/client_agent_prototype.py
sed -i "s|TEMPLATE_CENTRAL_URL|$CENTRAL_URL|g" /opt/seekright-agent/client_agent_prototype.py
sed -i "s|TEMPLATE_AGENT_TOKEN|$AGENT_TOKEN|g" /opt/seekright-agent/client_agent_prototype.py
# Agent file holds the shared secret; restrict it to root
chmod 600 /opt/seekright-agent/client_agent_prototype.py

# 4. Create Background systemd Service
echo "⚙️ Configuring background systemd daemon service..."
cat << 'EOF' > /etc/systemd/system/seekright-agent.service
[Unit]
Description=SeekRight RMM Background Agent
After=network-online.target
Wants=network-online.target

[Service]
Type=simple
# -u enables unbuffered stdout output for direct real-time logging via journalctl
ExecStart=/usr/bin/python3 -u /opt/seekright-agent/client_agent_prototype.py
WorkingDirectory=/opt/seekright-agent
Restart=always
RestartSec=5
User=root
Group=root
StandardOutput=syslog
StandardError=syslog
SyslogIdentifier=seekright-agent

[Install]
WantedBy=multi-user.target
EOF

chmod 644 /etc/systemd/system/seekright-agent.service

# 5. Enable and Launch Service
echo "🔄 Reloading system daemons and starting agent service..."
systemctl daemon-reload
systemctl enable seekright-agent.service
systemctl restart seekright-agent.service

echo ""
echo "✨ ============================================= ✨"
echo "✅ SeekRight RMM Agent successfully deployed!"
echo "📍 Location: /opt/seekright-agent/client_agent_prototype.py"
echo "📋 System Service: seekright-agent.service"
echo "✨ ============================================= ✨"
echo ""
echo "📈 Displaying real-time execution logs (Press Ctrl+C to exit log view):"
echo "--------------------------------------------------------"
journalctl -u seekright-agent.service -n 15 -f