#!/usr/bin/env python3
"""
Qwen Code Patcher — patches QwenCode CLI to route through custom AI proxy.

Targets:
  1. telemetry_flag       — force getTelemetryEnabled() → false
  2. telemetry_log_prompts — force getTelemetryLogPromptsEnabled() → false
  3. telemetry_init_guard  — early return in initializeTelemetry()
  4. dashscope_base_url    — DEFAULT_DASHSCOPE_BASE_URL → proxy
  5. coding_plan_urls      — coding.dashscope.aliyuncs.com → proxy
  6. default_model         — validate DEFAULT_QWEN_MODEL (no change)
  7. mainline_model        — validate MAINLINE_CODER_MODEL (no change)
  8. auto_update_registry  — registry.npmjs.org → npm.sensey24.ru
  9. auto_update_command   — add --registry to update commands
 10. user_settings         — ~/.qwen/settings.json (auth + telemetry)
 11. trusted_folders       — ~/.qwen/trustedFolders.json
 12. system_env            — env vars injection
"""

import json
import os
import re
import sys
import shutil
import platform
import argparse
from pathlib import Path

# ─── Constants ──────────────────────────────────────────────────────────

SCRIPT_DIR = Path(__file__).resolve().parent
CONFIG_PATH = SCRIPT_DIR / "qwen_config.json"

IS_WINDOWS = platform.system() == "Windows"
IS_MACOS = platform.system() == "Darwin"

NPM_PACKAGE = "@qwen-code/qwen-code"
CLI_JS_FILENAME = "cli.js"

PATCH_MARKER = "/* QWEN_PATCHED */"

# ANSI colors
GREEN = "\033[92m"
YELLOW = "\033[93m"
RED = "\033[91m"
CYAN = "\033[96m"
BOLD = "\033[1m"
RESET = "\033[0m"


# ─── Utilities ──────────────────────────────────────────────────────────

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


def load_config(config_path=None):
    """Load patcher configuration from JSON file."""
    path = Path(config_path) if config_path else CONFIG_PATH
    if not path.is_file():
        eprint(f"{RED}Config not found: {path}{RESET}")
        sys.exit(1)
    with open(path, "r") as f:
        return json.load(f)


def read_version(qwen_root):
    """Read version from package.json."""
    pkg = Path(qwen_root) / "package.json"
    if pkg.is_file():
        with open(pkg) as f:
            return json.load(f).get("version", "unknown")
    return "unknown"


# ─── Detection ──────────────────────────────────────────────────────────

def _candidate_paths():
    """Generate candidate paths for QwenCode CLI installation."""
    if IS_WINDOWS:
        appdata = os.environ.get("APPDATA", "")
        if appdata:
            yield Path(appdata) / "npm" / "node_modules" / "@qwen-code" / "qwen-code"
        localappdata = os.environ.get("LOCALAPPDATA", "")
        if localappdata:
            yield Path(localappdata) / "npm" / "node_modules" / "@qwen-code" / "qwen-code"
    else:
        yield Path("/usr/lib/node_modules/@qwen-code/qwen-code")
        yield Path("/usr/local/lib/node_modules/@qwen-code/qwen-code")
        if IS_MACOS:
            yield Path("/opt/homebrew/lib/node_modules/@qwen-code/qwen-code")
        home = Path.home()
        yield home / ".local" / "lib" / "node_modules" / "@qwen-code" / "qwen-code"
        yield home / ".npm-global" / "lib" / "node_modules" / "@qwen-code" / "qwen-code"
        # nvm
        nvm_dir = os.environ.get("NVM_DIR", str(home / ".nvm"))
        nvm_path = Path(nvm_dir)
        if nvm_path.is_dir():
            for ver_dir in sorted(nvm_path.glob("versions/node/v*"), reverse=True):
                yield ver_dir / "lib" / "node_modules" / "@qwen-code" / "qwen-code"


def detect_qwen():
    """Find QwenCode CLI installation. Returns (qwen_root, cli_js_path) or (None, None)."""
    for root in _candidate_paths():
        cli_js = root / CLI_JS_FILENAME
        if cli_js.is_file():
            return str(root), str(cli_js)
    return None, None


# ─── Patching: cli.js targets (1-9) ────────────────────────────────────

def _backup_file(filepath):
    """Create backup of a file."""
    backup = filepath + ".backup"
    if not os.path.isfile(backup):
        shutil.copy2(filepath, backup)
        print(f"  {CYAN}Backup:{RESET} {backup}")


