# modules/users.py
import os
import subprocess
import pwd
import re
import json

from config import (
    RUNNING_ON_MAIN_SERVER,
    MAIN_SERVER_IP, MAIN_SERVER_USER, MAIN_SERVER_KEY, MAIN_SERVER_PORT,
)


def _run(cmd, timeout=30):
    try:
        r = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=timeout)
        return r.stdout.strip(), r.stderr.strip()
    except subprocess.TimeoutExpired:
        return '', 'timeout'
    except Exception as e:
        return '', str(e)


def _ssh_main(remote_cmd, timeout=30):
    """
    Run a command ON THE MAIN SERVER.
    - If already on main server → run locally.
    - If on VM → SSH to main server first.
    """
    if RUNNING_ON_MAIN_SERVER:
        return _run(remote_cmd, timeout=timeout)
    else:
        # Escape single quotes in remote_cmd for safe shell wrapping
        escaped = remote_cmd.replace("'", "'\\''")
        ssh = (
            f"ssh -i {MAIN_SERVER_KEY} -p {MAIN_SERVER_PORT} "
            f"-o StrictHostKeyChecking=no -o ConnectTimeout=10 "
            f"-o BatchMode=yes "
            f"{MAIN_SERVER_USER}@{MAIN_SERVER_IP}"
        )
        return _run(f"{ssh} '{escaped}'", timeout=timeout)


# ────────────────────────────────────────────────────────────────
# USER LISTING  — single batched SSH call to main server
# ────────────────────────────────────────────────────────────────

def get_all_users():
    """
    Return list of non-system users (uid >= 1000) from the MAIN SERVER.
    Uses a SINGLE SSH call with a bash one-liner that collects all data at once,
    instead of making 7+ SSH calls per user.
    """

    # This script runs on the main server and emits one JSON line per user
    batch_script = r"""
python3 - <<'PYEOF'
import subprocess, json, os, pwd

def run(cmd):
    try:
        r = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=10)
        return r.stdout.strip()
    except:
        return ''

# Get all non-system users
passwd_lines = run("getent passwd | awk -F: '$3 >= 1000 && $1 != \"nobody\" {print}'")
users = []

for line in passwd_lines.splitlines():
    parts = line.split(':')
    if len(parts) < 6:
        continue
    name = parts[0]
    uid  = int(parts[2])
    home = parts[5]

    sock      = f"/run/user/{uid}/docker.sock"
    has_docker = os.path.exists(sock)
    disk_used = run(f"du -sh {home} 2>/dev/null | cut -f1") or 'N/A'

    linger_out = run(f"loginctl show-user {name} --property=Linger 2>/dev/null")
    linger = 'yes' in linger_out.lower()

    container_count = 0
    if has_docker:
        cnt = run(f"DOCKER_HOST=unix://{sock} docker ps -aq 2>/dev/null | wc -l")
        try:
            container_count = int(cnt)
        except:
            pass

    img_path  = f"/home/{name}.img"
    has_vdisk = os.path.exists(img_path)
    vdisk_mount = None
    vdisk_size  = None
    if has_vdisk:
        mnt = run(f"findmnt -S {img_path} -o TARGET --noheadings 2>/dev/null")
        vdisk_mount = mnt or None
        sz = run(f"du -sh {img_path} 2>/dev/null | cut -f1")
        vdisk_size = sz or None

    users.append({
        'name':            name,
        'uid':             uid,
        'home':            home,
        'has_docker':      has_docker,
        'docker_socket':   sock if has_docker else None,
        'disk_used':       disk_used,
        'linger':          linger,
        'container_count': container_count,
        'has_vdisk':       has_vdisk,
        'vdisk_mount':     vdisk_mount,
        'vdisk_size':      vdisk_size,
    })

print(json.dumps(users))
PYEOF
"""

    stdout, stderr = _ssh_main(batch_script, timeout=60)

    if not stdout:
        # Fallback: try a simpler approach if python3 one-liner fails
        return _get_all_users_fallback()

    # Find the JSON line (last non-empty line that starts with '[')
    for line in reversed(stdout.splitlines()):
        line = line.strip()
        if line.startswith('['):
            try:
                return json.loads(line)
            except json.JSONDecodeError:
                break

    return _get_all_users_fallback()


