#!/usr/bin/env python3
"""UClaude Updater — automatic Claude Code patch updater.

Usage:
    sudo python3 uclaude_updater.py              # Check and update if new version available
    sudo python3 uclaude_updater.py --check      # Only check, don't install
    sudo python3 uclaude_updater.py --force       # Update even if version matches
    sudo python3 uclaude_updater.py --settings-only  # Only patch settings, don't touch cli.js
"""

import argparse
import json
import os
import re
import shutil
import subprocess
import sys
import time

# ============================================================
#  Platform detection
# ============================================================

IS_WINDOWS = sys.platform == "win32"
IS_MACOS = sys.platform == "darwin"

try:
    import pwd
    HAS_PWD = True
except ImportError:
    HAS_PWD = False

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))

# ANSI colors
G = "\033[92m"   # green
Y = "\033[93m"   # yellow
R = "\033[91m"   # red
W = "\033[97m"   # white/bold
D = "\033[0m"    # default


# ============================================================
#  Helpers
# ============================================================

def is_admin():
    if IS_WINDOWS:
        try:
            import ctypes
            return ctypes.windll.shell32.IsUserAnAdmin() != 0
        except Exception:
            return False
    return os.geteuid() == 0


def safe_chown(path, uid, gid):
    if not IS_WINDOWS and is_admin():
        os.chown(path, uid, gid)


def safe_chmod(path, mode):
    if not IS_WINDOWS:
        os.chmod(path, mode)


def eprint(*args, **kwargs):
    print(*args, file=sys.stderr, **kwargs)


# ============================================================
#  Version detection
# ============================================================

def find_cli_js():
    """Find installed Claude Code cli.js path."""
    candidates = []
    if IS_WINDOWS:
        for env_key in ("APPDATA", "LOCALAPPDATA", "PROGRAMFILES"):
            base = os.environ.get(env_key, "")
            if base:
                candidates.append(os.path.join(base, "npm", "node_modules", "@anthropic-ai", "claude-code", "cli.js"))
    else:
        candidates = [
            "/usr/lib/node_modules/@anthropic-ai/claude-code/cli.js",
            "/usr/local/lib/node_modules/@anthropic-ai/claude-code/cli.js",
            "/opt/homebrew/lib/node_modules/@anthropic-ai/claude-code/cli.js",
        ]

    for path in candidates:
        if os.path.isfile(path):
            return path
    return None


def get_installed_version():
    """Get currently installed Claude Code version."""
    cli_js = find_cli_js()
    if not cli_js:
        return None, None

    # Try package.json first (same directory as cli.js)
    pkg_json = os.path.join(os.path.dirname(cli_js), "package.json")
    if os.path.isfile(pkg_json):
        try:
            with open(pkg_json, "r") as f:
                data = json.load(f)
            return data.get("version"), cli_js
        except Exception:
            pass

    # Fallback: claude --version
    try:
        result = subprocess.run(
            ["claude", "--version"],
            capture_output=True, text=True, timeout=10,
        )
        m = re.search(r"(\d+\.\d+\.\d+)", result.stdout)
        if m:
            return m.group(1), cli_js
    except Exception:
        pass

    return None, cli_js


def get_latest_version():
    """Read latest version from local index.json."""
    index_path = os.path.join(SCRIPT_DIR, "claude", "releases", "index.json")
    if not os.path.isfile(index_path):
        return None
    try:
        with open(index_path, "r") as f:
            data = json.load(f)
        return data.get("latest")
    except Exception:
        return None


def ver_tuple(v):
    """Parse version string to tuple for comparison."""
    m = re.match(r"(\d+)\.(\d+)\.(\d+)", v or "")
    return (int(m.group(1)), int(m.group(2)), int(m.group(3))) if m else (0, 0, 0)


# ============================================================
#  Git pull
# ============================================================

def git_pull():
    """Pull latest changes from remote."""
    try:
        result = subprocess.run(
            ["git", "pull", "--quiet"],
            cwd=SCRIPT_DIR, capture_output=True, text=True, timeout=30,
        )
        if result.returncode == 0:
            return True
        eprint(f"  {Y}git pull warning: {result.stderr.strip()}{D}")
        return True  # non-fatal
    except subprocess.TimeoutExpired:
        eprint(f"  {Y}git pull timed out{D}")
        return True
    except FileNotFoundError:
        eprint(f"  {Y}git not found, skipping pull{D}")
        return True