def _already_patched(content):
    """Check if file already has patch marker."""
    return PATCH_MARKER in content


def patch_cli_js(cli_js_path, config):
    """Apply all 9 cli.js patch targets. Returns dict of {target: status}."""
    _backup_file(cli_js_path)

    with open(cli_js_path, "r", encoding="utf-8", errors="replace") as f:
        content = f.read()

    if _already_patched(content):
        print(f"  {GREEN}cli.js already patched (marker found){RESET}")
        return {"cli_js": "already_patched"}

    base_url = config["base_url"]
    npm_registry = config.get("npm_registry", "https://npm.sensey24.ru")
    results = {}
    original = content

    # Target 1: TELEMETRY_FLAG
    pat1 = r'(getTelemetryEnabled\(\)\s*\{)\s*return\s+this\.telemetrySettings\.enabled\s*\?\?\s*false;'
    rep1 = r'\1 return false; /* QWEN_PATCHED */'
    content, n = re.subn(pat1, rep1, content)
    results["telemetry_flag"] = f"OK ({n})" if n > 0 else "SKIP"

    # Target 2: TELEMETRY_LOG_PROMPTS
    pat2 = r'(getTelemetryLogPromptsEnabled\(\)\s*\{)\s*return\s+this\.telemetrySettings\.logPrompts\s*\?\?\s*true;'
    rep2 = r'\1 return false; /* QWEN_PATCHED */'
    content, n = re.subn(pat2, rep2, content)
    results["telemetry_log_prompts"] = f"OK ({n})" if n > 0 else "SKIP"

    # Target 3: TELEMETRY_INIT_GUARD
    pat3 = r'(function initializeTelemetry\(config2\)\s*\{)\s*\n(\s*)if\s*\(telemetryInitialized'
    rep3 = r'\1\n\2return; /* QWEN_PATCHED: telemetry disabled */\n\2if (telemetryInitialized'
    content, n = re.subn(pat3, rep3, content)
    results["telemetry_init_guard"] = f"OK ({n})" if n > 0 else "SKIP"

    # Target 4: DASHSCOPE_BASE_URL — main + regional (HK/intl/us) endpoints
    count4 = 0
    for old_url in [
        "https://dashscope.aliyuncs.com/compatible-mode/v1",
        "https://cn-hongkong.dashscope.aliyuncs.com/compatible-mode/v1",
        "https://dashscope-intl.aliyuncs.com/compatible-mode/v1",
        "https://dashscope-us.aliyuncs.com/compatible-mode/v1",
    ]:
        new_url = f"{base_url}/v1"
        c = content.count(old_url)
        content = content.replace(old_url, new_url)
        count4 += c
    results["dashscope_base_url"] = f"OK ({count4})" if count4 > 0 else "SKIP"

    # Target 5: CODING_PLAN_URLS (string replace, not regex)
    count5 = 0
    for old_url in [
        "https://coding.dashscope.aliyuncs.com/v1",
        "https://coding-intl.dashscope.aliyuncs.com/v1",
    ]:
        new_url = f"{base_url}/v1"
        c = content.count(old_url)
        content = content.replace(old_url, new_url)
        count5 += c
    results["coding_plan_urls"] = f"OK ({count5})" if count5 > 0 else "SKIP"

    # Target 6: DEFAULT_MODEL (validate only)
    if re.search(r'DEFAULT_QWEN_MODEL\s*=\s*"coder-model"', content):
        results["default_model"] = "OK (validated)"
    else:
        results["default_model"] = "WARN (unexpected value)"

    # Target 7: MAINLINE_MODEL (validate only)
    if re.search(r'MAINLINE_CODER_MODEL\s*=\s*"qwen3\.5-plus"', content):
        results["mainline_model"] = "OK (validated)"
    else:
        results["mainline_model"] = "WARN (unexpected value)"

    # Target 8: AUTO_UPDATE_REGISTRY
    count8 = 0
    old_reg = '"https://registry.npmjs.org/"'
    new_reg = f'"{npm_registry}/"'
    c = content.count(old_reg)
    content = content.replace(old_reg, new_reg)
    count8 += c
    # Also handle single-quoted variant
    old_reg_sq = "'https://registry.npmjs.org/'"
    new_reg_sq = f"'{npm_registry}/'"
    c = content.count(old_reg_sq)
    content = content.replace(old_reg_sq, new_reg_sq)
    count8 += c
    results["auto_update_registry"] = f"OK ({count8})" if count8 > 0 else "SKIP"

    # Target 9: AUTO_UPDATE_COMMAND
    old_cmd = '"npm install -g @qwen-code/qwen-code@latest"'
    new_cmd = f'"npm install -g @qwen-code/qwen-code@latest --registry {npm_registry}"'
    c = content.count(old_cmd)
    content = content.replace(old_cmd, new_cmd)
    results["auto_update_command"] = f"OK ({c})" if c > 0 else "SKIP"

    # Add patch marker after shebang line (preserve shebang on line 1)
    if content != original:
        if content.startswith("#!"):
            first_nl = content.index("\n")
            content = content[:first_nl + 1] + PATCH_MARKER + "\n" + content[first_nl + 1:]
        else:
            content = PATCH_MARKER + "\n" + content

    # Write patched file
    with open(cli_js_path, "w", encoding="utf-8") as f:
        f.write(content)

    return results