def _get_all_users_fallback():
    """
    Simpler fallback: just get user names/UIDs, skip slow per-user checks.
    Returns basic user list without docker/disk details.
    """
    users = []
    stdout, _ = _ssh_main(
        "getent passwd | awk -F: '$3 >= 1000 && $1 != \"nobody\" {print $1\"|\"$3\"|\"$6}'"
    )
    if not stdout:
        return users

    for line in stdout.splitlines():
        line = line.strip()
        if not line:
            continue
        parts = line.split('|')
        if len(parts) < 3:
            continue
        try:
            uid = int(parts[1])
        except ValueError:
            continue
        users.append({
            'name':            parts[0],
            'uid':             uid,
            'home':            parts[2],
            'has_docker':      False,
            'docker_socket':   None,
            'disk_used':       'N/A',
            'linger':          False,
            'container_count': 0,
            'has_vdisk':       False,
            'vdisk_mount':     None,
            'vdisk_size':      None,
        })
    return users


# ────────────────────────────────────────────────────────────────
# USER CONTAINERS  — always from main server
# ────────────────────────────────────────────────────────────────

def get_user_containers(username):
    """Get containers running under a specific user's rootless docker on the main server."""
    uid_out, _ = _ssh_main(f"id -u {username} 2>/dev/null")
    try:
        uid = int(uid_out.strip())
    except ValueError:
        return []

    sock = f"/run/user/{uid}/docker.sock"
    sock_check, _ = _ssh_main(f"test -S {sock} && echo yes || echo no")
    if sock_check.strip() != 'yes':
        return []

    out, _ = _ssh_main(
        f"DOCKER_HOST=unix://{sock} "
        f"docker ps -a --format '{{{{.Names}}}}|{{{{.Status}}}}|{{{{.Image}}}}|{{{{.Ports}}}}' 2>/dev/null"
    )
    containers = []
    if out:
        for line in out.split('\n'):
            if '|' not in line:
                continue
            parts = line.split('|')
            containers.append({
                'name':   parts[0] if len(parts) > 0 else '',
                'status': parts[1] if len(parts) > 1 else '',
                'image':  parts[2] if len(parts) > 2 else '',
                'ports':  parts[3] if len(parts) > 3 else '',
                'owner':  username,
            })
    return containers


def get_all_users_containers():
    """Get containers from ALL users' rootless docker instances on main server."""
    all_containers = []
    for user in get_all_users():
        if user['has_docker']:
            all_containers.extend(get_user_containers(user['name']))
    return all_containers


# ────────────────────────────────────────────────────────────────
# USER DISK USAGE  — always from main server
# ────────────────────────────────────────────────────────────────

def get_user_disk_usage(username):
    home_out, _ = _ssh_main(
        f"getent passwd {username} | cut -d: -f6"
    )
    home = home_out.strip()
    if not home:
        return {}

    total_out, _ = _ssh_main(f"du -sh {home} 2>/dev/null | cut -f1")

    img_path = f"/home/{username}.img"
    vdisk_out, _ = _ssh_main(f"test -f {img_path} && echo yes || echo no")
    vdisk_info = {}

    if vdisk_out.strip() == 'yes':
        df_out, _ = _ssh_main(f"df -h {home} 2>/dev/null | tail -1")
        if df_out:
            parts = df_out.split()
            if len(parts) >= 4:
                vdisk_info = {
                    'size':      parts[1],
                    'used':      parts[2],
                    'available': parts[3],
                    'use_pct':   parts[4] if len(parts) > 4 else '?',
                }

    return {
        'home':  home,
        'total': total_out or 'N/A',
        'vdisk': vdisk_info,
    }


# ────────────────────────────────────────────────────────────────
# CREATE / DELETE USER  — only works on main server
# ────────────────────────────────────────────────────────────────