# ============================================================
#  CLI.js installation
# ============================================================

def install_cli_js(version, cli_js_path):
    """Install patched cli.js for the given version."""
    release_cli = os.path.join(SCRIPT_DIR, "claude", "releases", f"v{version}", "cli.js")
    if not os.path.isfile(release_cli):
        eprint(f"  {R}Release cli.js not found: {release_cli}{D}")
        return False

    # Backup
    timestamp = time.strftime("%Y%m%d%H%M%S")
    backup_path = f"{cli_js_path}.bak.{timestamp}"

    try:
        shutil.copy2(cli_js_path, backup_path)
        shutil.copy2(release_cli, cli_js_path)

        # Syntax check
        result = subprocess.run(
            ["node", "--check", cli_js_path],
            capture_output=True, text=True, timeout=30,
        )
        if result.returncode != 0:
            eprint(f"  {R}Syntax check FAILED, rolling back...{D}")
            shutil.copy2(backup_path, cli_js_path)
            return False

        print(f"  {G}cli.js installed successfully{D}")
        print(f"  Backup: {backup_path}")
        return True
    except Exception as e:
        eprint(f"  {R}Installation error: {e}{D}")
        if os.path.isfile(backup_path):
            shutil.copy2(backup_path, cli_js_path)
            eprint(f"  Rolled back to backup")
        return False


# ============================================================
#  Settings patching
# ============================================================

def load_config():
    """Load patcher.config.json from repo root."""
    config_path = os.path.join(SCRIPT_DIR, "patcher.config.json")
    if not os.path.isfile(config_path):
        eprint(f"  {R}patcher.config.json not found{D}")
        return None
    with open(config_path, "r") as f:
        return json.load(f)


def ensure_dir(path, uid, gid):
    os.makedirs(path, mode=0o700, exist_ok=True)
    safe_chmod(path, 0o700)
    safe_chown(path, uid, gid)


def read_settings(path):
    if not os.path.exists(path):
        return {}
    try:
        with open(path, "r", encoding="utf-8") as handle:
            data = json.load(handle)
        return data if isinstance(data, dict) else {}
    except Exception:
        timestamp = time.strftime("%Y%m%d%H%M%S")
        backup_path = f"{path}.bak.{timestamp}"
        os.rename(path, backup_path)
        print(f"  Backed up invalid settings to {backup_path}")
        return {}


def write_settings(path, data, uid, gid):
    tmp_path = f"{path}.tmp"
    with open(tmp_path, "w", encoding="utf-8") as handle:
        json.dump(data, handle, indent=2, ensure_ascii=True)
        handle.write("\n")
    os.replace(tmp_path, path)
    safe_chmod(path, 0o600)
    safe_chown(path, uid, gid)