# ─── Patching: settings targets (10-12) ────────────────────────────────

def patch_user_settings(config):
    """Configure ~/.qwen/settings.json (Target 10)."""
    qwen_dir = Path.home() / ".qwen"
    qwen_dir.mkdir(parents=True, exist_ok=True)
    settings_path = qwen_dir / "settings.json"

    existing = {}
    if settings_path.is_file():
        try:
            with open(settings_path) as f:
                existing = json.load(f)
        except (json.JSONDecodeError, OSError):
            pass

    # Deep merge
    if "security" not in existing:
        existing["security"] = {}
    if "auth" not in existing["security"]:
        existing["security"]["auth"] = {}
    existing["security"]["auth"]["selectedType"] = "openai"

    if "telemetry" not in existing:
        existing["telemetry"] = {}
    existing["telemetry"]["enabled"] = False
    existing["telemetry"]["logPrompts"] = False

    if "model" not in existing:
        existing["model"] = {}
    existing["model"]["name"] = config.get("default_model", "qwen3.5-plus")

    with open(settings_path, "w") as f:
        json.dump(existing, f, indent=2)

    print(f"  {GREEN}Settings:{RESET} {settings_path}")
    return "OK"


def patch_trusted_folders(config):
    """Create/update ~/.qwen/trustedFolders.json (Target 11)."""
    qwen_dir = Path.home() / ".qwen"
    qwen_dir.mkdir(parents=True, exist_ok=True)
    tf_path = qwen_dir / "trustedFolders.json"

    existing = {}
    if tf_path.is_file():
        try:
            with open(tf_path) as f:
                existing = json.load(f)
        except (json.JSONDecodeError, OSError):
            pass

    trust_paths = config.get("trust_paths", ["/home", "/root", "/tmp"])
    for p in trust_paths:
        if p not in existing:
            existing[p] = "TRUST_PARENT"

    # Also trust home directory
    home = str(Path.home())
    if home not in existing:
        existing[home] = "TRUST_PARENT"

    with open(tf_path, "w") as f:
        json.dump(existing, f, indent=2)

    print(f"  {GREEN}Trusted folders:{RESET} {tf_path}")
    return "OK"


def setup_env_vars(config):
    """Set environment variables (Target 12)."""
    base_url = config["base_url"]
    api_key = config.get("api_key", "")
    default_model = config.get("default_model", "qwen3.5-plus")

    env_vars = {
        "OPENAI_API_KEY": api_key,
        "OPENAI_BASE_URL": f"{base_url}/v1",
        "OPENAI_MODEL": default_model,
        "GEMINI_TELEMETRY_ENABLED": "false",
        "GEMINI_TELEMETRY_LOG_PROMPTS": "false",
    }

    if IS_WINDOWS:
        import subprocess
        for k, v in env_vars.items():
            subprocess.run(["setx", k, v], capture_output=True)
        print(f"  {GREEN}Env vars:{RESET} Set via setx (Windows)")
        return "OK"

    # Linux/macOS: write to /etc/environment
    env_file = Path("/etc/environment")
    if not env_file.is_file():
        # Try creating it
        try:
            env_file.touch()
        except PermissionError:
            eprint(f"  {YELLOW}Cannot write /etc/environment (no root){RESET}")
            _print_env_export(env_vars)
            return "MANUAL"

    try:
        existing = env_file.read_text()
    except PermissionError:
        _print_env_export(env_vars)
        return "MANUAL"

    lines = existing.splitlines()
    updated = False

    for key, value in env_vars.items():
        found = False
        for i, line in enumerate(lines):
            if line.startswith(f"{key}="):
                lines[i] = f'{key}="{value}"'
                found = True
                break
        if not found:
            lines.append(f'{key}="{value}"')
            updated = True

    new_content = "\n".join(lines)
    if not new_content.endswith("\n"):
        new_content += "\n"

    try:
        env_file.write_text(new_content)
        print(f"  {GREEN}Env vars:{RESET} Written to /etc/environment")
        return "OK"
    except PermissionError:
        eprint(f"  {YELLOW}Cannot write /etc/environment (no root){RESET}")
        _print_env_export(env_vars)
        return "MANUAL"