def create_user(username, password=None, setup_docker=True, disk_quota_mb=None):
    """
    Create a new system user on the MAIN SERVER.
    If called from VM, all commands SSH to main server.
    Returns (success: bool, log_text: str)
    """
    logs = []

    if not re.match(r'^[a-z][a-z0-9_-]{1,30}$', username):
        return False, "Invalid username. Use lowercase letters, numbers, _ or -"

    check_out, _ = _ssh_main(f"id {username} 2>/dev/null && echo exists || echo notfound")
    if 'exists' in check_out:
        return False, f"User '{username}' already exists"

    out, err = _ssh_main(f"useradd -m -s /bin/bash {username}")
    if err and 'already exists' not in err:
        return False, f"useradd failed: {err}"
    logs.append(f"✅ User {username} created")

    if password:
        out, err = _ssh_main(f"echo '{username}:{password}' | chpasswd")
        if err:
            logs.append(f"⚠️ Password set failed: {err}")
        else:
            logs.append("✅ Password set")

    _ssh_main(
        "apt-get install -y uidmap dbus-user-session curl 2>/dev/null",
        timeout=60
    )
    logs.append("✅ Prerequisites ready")

    _ssh_main(f"loginctl enable-linger {username}")
    logs.append("✅ Linger enabled")

    if disk_quota_mb:
        if RUNNING_ON_MAIN_SERVER:
            success, msg = _setup_virtual_disk(username, disk_quota_mb, logs)
            if not success:
                logs.append(f"⚠️ Virtual disk setup failed: {msg}")
        else:
            logs.append("⚠️ Virtual disk setup must be run directly on main server")

    if setup_docker:
        if RUNNING_ON_MAIN_SERVER:
            success, msg = _setup_rootless_docker_via_script(username, logs)
            if not success:
                logs.append(f"⚠️ Docker setup incomplete: {msg}")
        else:
            logs.append("⚠️ Rootless Docker setup must be run directly on main server")

    return True, '\n'.join(logs)


def delete_user(username, remove_home=False):
    """Remove a user from the MAIN SERVER. Returns (success, message)."""
    logs = []

    check_out, _ = _ssh_main(f"id {username} 2>/dev/null && echo exists || echo notfound")
    if 'notfound' in check_out:
        return False, f"User '{username}' does not exist"

    uid_out, _ = _ssh_main(f"id -u {username} 2>/dev/null")
    try:
        uid = int(uid_out.strip())
        _ssh_main(
            f"XDG_RUNTIME_DIR=/run/user/{uid} "
            f"su --login {username} --command "
            f"'systemctl --user stop docker-rootless.service 2>/dev/null' 2>/dev/null"
        )
    except Exception:
        pass

    _ssh_main(f"loginctl disable-linger {username} 2>/dev/null")

    if RUNNING_ON_MAIN_SERVER:
        _remove_virtual_disk(username, logs)
    else:
        img_path = f"/home/{username}.img"
        home = f"/home/{username}"
        _ssh_main(f"umount {home} 2>/dev/null || true")
        _ssh_main(f"rm -f {img_path} 2>/dev/null || true")

    flag = '-r' if remove_home else ''
    out, err = _ssh_main(f"userdel {flag} {username}")
    if err and 'does not exist' not in err and 'mail spool' not in err:
        return False, f"userdel error: {err}"

    msg = f"✅ User {username} deleted" + (" (home removed)" if remove_home else "")
    logs.append(msg)
    return True, '\n'.join(logs)


# ────────────────────────────────────────────────────────────────
# LOCAL-ONLY HELPERS (only called when RUNNING_ON_MAIN_SERVER)
# ────────────────────────────────────────────────────────────────

def _setup_virtual_disk(username, disk_mb, logs):
    try:
        pw = pwd.getpwnam(username)
    except KeyError as e:
        return False, str(e)

    home     = pw.pw_dir
    img_path = f"/home/{username}.img"

    logs.append(f"📦 Creating {disk_mb} MB virtual disk at {img_path} ...")

    out, err = _run(f"fallocate -l {disk_mb}M {img_path}", timeout=60)
    if err and 'fallocate' in err:
        logs.append("  ↳ fallocate not available, using dd...")
        out, err = _run(
            f"dd if=/dev/zero of={img_path} bs=1M count={disk_mb} status=none",
            timeout=600
        )
    if err:
        return False, f"Failed to create image: {err}"
    logs.append(f"  ✅ Image file created ({disk_mb} MB)")

    out, err = _run(f"mkfs.ext4 -F {img_path}", timeout=60)
    if err and 'error' in err.lower() and 'failed' in err.lower():
        return False, f"mkfs.ext4 failed: {err}"
    logs.append("  ✅ Formatted as ext4")

    tmp_backup = f"/tmp/{username}_home_backup"
    _run(f"cp -a {home}/. {tmp_backup}/ 2>/dev/null")

    out, err = _run(f"mount -o loop {img_path} {home}")
    if err:
        return False, f"mount failed: {err}"
    logs.append(f"  ✅ Mounted at {home}")

    _run(f"cp -a {tmp_backup}/. {home}/ 2>/dev/null")
    _run(f"rm -rf {tmp_backup}")
    _run(f"chown -R {username}:{username} {home}")
    logs.append("  ✅ Ownership set")

    fstab_line = f"{img_path} {home} ext4 loop,defaults 0 0\n"
    try:
        with open('/etc/fstab', 'r') as f:
            fstab = f.read()
        if img_path not in fstab:
            with open('/etc/fstab', 'a') as f:
                f.write(fstab_line)
            logs.append("  ✅ Added to /etc/fstab")
    except Exception as e:
        logs.append(f"  ⚠️ Could not update /etc/fstab: {e}")

    logs.append(f"✅ Virtual disk ready: {disk_mb} MB for {username}")
    return True, "ok"


