
import os
import re
import json

BASE = "/opt/agent-jae-cli"

# ============================================================
# Ordered replacement pairs (most specific first!)
# ============================================================
REPLACEMENTS = [
    # === Package scope & names (must be first - most specific) ===
    ("@mariozechner/pi-coding-agent", "@jaeswift/jae-coding-agent"),
    ("@mariozechner/pi-agent-core", "@jaeswift/jae-agent-core"),
    ("@mariozechner/pi-web-ui", "@jaeswift/jae-web-ui"),
    ("@mariozechner/pi-ai", "@jaeswift/jae-ai"),
    ("@mariozechner/pi-tui", "@jaeswift/jae-tui"),
    ("@mariozechner/pi-mom", "@jaeswift/jae-mom"),
    ("@mariozechner/pi", "@jaeswift/jae"),  # pods package name

    # === GitHub / repo references ===
    ("badlogic/pi-mono", "jae/agent-jae-cli"),

    # === npm registry URL ===
    ("registry.npmjs.org/@mariozechner/pi-coding-agent", "registry.npmjs.org/@jaeswift/jae-coding-agent"),

    # === Monorepo name ===
    ("pi-monorepo", "jae-monorepo"),
    ("pi-web-ui-example", "jae-web-ui-example"),

    # === piConfig block ===
    ('"piConfig"', '"jaeConfig"'),
    ("'piConfig'", "'jaeConfig'"),
    ("piConfig", "jaeConfig"),  # catch any other references

    # === configDir ===
    ('"configDir": ".pi"', '"configDir": ".jae"'),
    ("configDir: \".pi\"", 'configDir: ".jae"'),

    # === Environment variables ===
    ("PI_PACKAGE_DIR", "JAE_PACKAGE_DIR"),
    ("PI_SKIP_VERSION_CHECK", "JAE_SKIP_VERSION_CHECK"),
    ("PI_OFFLINE", "JAE_OFFLINE"),
    ("PI_STARTUP_BENCHMARK", "JAE_STARTUP_BENCHMARK"),

    # === Binary names in bin fields ===
    ('"pi-ai": ', '"jae-ai": '),
    ('"pi-pods": ', '"jae-pods": '),

    # === Build script: bun outfile ===
    ("--outfile dist/pi ", "--outfile dist/jae "),
    ("--outfile dist/pi\n", "--outfile dist/jae\n"),

    # === Temp file prefixes ===
    ("pi-clipboard-", "jae-clipboard-"),
    ("pi-editor-", "jae-editor-"),

    # === Workspace / script file references ===
    ("pi-mono.code-workspace", "jae-mono.code-workspace"),
    ("pi-test.sh", "jae-test.sh"),
    ("pi-mono", "jae-mono"),  # catch any remaining pi-mono refs

    # === process.title ===
    ('process.title = "pi"', 'process.title = "jae"'),
    ("process.title = 'pi'", "process.title = 'jae'"),

    # === System prompt - specific phrases ===
    ("operating inside pi, a coding agent harness. You help users by reading files, executing commands, editing code, and writing new files.",
     "Agent JAE, an expert AI coding assistant. You help users write, edit, debug, and understand code."),
    ("Pi documentation (read only when the user asks about pi itself",
     "Agent JAE documentation (read only when the user asks about Agent JAE itself"),
    ("pi itself, its SDK", "Agent JAE itself, its SDK"),
    ("pi topics", "Agent JAE topics"),
    ("pi .md files", "Agent JAE .md files"),
    ("pi packages (docs/packages.md)", "Agent JAE packages (docs/packages.md)"),
    ("about pi itself", "about Agent JAE itself"),
    ("When working on pi", "When working on Agent JAE"),
    ("read pi .md", "read Agent JAE .md"),
    ("read the pi docs", "read the Agent JAE docs"),

    # === settings references to .pi/ paths in help text ===
    (".pi/settings.json", ".jae/settings.json"),
    (".pi/extensions", ".jae/extensions"),
    (".pi/prompts", ".jae/prompts"),
    (".pi/skills", ".jae/skills"),

    # === tmux message mentioning Pi ===
    ("Pi works best with", "Agent JAE works best with"),
]