def patch_user(user_home, user_name, uid, gid, config):
    """Patch settings for a single user."""
    settings_dir = os.path.join(user_home, ".claude")
    settings_path = os.path.join(settings_dir, "settings.json")

    ensure_dir(settings_dir, uid, gid)
    data = read_settings(settings_path)

    env = data.get("env")
    if not isinstance(env, dict):
        env = {}

    env["ANTHROPIC_AUTH_TOKEN"] = config["api_key"]
    env["ANTHROPIC_BASE_URL"] = config["base_url"]
    env.pop("ANTHROPIC_MODEL", None)
    if "timeout_ms" in config and config["timeout_ms"] is not None:
        env["API_TIMEOUT_MS"] = str(config["timeout_ms"])
    if config.get("models"):
        env["CLAUDE_CUSTOM_MODELS"] = ",".join(config["models"])
    if config.get("default_sonnet_model"):
        env["ANTHROPIC_DEFAULT_SONNET_MODEL"] = config["default_sonnet_model"]
    if config.get("default_opus_model"):
        env["ANTHROPIC_DEFAULT_OPUS_MODEL"] = config["default_opus_model"]
    env["DISABLE_AUTOUPDATER"] = "1"
    env["DISABLE_TELEMETRY"] = "1"
    env["DISABLE_ERROR_REPORTING"] = "1"
    env["CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC"] = "1"
    env["CLAUDE_CODE_DISABLE_FEEDBACK_SURVEY"] = "1"
    env["CLAUDE_CODE_EFFORT_LEVEL"] = "medium"

    data["env"] = env
    data["model"] = config["model"]
    data.pop("models", None)
    data.pop("availableModels", None)
    data["effortLevel"] = "medium"

    theme = config.get("theme")
    if theme:
        data["theme"] = theme

    if config.get("complete_onboarding"):
        data["hasCompletedOnboarding"] = True
        version = config.get("target_version") or "2.1.50"
        data["lastOnboardingVersion"] = {
            "ISSUES_EXPLAINER": "report the issue at https://github.com/anthropics/claude-code/issues",
            "PACKAGE_URL": "@anthropic-ai/claude-code",
            "README_URL": "https://code.claude.com/docs/en/overview",
            "VERSION": version,
        }

    perms = data.get("permissions")
    if not isinstance(perms, dict):
        perms = {}
    perms["defaultMode"] = "bypassPermissions"
    data["permissions"] = perms

    write_settings(settings_path, data, uid, gid)

    # --- settings.local.json ---
    local_path = os.path.join(settings_dir, "settings.local.json")
    local_data = read_settings(local_path)
    local_perms = local_data.get("permissions")
    if not isinstance(local_perms, dict):
        local_perms = {}
    local_perms["defaultMode"] = "bypassPermissions"

    base_allow = [
        "Bash", "Edit", "Write", "Read", "Glob", "Grep",
        "NotebookEdit", "WebFetch", "WebSearch",
        "mcp__memory-bank__list_projects",
        "mcp__memory-bank__list_project_files",
        "mcp__memory-bank__memory_bank_read",
        "mcp__memory-bank__memory_bank_write",
        "mcp__memory-bank__memory_bank_update",
    ]
    bash_cmds = [
        "git", "python3", "python", "node", "npm", "npx", "bash", "sh",
        "ls", "cat", "wc", "ln", "cp", "mv", "rm", "mkdir", "chmod",
        "chown", "tail", "head", "touch", "tee", "echo", "printf",
        "date", "sleep", "sort", "uniq", "tr", "cut", "xargs", "find",
        "grep", "sed", "awk", "jq", "diff", "curl", "wget", "tar",
        "gzip", "gunzip", "unzip", "sha256sum", "md5sum", "du", "df",
        "free", "ps", "kill", "whoami", "hostname", "uname", "go",
        "make", "systemctl", "journalctl", "docker", "docker-compose",
        "ssh", "scp", "rsync", "pip", "pip3", "gh", "claude", "entire",
        "cd",
    ]
    for cmd in bash_cmds:
        base_allow.append(f"Bash({cmd}:*)")

    existing_allow = local_perms.get("allow", [])
    existing_set = set(existing_allow)
    for item in base_allow:
        if item not in existing_set:
            existing_allow.append(item)
    local_perms["allow"] = existing_allow
    if "deny" not in local_perms:
        local_perms["deny"] = []
    if "ask" not in local_perms:
        local_perms["ask"] = []

    local_data["permissions"] = local_perms
    write_settings(local_path, local_data, uid, gid)

    return settings_path


def discover_users():
    """Find all users with home directories."""
    if IS_WINDOWS or not HAS_PWD:
        home = os.path.expanduser("~")
        try:
            username = os.getlogin()
        except Exception:
            username = os.environ.get("USER", "user")
        # Return as simple namespace
        class User:
            def __init__(self, name, home, uid, gid):
                self.name = name
                self.home = home
                self.uid = uid
                self.gid = gid
        return [User(username, home, os.getuid() if not IS_WINDOWS else 0, os.getgid() if not IS_WINDOWS else 0)]

    users = []
    for entry in pwd.getpwall():
        if entry.pw_uid < 500 and entry.pw_name != "root":
            continue
        if entry.pw_shell in ("/usr/sbin/nologin", "/bin/false", "/sbin/nologin"):
            continue
        if not os.path.isdir(entry.pw_dir):
            continue
        class User:
            def __init__(self, name, home, uid, gid):
                self.name = name
                self.home = home
                self.uid = uid
                self.gid = gid
        users.append(User(entry.pw_name, entry.pw_dir, entry.pw_uid, entry.pw_gid))

    return users


def patch_all_users(config):
    """Patch settings for all discovered users."""
    users = discover_users()
    if not users:
        eprint(f"  {Y}No users found{D}")
        return

    for user in users:
        try:
            path = patch_user(user.home, user.name, user.uid, user.gid, config)
            print(f"  {G}Patched {user.name}{D}: {path}")
        except Exception as e:
            eprint(f"  {R}Failed to patch {user.name}: {e}{D}")

    # Windows extras
    if IS_WINDOWS:
        _set_user_env_windows(config)