def _remove_virtual_disk(username, logs):
    try:
        pw = pwd.getpwnam(username)
        home = pw.pw_dir
    except KeyError:
        return

    img_path = f"/home/{username}.img"
    _run(f"umount {home} 2>/dev/null")
    logs.append(f"  ↳ Unmounted {home}")

    if os.path.exists(img_path):
        os.remove(img_path)
        logs.append(f"  ↳ Removed {img_path}")

    try:
        with open('/etc/fstab', 'r') as f:
            lines = f.readlines()
        with open('/etc/fstab', 'w') as f:
            for line in lines:
                if img_path not in line:
                    f.write(line)
        logs.append("  ↳ Removed from /etc/fstab")
    except Exception as e:
        logs.append(f"  ⚠️ fstab cleanup error: {e}")


def _setup_rootless_docker_via_script(username, logs):
    try:
        pw = pwd.getpwnam(username)
        uid  = pw.pw_uid
        home = pw.pw_dir
    except KeyError as e:
        return False, str(e)

    _run("sysctl -w kernel.apparmor_restrict_unprivileged_userns=0")
    _run("echo 'kernel.apparmor_restrict_unprivileged_userns=0' >> /etc/sysctl.conf")

    runtime_dir = f"/run/user/{uid}"
    os.makedirs(runtime_dir, exist_ok=True)
    _run(f"chown {username}:{username} {runtime_dir}")
    _run(f"chmod 700 {runtime_dir}")

    logs.append(f"📝 Installing rootless Docker for {username}...")

    install_cmd = f"""bash -c '
export XDG_RUNTIME_DIR=/run/user/{uid}
export PATH=$HOME/bin:$PATH
mkdir -p $XDG_RUNTIME_DIR
chmod 700 $XDG_RUNTIME_DIR
curl -fsSL https://get.docker.com/rootless | sh
echo "export PATH=$HOME/bin:\\$PATH" >> ~/.bashrc
echo "export DOCKER_HOST=unix:///run/user/{uid}/docker.sock" >> ~/.bashrc
echo "export XDG_RUNTIME_DIR=/run/user/{uid}" >> ~/.bashrc
mkdir -p ~/.config/systemd/user
cat > ~/.config/systemd/user/docker.service << EOF
[Unit]
Description=Docker Rootless Daemon
After=network.target
[Service]
Type=simple
ExecStart=$HOME/bin/dockerd-rootless.sh
Restart=always
RestartSec=10
Environment=PATH=$HOME/bin:/usr/local/bin:/usr/bin:/bin
Environment=XDG_RUNTIME_DIR=/run/user/{uid}
[Install]
WantedBy=default.target
EOF
systemctl --user daemon-reload
systemctl --user enable docker.service
systemctl --user start docker.service
sleep 5
'"""

    try:
        result = subprocess.run(
            ['su', '-', username, '-c', install_cmd],
            capture_output=True, text=True, timeout=300
        )
        for line in result.stdout.split('\n'):
            if line.strip():
                logs.append(line.strip())
        if result.stderr:
            for line in result.stderr.split('\n')[-5:]:
                if line.strip():
                    logs.append(f"  stderr: {line.strip()}")

        sock = f"/run/user/{uid}/docker.sock"
        if os.path.exists(sock):
            logs.append(f"✅ Rootless Docker ready for {username}")
            return True, "ok"
        else:
            logs.append(f"⚠️ Socket not found at {sock}")
            return True, "socket_pending"

    except subprocess.TimeoutExpired:
        logs.append("⚠️ Setup timed out after 300s")
        return False, "timeout"
    except Exception as e:
        logs.append(f"⚠️ Setup failed: {e}")
        return False, str(e)