# Files/dirs to skip
SKIP_DIRS = {"node_modules", ".git", "dist", "package-lock.json"}
TARGET_EXTS = {".ts", ".json", ".md", ".sh", ".mjs", ".js", ".mts"}

def should_process(filepath):
    """Check if file should be processed."""
    # Skip package-lock.json - will be regenerated
    if os.path.basename(filepath) == "package-lock.json":
        return False
    _, ext = os.path.splitext(filepath)
    return ext in TARGET_EXTS

def apply_replacements(content):
    """Apply all replacements in order."""
    for old, new in REPLACEMENTS:
        content = content.replace(old, new)
    return content

def process_file(filepath):
    """Process a single file."""
    try:
        with open(filepath, "r", encoding="utf-8", errors="ignore") as f:
            original = f.read()
    except Exception as e:
        print(f"  SKIP (read error): {filepath}: {e}")
        return False

    modified = apply_replacements(original)

    if modified != original:
        with open(filepath, "w", encoding="utf-8") as f:
            f.write(modified)
        return True
    return False

def walk_and_process(base_dir):
    """Walk directory tree and process all matching files."""
    changed = 0
    total = 0
    for root, dirs, files in os.walk(base_dir):
        # Filter out skip dirs
        dirs[:] = [d for d in dirs if d not in SKIP_DIRS]
        for fname in files:
            filepath = os.path.join(root, fname)
            if should_process(filepath):
                total += 1
                if process_file(filepath):
                    changed += 1
                    print(f"  CHANGED: {filepath[len(base_dir):]}")  
    print(f"\nProcessed {total} files, changed {changed} files")

# === Handle bin field "pi": carefully in coding-agent package.json ===
def fix_coding_agent_bin():
    """Fix the bin field in coding-agent/package.json - needs special regex."""
    pkg_path = os.path.join(BASE, "packages/coding-agent/package.json")
    with open(pkg_path, "r") as f:
        content = f.read()
    # Replace "pi": "dist/cli.js" in bin section
    # This is tricky - the bin field has "pi": "dist/cli.js"
    content = content.replace('"pi": "dist/cli.js"', '"jae": "dist/cli.js"')
    with open(pkg_path, "w") as f:
        f.write(content)
    print("  FIXED: coding-agent bin field pi -> jae")

print("=== Starting rebrand ===")
print(f"Base: {BASE}")
print()

# Step 1: Apply all text replacements
print("--- Applying text replacements ---")
walk_and_process(BASE)

# Step 2: Fix coding-agent bin field specifically
print("\n--- Fixing coding-agent bin field ---")
fix_coding_agent_bin()

# Step 3: Rename .pi directory to .jae
pi_dir = os.path.join(BASE, ".pi")
jae_dir = os.path.join(BASE, ".jae")
if os.path.isdir(pi_dir):
    os.rename(pi_dir, jae_dir)
    print(f"  RENAMED: .pi -> .jae")

# Step 4: Rename pi-mono.code-workspace
old_ws = os.path.join(BASE, "pi-mono.code-workspace")
new_ws = os.path.join(BASE, "jae-mono.code-workspace")
if os.path.exists(old_ws):
    os.rename(old_ws, new_ws)
    print(f"  RENAMED: pi-mono.code-workspace -> jae-mono.code-workspace")

# Step 5: Rename pi-test.sh
old_test = os.path.join(BASE, "pi-test.sh")
new_test = os.path.join(BASE, "jae-test.sh")
if os.path.exists(old_test):
    os.rename(old_test, new_test)
    print(f"  RENAMED: pi-test.sh -> jae-test.sh")

print("\n=== Rebrand complete ===")