def _print_env_export(env_vars):
    """Print export commands for manual setup."""
    print(f"\n  {YELLOW}Add these to your shell profile:{RESET}")
    for k, v in env_vars.items():
        print(f'  export {k}="{v}"')
    print()


# ─── Orchestration ─────────────────────────────────────────────────────

def apply_all_patches(cli_js_path, config, settings_only=False):
    """Apply all patches. Returns overall results dict."""
    results = {}

    if not settings_only:
        print(f"\n{BOLD}Patching cli.js...{RESET}")
        cli_results = patch_cli_js(cli_js_path, config)
        results.update(cli_results)

    print(f"\n{BOLD}Configuring settings...{RESET}")
    results["user_settings"] = patch_user_settings(config)
    results["trusted_folders"] = patch_trusted_folders(config)
    results["system_env"] = setup_env_vars(config)

    return results


def rollback(cli_js_path):
    """Restore cli.js from backup."""
    backup = cli_js_path + ".backup"
    if os.path.isfile(backup):
        shutil.copy2(backup, cli_js_path)
        print(f"  {GREEN}Restored:{RESET} {cli_js_path}")
        return True
    else:
        eprint(f"  {RED}No backup found:{RESET} {backup}")
        return False


# ─── Validation (standalone) ───────────────────────────────────────────

def run_validation(cli_js_path):
    """Lightweight validation — checks patch markers and key invariants in
    cli.js + ~/.qwen/settings.json + trustedFolders.json."""
    counts = {"green": 0, "yellow": 0, "red": 0}
    print()
    print(f"  {BOLD}Validating Qwen patches{RESET}")
    print(f"  {'─' * 56}")

    def report(name, status, detail=""):
        if status == "GREEN":
            counts["green"] += 1
            color = GREEN
        elif status == "YELLOW":
            counts["yellow"] += 1
            color = YELLOW
        else:
            counts["red"] += 1
            color = RED
        suffix = f" — {detail}" if detail else ""
        print(f"  {color}[{status}]{RESET} {name}{suffix}")

    # cli.js checks
    if cli_js_path and os.path.isfile(cli_js_path):
        with open(cli_js_path, "r", encoding="utf-8", errors="replace") as f:
            content = f.read()

        report("cli.js patch_marker",
               "GREEN" if PATCH_MARKER in content else "RED",
               "" if PATCH_MARKER in content else "marker not found")

        report("telemetry_flag",
               "GREEN" if "getTelemetryEnabled() { return false" in content else "YELLOW")

        report("telemetry_log_prompts",
               "GREEN" if "getTelemetryLogPromptsEnabled() { return false" in content else "YELLOW")

        # dashscope_base_url: any /compatible-mode/v1 endpoint left = RED
        leftover_ds = [u for u in [
            "dashscope.aliyuncs.com/compatible-mode",
            "cn-hongkong.dashscope.aliyuncs.com/compatible-mode",
            "dashscope-intl.aliyuncs.com/compatible-mode",
            "dashscope-us.aliyuncs.com/compatible-mode",
        ] if u in content]
        report("dashscope_base_url",
               "GREEN" if not leftover_ds else "RED",
               f"unpatched: {', '.join(leftover_ds)}" if leftover_ds else "")

        leftover_cp = [u for u in [
            "coding.dashscope.aliyuncs.com",
            "coding-intl.dashscope.aliyuncs.com",
        ] if u in content]
        report("coding_plan_urls",
               "GREEN" if not leftover_cp else "RED",
               f"unpatched: {', '.join(leftover_cp)}" if leftover_cp else "")

        report("default_model",
               "GREEN" if 'DEFAULT_QWEN_MODEL = "coder-model"' in content else "YELLOW")

        report("mainline_model",
               "GREEN" if 'MAINLINE_CODER_MODEL = "qwen3.5-plus"' in content else "YELLOW")

        report("auto_update_registry",
               "GREEN" if '"https://registry.npmjs.org/"' not in content else "YELLOW")
    else:
        report("cli.js", "RED", f"not found: {cli_js_path}")

    # settings.json checks
    settings_path = Path.home() / ".qwen" / "settings.json"
    if settings_path.is_file():
        try:
            with open(settings_path, "r", encoding="utf-8") as f:
                settings = json.load(f)
            sec = settings.get("security", {})
            auth = sec.get("auth", {})
            telemetry = settings.get("telemetry", {})
            report("user_settings.auth.selectedType",
                   "GREEN" if auth.get("selectedType") == "openai" else "YELLOW",
                   f"value={auth.get('selectedType')!r}")
            report("user_settings.telemetry.enabled",
                   "GREEN" if telemetry.get("enabled") is False else "YELLOW")
        except Exception as e:
            report("user_settings", "RED", f"parse error: {e}")
    else:
        report("user_settings", "YELLOW", f"missing: {settings_path}")

    # trustedFolders.json checks
    trusted_path = Path.home() / ".qwen" / "trustedFolders.json"
    if trusted_path.is_file():
        report("trusted_folders", "GREEN", "present")
    else:
        report("trusted_folders", "YELLOW", "missing")

    # env vars
    has_api = bool(os.environ.get("OPENAI_API_KEY"))
    has_url = bool(os.environ.get("OPENAI_BASE_URL"))
    report("env.OPENAI_API_KEY",
           "GREEN" if has_api else "YELLOW",
           "" if has_api else "not set in current shell")
    report("env.OPENAI_BASE_URL",
           "GREEN" if has_url else "YELLOW",
           "" if has_url else "not set in current shell")

    print(f"  {'─' * 56}")
    print(f"  {GREEN}GREEN: {counts['green']}{RESET}  "
          f"{YELLOW}YELLOW: {counts['yellow']}{RESET}  "
          f"{RED}RED: {counts['red']}{RESET}")
    return counts, f"{counts['green']}G/{counts['yellow']}Y/{counts['red']}R"