def _set_user_env_windows(config):
    """Set user-level environment variables via setx (Windows only)."""
    env_vars = {
        "ANTHROPIC_BASE_URL": config["base_url"],
        "ANTHROPIC_AUTH_TOKEN": config["api_key"],
        "DISABLE_TELEMETRY": "1",
        "DISABLE_ERROR_REPORTING": "1",
        "DISABLE_AUTOUPDATER": "1",
        "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1",
        "CLAUDE_CODE_DISABLE_FEEDBACK_SURVEY": "1",
    }
    if config.get("models"):
        env_vars["CLAUDE_CUSTOM_MODELS"] = ",".join(config["models"])
    if config.get("timeout_ms"):
        env_vars["API_TIMEOUT_MS"] = str(config["timeout_ms"])

    for key, val in env_vars.items():
        try:
            subprocess.run(["setx", key, val], capture_output=True, check=True)
            os.environ[key] = val
        except Exception:
            pass


# ============================================================
#  Main commands
# ============================================================

def cmd_check():
    """Check for updates without installing."""
    installed, cli_js = get_installed_version()
    latest = get_latest_version()

    print(f"\n{W}=== UClaude Update Check ==={D}")
    print(f"  Installed: {installed or 'not found'}")
    print(f"  Latest:    {latest or 'unknown'}")

    if not installed:
        print(f"  {Y}Claude Code not found. Install it first: npm install -g @anthropic-ai/claude-code{D}")
        return 1

    if not latest:
        print(f"  {R}Cannot determine latest version. Run 'git pull' first.{D}")
        return 1

    if ver_tuple(latest) > ver_tuple(installed):
        print(f"  {Y}Update available: {installed} → {latest}{D}")
        return 0
    else:
        print(f"  {G}Up to date.{D}")
        return 0


def cmd_update(force=False, settings_only=False):
    """Full update: git pull → install cli.js → patch settings."""
    installed, cli_js = get_installed_version()

    print(f"\n{W}=== UClaude Updater ==={D}")

    # Git pull to get latest artifacts
    print(f"\n  Pulling latest updates...")
    git_pull()

    latest = get_latest_version()
    print(f"  Installed: {installed or 'not found'}")
    print(f"  Latest:    {latest or 'unknown'}")

    if not latest:
        eprint(f"  {R}Cannot determine latest version.{D}")
        return 1

    needs_update = force or not installed or ver_tuple(latest) > ver_tuple(installed)

    if not needs_update and not settings_only:
        print(f"\n  {G}Already up to date.{D}")
        # Still patch settings in case config changed
        config = load_config()
        if config:
            print(f"\n{W}--- Patching settings ---{D}")
            patch_all_users(config)
        return 0

    # Install cli.js
    if not settings_only:
        if not cli_js:
            cli_js = find_cli_js()
        if not cli_js:
            eprint(f"  {R}Claude Code cli.js not found. Install Claude Code first.{D}")
            return 1

        if not is_admin():
            eprint(f"  {R}Root/admin privileges required to update cli.js.{D}")
            eprint(f"  Run with: sudo python3 {sys.argv[0]}")
            return 1

        print(f"\n{W}--- Installing cli.js v{latest} ---{D}")
        ok = install_cli_js(latest, cli_js)
        if not ok:
            return 1

    # Patch settings
    config = load_config()
    if config:
        print(f"\n{W}--- Patching settings ---{D}")
        patch_all_users(config)
    else:
        eprint(f"  {Y}No config found, skipping settings patch{D}")

    # Verify
    new_ver, _ = get_installed_version()
    print(f"\n{W}=== Done ==={D}")
    print(f"  Version: {new_ver or 'unknown'}")
    print(f"  {G}Update complete.{D}")
    return 0


def main():
    parser = argparse.ArgumentParser(
        description="UClaude Updater — automatic Claude Code patch updater.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    parser.add_argument("--check", action="store_true", help="Only check for updates")
    parser.add_argument("--force", action="store_true", help="Force update even if version matches")
    parser.add_argument("--settings-only", action="store_true", help="Only patch settings, don't touch cli.js")

    args = parser.parse_args()

    if args.check:
        return cmd_check()
    else:
        return cmd_update(force=args.force, settings_only=args.settings_only)


if __name__ == "__main__":
    raise SystemExit(main())