# ─── CLI ───────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(
        description="Qwen Code Patcher — patches QwenCode CLI for custom AI proxy"
    )
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--detect", action="store_true", help="Find QwenCode CLI installation")
    group.add_argument("--apply", action="store_true", help="Apply all patches")
    group.add_argument("--settings-only", action="store_true", help="Only settings + env (no cli.js)")
    group.add_argument("--rollback", action="store_true", help="Restore from backup")
    group.add_argument("--validate", action="store_true", help="Validate all targets")
    parser.add_argument("--config", type=str, help="Path to custom config file")

    args = parser.parse_args()
    config = load_config(args.config)

    # Detection
    qwen_root, cli_js_path = detect_qwen()

    if args.detect:
        if qwen_root:
            version = read_version(qwen_root)
            print(f"\n  {GREEN}Found QwenCode CLI{RESET}")
            print(f"  Root:    {qwen_root}")
            print(f"  cli.js:  {cli_js_path}")
            print(f"  Version: {version}")
        else:
            eprint(f"\n  {RED}QwenCode CLI not found{RESET}")
            sys.exit(1)
        return

    if not qwen_root:
        eprint(f"{RED}QwenCode CLI not found. Install: npm install -g @qwen-code/qwen-code{RESET}")
        sys.exit(1)

    version = read_version(qwen_root)
    print(f"\n{BOLD}QwenCode CLI v{version}{RESET} — {qwen_root}")

    if args.validate:
        counts, summary = run_validation(cli_js_path)
        # Save report
        report_dir = SCRIPT_DIR / "reports"
        report_dir.mkdir(exist_ok=True)
        report_path = report_dir / f"validation_{version}.json"
        with open(report_path, "w") as f:
            json.dump(summary, f, indent=2)
        print(f"\n  Report saved: {report_path}")
        if counts.get("RED", 0) > 0:
            sys.exit(2)
        return

    if args.rollback:
        rollback(cli_js_path)
        return

    if args.settings_only:
        results = apply_all_patches(cli_js_path, config, settings_only=True)
    else:
        results = apply_all_patches(cli_js_path, config, settings_only=False)

    # Print summary
    print(f"\n{BOLD}Results:{RESET}")
    for target, status in results.items():
        if "OK" in str(status) or status == "already_patched":
            print(f"  {GREEN}[OK]{RESET}  {target}: {status}")
        elif "SKIP" in str(status):
            print(f"  {YELLOW}[SKIP]{RESET}  {target}: {status}")
        else:
            print(f"  {CYAN}[INFO]{RESET}  {target}: {status}")

    print(f"\n{GREEN}Done!{RESET} Restart QwenCode CLI to apply changes.\n")


if __name__ == "__main__":
    main()