#!/usr/bin/env python3
"""claude-sub: run two or more Claude subscriptions in Claude Code with ONE shared workspace.

Every account is an official Claude Code login. There is no proxy and no token handling: each
sign-in goes through Anthropic's own browser flow, the pattern Anthropic permits ("signing in to
the unmodified Claude Code binary with their own Claude subscription",
https://code.claude.com/docs/en/legal-and-compliance).

How it works
  account a   your normal login in ~/.claude (plain `claude`, or `claude-a`)
  account b   a second login in ~/.claude-b (`claude-b`); any short name works (c, work, ...)

~/.claude-<name> is a "shadow" config dir. Every entry in it is a symlink back into ~/.claude
(CLAUDE.md, settings.json, skills, commands, agents, hooks, plugins, and projects = memory and
chat history) EXCEPT the per-login state in NEVER_SHARE (credentials, the background daemon,
telemetry, Claude Code's own .claude.json backups). Claude Code keeps the login for each config
dir in its own slot (macOS: keychain item "Claude Code-credentials-<sha256(dir)[:8]>"; Linux:
<dir>/.credentials.json), so each dir is a separate account.

Commands
  claude-sub setup <name>       create the account dir, shell commands and status line (safe to re-run)
  claude-sub login <name>       sign <name> in (opens the browser: pick the RIGHT account there)
  claude-sub logout <name>      sign <name> out
  claude-sub status             accounts, emails, last-seen 5h / 7d use, which one `pick` prefers
  claude-sub doctor [name]      health check (read-only)
  claude-sub pick               print the account with the most room left
  claude-sub exec <name> [--] <claude args>   run claude on that account (scripts, cron, agents)
  claude-sub sync <name>        refresh the links (the shell commands run this at every launch)
  claude-sub statusline         the status line command (setup wires it; wraps any existing one)
  claude-sub uninstall <name>   undo setup for <name> (nothing is hard-deleted)

Shell commands that setup adds: claude-a, claude-<name> for each account, claude-auto.
When a limit hits mid-task: /exit, then `claude-b -c` continues the same chat on account b.
Needs: macOS or Linux (WSL counts), Python 3.8+, Claude Code 2.x, zsh or bash for the shell commands.
"""
import hashlib
import json
import os
import platform
import re
import shutil
import subprocess
import sys
import time

HOME = os.path.expanduser("~")
MAIN = os.path.join(HOME, ".claude")
MAIN_JSON = os.path.join(HOME, ".claude.json")
STATE = os.path.join(os.environ.get("XDG_STATE_HOME") or os.path.join(HOME, ".local", "state"), "claude-sub")
USAGE = os.path.join(STATE, "usage")
CONFIG = os.path.join(STATE, "config.json")
BACKUPS = os.path.join(STATE, "backups")
MARKER = ".claude-sub.json"
NAME_RE = re.compile(r"^[a-z0-9][a-z0-9_-]{0,15}$")
RC_BEGIN = "# >>> claude-sub >>>"
RC_END = "# <<< claude-sub <<<"
IS_MAC = platform.system() == "Darwin"
SELF = os.path.realpath(os.path.abspath(sys.argv[0])) if sys.argv and sys.argv[0] else "claude-sub"

# Per-login state: never linked into a shadow dir. Each login keeps its own copy.
NEVER_SHARE = {
    ".credentials.json", ".device-keys.json", "hfi-auth.json", ".session_ingress_token",  # tokens
    "policy-limits.json", "remote-settings.json",  # org policy fetched per login
    "statsig", "telemetry",  # per-user feature flags and event queues
    "daemon", "daemon.log", "jobs", "bridge-spawn",  # background daemon + Remote Control run on one login
    "backups",  # Claude Code writes .claude.json backups here; b's must never mix with a's
    "mcp-needs-auth-cache.json", "stats-cache.json", "usage-data", "shares",
    ".claude.json", ".claude.json.backup", ".claude.json.lock",
    MARKER, ".diverged", ".DS_Store", ".last-cleanup",
}
# Belt and braces: anything credential-shaped that a future Claude Code adds is never linked.
CRED_RE = re.compile(r"(credential|auth|token|secret|device-key|\.key$|\.pem$)", re.I)

# .claude.json keys copied once when an account dir is created (no identity, no consents).
SEED_KEYS = ("hasCompletedOnboarding", "lastOnboardingVersion", "installMethod", "autoUpdates",
             "autoUpdatesProtectedForNative", "shiftEnterKeyBindingInstalled", "theme",
             "hasCompletedClaudeInChromeOnboarding", "claudeInChromeDefaultEnabled",
             "githubRepoPaths", "lastReleaseNotesSeen")
# Per-folder keys kept in step on every sync (folder trust + folder-scoped MCP servers).
PROJECT_KEYS = ("hasTrustDialogAccepted", "mcpServers", "enabledMcpjsonServers",
                "disabledMcpjsonServers", "allowedTools", "mcpContextUris",
                "hasClaudeMdExternalIncludesApproved", "hasClaudeMdExternalIncludesWarningShown")


# ------------------------------------------------------------------ helpers

def die(msg, code=1):
    print(f"claude-sub: {msg}", file=sys.stderr)
    sys.exit(code)


def check_name(name, allow_main=False):
    if name == "a":
        if allow_main:
            return name
        die("account a is your main login in ~/.claude; pick another name for the second account (like b)")
    if not NAME_RE.match(name or ""):
        die(f"bad account name {name!r}: use a short lowercase name like b, c or work")
    return name


def shadow_dir(name):
    # The exact spelling matters: Claude Code hashes this string to name the login slot.
    return os.path.join(HOME, f".claude-{name}")


def config_json(name):
    return MAIN_JSON if name == "a" else os.path.join(shadow_dir(name), ".claude.json")


def keychain_service(name):
    if name == "a":
        return "Claude Code-credentials"
    return "Claude Code-credentials-" + hashlib.sha256(shadow_dir(name).encode()).hexdigest()[:8]


def read_json(path, default=None):
    try:
        with open(path) as fh:
            return json.load(fh)
    except Exception:
        return default


def write_json_atomic(path, data, mode=0o600):
    tmp = f"{path}.claude-sub.{os.getpid()}.tmp"
    fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, mode)
    with os.fdopen(fd, "w") as fh:
        json.dump(data, fh, indent=2)
        fh.write("\n")
    os.replace(tmp, path)


def write_in_place(path, text):
    # In place, not rename: rc files and settings.json are often symlinks into a dotfiles repo.
    with open(path, "w") as fh:
        fh.write(text)


def backup(path):
    if not os.path.exists(path):
        return None
    os.makedirs(BACKUPS, mode=0o700, exist_ok=True)
    dst = os.path.join(BACKUPS, f"{os.path.basename(path)}.{time.strftime('%Y%m%d-%H%M%S')}")
    shutil.copy2(path, dst)
    return dst


def accounts():
    """a first, then every account dir this tool made (marker present)."""
    names = ["a"]
    for entry in sorted(os.listdir(HOME)):
        if entry.startswith(".claude-") and os.path.isfile(os.path.join(HOME, entry, MARKER)):
            n = entry[len(".claude-"):]
            if NAME_RE.match(n) and n != "a":
                names.append(n)
    return names


def login_of(name):
    oa = (read_json(config_json(name), {}) or {}).get("oauthAccount") or {}
    return oa.get("emailAddress"), oa.get("organizationType") or oa.get("billingType")


def slot_present(name):
    """Is a saved login there for this account? Reads metadata only, never the secret."""
    if IS_MAC:
        try:
            return subprocess.call(["security", "find-generic-password", "-s", keychain_service(name)],
                                   stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) == 0
        except Exception:
            return False
    d = MAIN if name == "a" else shadow_dir(name)
    p = os.path.join(d, ".credentials.json")
    return os.path.isfile(p) and not os.path.islink(p)


def shared(entry):
    return entry not in NEVER_SHARE and not CRED_RE.search(entry)


def ready(name):
    return bool(login_of(name)[0])


def claude_bin():
    return shutil.which("claude")


def claude_version():
    b = claude_bin()
    if not b:
        return None
    try:
        out = subprocess.run([b, "--version"], capture_output=True, text=True, timeout=30).stdout
        m = re.search(r"(\d+)\.(\d+)\.(\d+)", out)
        return tuple(int(x) for x in m.groups()) if m else None
    except Exception:
        return None


# ------------------------------------------------------------------ sync

def sync(name, quiet=False):
    check_name(name)
    if not os.path.isdir(MAIN):
        die(f"{MAIN} not found: run `claude` once and sign in with your main account first")
    d = shadow_dir(name)
    os.makedirs(d, mode=0o700, exist_ok=True)
    notes = []
    stamp = time.strftime("%Y%m%d-%H%M%S")

    for entry in sorted(os.listdir(MAIN)):
        src = os.path.join(MAIN, entry)
        dst = os.path.join(d, entry)
        if not shared(entry):
            # Undo any earlier link of per-login state (it would hand a's login to this account).
            if os.path.islink(dst) and os.path.realpath(dst).startswith(os.path.realpath(MAIN) + os.sep):
                os.unlink(dst)
                notes.append(f"removed a shared link to per-login item {entry}")
            continue
        if os.path.islink(dst):
            if os.readlink(dst) == src:
                continue
            os.unlink(dst)
        elif os.path.lexists(dst):
            # A real file or folder where a link belongs. Keep it (never lose data), then relink
            # so both logins see one workspace again.
            div = os.path.join(d, ".diverged")
            os.makedirs(div, mode=0o700, exist_ok=True)
            kept = os.path.join(div, f"{entry}.{stamp}")
            shutil.move(dst, kept)
            notes.append(f"WARNING {entry} had split off from ~/.claude; this account's copy is kept at {kept}")
        os.symlink(src, dst)

    for entry in os.listdir(d):  # links whose ~/.claude target is gone
        p = os.path.join(d, entry)
        if os.path.islink(p) and os.readlink(p).startswith(MAIN + os.sep) and not os.path.lexists(os.readlink(p)):
            os.unlink(p)

    marker = os.path.join(d, MARKER)
    if not os.path.exists(marker):
        write_json_atomic(marker, {"name": name, "main": MAIN, "created": stamp,
                                   "keychain_service": keychain_service(name) if IS_MAC else None})

    # .claude.json: identity stays per login; MCP servers, folder trust and refused API keys follow a.
    a = read_json(MAIN_JSON, {}) or {}
    bpath = config_json(name)
    b = read_json(bpath)
    fresh = b is None
    if fresh:
        b = {k: a[k] for k in SEED_KEYS if k in a}
    before = json.dumps(b, sort_keys=True)
    if "mcpServers" in a:
        merged = dict(b.get("mcpServers") or {})
        merged.update(a["mcpServers"])
        b["mcpServers"] = merged
    rejected = list((b.get("customApiKeyResponses") or {}).get("rejected") or [])
    for k in (a.get("customApiKeyResponses") or {}).get("rejected") or []:
        if k not in rejected:
            rejected.append(k)
    if rejected:
        car = dict(b.get("customApiKeyResponses") or {})
        car["rejected"] = rejected
        car.setdefault("approved", [])
        b["customApiKeyResponses"] = car
    projects = dict(b.get("projects") or {})
    for path, pa in (a.get("projects") or {}).items():
        pb = dict(projects.get(path) or {})
        for k in PROJECT_KEYS:
            if k in (pa or {}):
                pb[k] = pa[k]
        if pb:
            projects[path] = pb
    if projects:
        b["projects"] = projects
    if fresh or json.dumps(b, sort_keys=True) != before:
        write_json_atomic(bpath, b)
        notes.append("created .claude.json" if fresh else "updated .claude.json (MCP servers, folder trust, refused API keys)")

    for n in notes:
        if not quiet or n.startswith("WARNING"):
            print(f"claude-sub sync {name}: {n}", file=sys.stderr if quiet else sys.stdout)
    return d


# ------------------------------------------------------------------ usage, pick, status

def record_usage(acct, rl):
    snap = {}
    for key in ("five_hour", "seven_day"):
        w = (rl or {}).get(key) or {}
        if isinstance(w.get("used_percentage"), (int, float)):
            snap[key] = {"used": round(float(w["used_percentage"]), 1), "resets_at": w.get("resets_at")}
    if not snap:
        return
    path = os.path.join(USAGE, f"{acct}.json")
    old = read_json(path, {}) or {}
    if all(old.get(k) == v for k, v in snap.items()) and time.time() - old.get("at", 0) < 300:
        return
    snap["at"] = int(time.time())
    os.makedirs(USAGE, exist_ok=True)
    tmp = f"{path}.{os.getpid()}.tmp"
    with open(tmp, "w") as fh:
        json.dump(snap, fh)
    os.replace(tmp, path)


def usage(name):
    """Last-seen use for an account; a window whose reset time has passed counts as 0."""
    snap = read_json(os.path.join(USAGE, f"{name}.json"), {}) or {}
    now = time.time()
    out = {}
    for key in ("five_hour", "seven_day"):
        w = snap.get(key) or {}
        used, resets = w.get("used"), w.get("resets_at")
        if isinstance(resets, (int, float)) and resets > 1e12:
            resets = resets / 1000.0
        if not isinstance(used, (int, float)):
            out[key] = (None, None)
        elif isinstance(resets, (int, float)) and resets <= now:
            out[key] = (0.0, None)
        else:
            out[key] = (float(used), resets)
    return out, snap.get("at")


def pick():
    candidates = [n for n in accounts() if ready(n)] or ["a"]
    scored = []
    for i, n in enumerate(candidates):
        u, _ = usage(n)
        five = u["five_hour"][0] or 0.0
        seven = u["seven_day"][0] or 0.0
        blocked = five >= 95 or seven >= 99
        scored.append((blocked, -(100 - max(five, seven)), i, n))
    scored.sort()
    return scored[0][3]


def fmt_when(ts):
    if not ts:
        return ""
    t, today = time.localtime(ts), time.localtime()
    hm = time.strftime("%I:%M %p", t).lstrip("0")
    return hm if (t.tm_yday, t.tm_year) == (today.tm_yday, today.tm_year) else time.strftime("%a ", t) + hm


def fmt_ago(ts):
    if not ts:
        return "no usage seen yet"
    s = max(0, int(time.time() - ts))
    ago = f"{s // 60}m ago" if s < 3600 else (f"{s // 3600}h ago" if s < 172800 else f"{s // 86400}d ago")
    return f"seen {ago}"


def status():
    emails = {}
    for n in accounts():
        email, plan = login_of(n)
        emails[n] = email
        u, at = usage(n)
        parts = []
        for key, tag in (("five_hour", "5h"), ("seven_day", "7d")):
            used, resets = u[key]
            if used is None:
                parts.append(f"{tag} ?")
            else:
                parts.append(f"{tag} {int(used)}%" + (f" (resets {fmt_when(resets)})" if resets else ""))
        who = f"{email} [{plan}]" if email else f"NOT SIGNED IN (run: claude-sub login {n})"
        where = "~/.claude" if n == "a" else f"~/.claude-{n}"
        print(f"{n.upper()}  {where:<12} {who}")
        print(f"    {'   '.join(parts)}   {fmt_ago(at)}")
    seen = [e for e in emails.values() if e]
    if len(seen) != len(set(seen)):
        print("WARNING: two accounts are signed in to the same Claude account. "
              "Run `claude-sub logout <name>`, then `claude-sub login <name>` and pick the other account in the browser.")
    print(f"pick: {pick()}")


# ------------------------------------------------------------------ login / logout / exec

def account_env(name):
    env = dict(os.environ)
    env.pop("CLAUDE_SECURESTORAGE_CONFIG_DIR", None)
    if name == "a":
        env.pop("CLAUDE_CONFIG_DIR", None)
    else:
        env["CLAUDE_CONFIG_DIR"] = shadow_dir(name)
    return env


def login(name, email=None):
    check_name(name)
    if not claude_bin():
        die("the `claude` command is not on PATH")
    d = sync(name, quiet=True)
    env = account_env(name)
    has_auth = subprocess.run([claude_bin(), "auth", "--help"], capture_output=True, text=True, env=env).returncode == 0
    print(f"Signing in account {name} (its own login slot: {keychain_service(name) if IS_MAC else d + '/.credentials.json'}).")
    print("When the browser opens, look at the email on the page. It must be the account you want as")
    print(f"'{name}'. If it shows your main account, switch account there, or paste the printed link")
    print("into a private (incognito) window.")
    if has_auth:
        rc = subprocess.call([claude_bin(), "auth", "login", "--claudeai"] + (["--email", email] if email else []), env=env)
    else:
        print(f"\nThis Claude Code has no `claude auth login`. Run `claude update`, or sign in by hand:\n"
              f"  CLAUDE_CONFIG_DIR={d} claude     then type /login")
        return 1
    got, plan = login_of(name)
    main_email, _ = login_of("a")
    if rc != 0 or not got:
        die(f"sign-in did not finish (exit {rc}). Try again: claude-sub login {name}")
    if IS_MAC and not slot_present(name):
        die("the login did not land in its own keychain slot, so this Claude Code version may share one login\n"
            "between folders. Run `claude update`, then `claude-sub logout " + name + "` and log in again.")
    if main_email and got == main_email:
        die(f"account {name} signed in as {got}, the SAME account as a.\n"
            f"Fix: claude-sub logout {name}; then claude-sub login {name} and pick the other account in the browser.")
    print(f"OK: account {name} = {got} [{plan}]. Start it with: claude-{name}")
    return 0


def logout(name):
    check_name(name)
    return subprocess.call([claude_bin() or "claude", "auth", "logout"], env=account_env(name))


def exec_account(name, args):
    check_name(name, allow_main=True)
    if args and args[0] == "--":
        args = args[1:]
    b = claude_bin()
    if not b:
        die("the `claude` command is not on PATH")
    if name != "a":
        sync(name, quiet=True)
        if not (args and args[0] == "auth") and not ready(name):
            die(f"account {name} is not signed in yet. Run: claude-sub login {name}")
    env = account_env(name)
    if any(x in ("-p", "--print") for x in args):
        # In -p mode Claude Code bills an exported ANTHROPIC_API_KEY before the login. Drop it so the plan pays.
        env.pop("ANTHROPIC_API_KEY", None)
    os.execvpe(b, [b] + list(args), env)


# ------------------------------------------------------------------ status line

def current_account():
    base = os.path.basename(os.environ.get("CLAUDE_CONFIG_DIR", "").rstrip("/"))
    return base[len(".claude-"):] if base.startswith(".claude-") else "a"


def statusline(wrapped=False):
    raw = ""
    try:
        raw = sys.stdin.read()
        d = json.loads(raw) if raw.strip() else {}
    except Exception:
        d = {}
    acct = current_account()
    try:
        record_usage(acct, d.get("rate_limits") or {})
    except Exception:
        pass
    dim, yel, bold, rst = "\033[2m", "\033[33m", "\033[1m", "\033[0m"
    tag = f"{dim}[A]{rst}" if acct == "a" else f"{yel}{bold}[{acct.upper()}]{rst}"
    if wrapped:
        old = ((read_json(CONFIG, {}) or {}).get("wrapped_statusline") or {}).get("command")
        if old:
            try:
                p = subprocess.run(old, shell=True, input=raw, capture_output=True, text=True, timeout=10)
                lines = p.stdout.rstrip("\n").split("\n")
                lines[0] = f"{tag} {lines[0]}"
                print("\n".join(lines))
                return 0
            except Exception:
                pass
    model = ((d.get("model") or {}).get("display_name") or (d.get("model") or {}).get("id") or "")
    parts = [tag]
    if model:
        parts.append(model)
    ctx = (d.get("context_window") or {}).get("used_percentage")
    if isinstance(ctx, (int, float)):
        parts.append(f"ctx {int(ctx)}%")
    rl = d.get("rate_limits") or {}
    lim = []
    for key, t in (("five_hour", "5h"), ("seven_day", "7d")):
        v = (rl.get(key) or {}).get("used_percentage")
        if isinstance(v, (int, float)):
            lim.append(f"{t} {int(v)}%")
    if lim:
        parts.append(" ".join(lim))
    print(" | ".join(parts))
    return 0


def settings_path():
    return os.path.join(MAIN, "settings.json")


def install_statusline():
    sp = settings_path()
    raw = open(sp).read() if os.path.exists(sp) else "{}"
    try:
        s = json.loads(raw) if raw.strip() else {}
    except Exception:
        print(f"status line: {sp} is not plain JSON, so it was left alone. To add it by hand, set\n"
              f'  "statusLine": {{"type": "command", "command": "{SELF} statusline"}}')
        return
    cur = s.get("statusLine") or {}
    if "claude-sub" in str(cur.get("command", "")) and "statusline" in str(cur.get("command", "")):
        print("status line: already wired")
        return
    backup(sp)
    cfg = read_json(CONFIG, {}) or {}
    if cur.get("command"):
        cfg["wrapped_statusline"] = cur
        write_json_atomic(CONFIG, cfg)
        s["statusLine"] = {"type": "command", "command": f"'{SELF}' statusline --wrapped", "padding": cur.get("padding", 0)}
        print("status line: your existing status line now shows [A] / [B] in front and records usage")
    else:
        s["statusLine"] = {"type": "command", "command": f"'{SELF}' statusline", "padding": 0}
        print("status line: added (shows [A] / [B], model, context, 5h / 7d use)")
    write_in_place(sp, json.dumps(s, indent=2) + "\n")


def restore_statusline():
    sp = settings_path()
    s = read_json(sp)
    if not isinstance(s, dict):
        return
    cur = s.get("statusLine") or {}
    if "claude-sub" not in str(cur.get("command", "")):
        return
    backup(sp)
    cfg = read_json(CONFIG, {}) or {}
    if cfg.get("wrapped_statusline"):
        s["statusLine"] = cfg.pop("wrapped_statusline")
        write_json_atomic(CONFIG, cfg)
    else:
        s.pop("statusLine", None)
    write_in_place(sp, json.dumps(s, indent=2) + "\n")
    print("status line: restored")


# ------------------------------------------------------------------ shell block

def default_rc():
    sh = os.path.basename(os.environ.get("SHELL", ""))
    if sh == "zsh":
        return os.path.join(HOME, ".zshrc")
    if sh == "bash":
        return os.path.join(HOME, ".bash_profile" if IS_MAC else ".bashrc")
    return None


def rc_block(names):
    bindir = os.path.dirname(SELF)
    lines = [RC_BEGIN,
             "# Two or more Claude subscriptions, one workspace. Managed by claude-sub; remove with `claude-sub uninstall <name>`.",
             "# claude / claude-a = main login (~/.claude). claude-<name> = that login (~/.claude-<name>).",
             "# claude-auto = the account with the most room left. Limit hit? /exit, then `claude-b -c`."]
    if bindir not in os.environ.get("PATH", "").split(os.pathsep):
        lines.append(f'export PATH="{bindir}:$PATH"')
    lines += [
        "_claude_sub_print() { local x; for x in \"$@\"; do case \"$x\" in -p|--print) return 0;; esac; done; return 1; }",
        "claude-a() { (unset CLAUDE_CONFIG_DIR; _claude_sub_print \"$@\" && unset ANTHROPIC_API_KEY; claude \"$@\"); }",
    ]
    for n in names:
        if n != "a":
            lines.append(f"claude-{n}() {{ _claude_sub_run {n} \"$@\"; }}")
    lines += [
        "claude-auto() {",
        "  local who; who=$(claude-sub pick 2>/dev/null) || who=a",
        "  echo \"claude-auto: account $who\" >&2",
        "  if [ \"$who\" = a ]; then claude-a \"$@\"; else _claude_sub_run \"$who\" \"$@\"; fi",
        "}",
        "_claude_sub_run() {",
        "  local who=\"$1\"; shift",
        "  claude-sub sync \"$who\" -q || return 1",
        "  if [ \"$1\" != auth ] && ! claude-sub ready \"$who\"; then",
        "    echo \"Account $who is not signed in yet. Run: claude-sub login $who\" >&2; return 1",
        "  fi",
        "  (export CLAUDE_CONFIG_DIR=\"$HOME/.claude-$who\"; _claude_sub_print \"$@\" && unset ANTHROPIC_API_KEY; claude \"$@\")",
        "}",
        RC_END,
    ]
    return "\n".join(lines) + "\n"


def strip_block(text):
    return re.sub(re.escape(RC_BEGIN) + r".*?" + re.escape(RC_END) + r"\n?", "", text, flags=re.S)


def install_rc(rc, names):
    old = open(rc).read() if os.path.exists(rc) else ""
    new = strip_block(old)
    if new and not new.endswith("\n"):
        new += "\n"
    new += ("\n" if new else "") + rc_block(names)
    if new == old:
        print(f"shell: {rc} already up to date")
        return
    backup(rc)
    write_in_place(rc, new)
    print(f"shell: added claude-a, {', '.join('claude-' + n for n in names if n != 'a')}, claude-auto to {rc} "
          "(open a new terminal tab to use them)")


def remove_rc(rc):
    if not rc or not os.path.exists(rc):
        return
    old = open(rc).read()
    new = strip_block(old)
    if new != old:
        backup(rc)
        write_in_place(rc, new)
        print(f"shell: removed the claude-sub block from {rc}")


# ------------------------------------------------------------------ setup / uninstall / doctor

def setup(name, rc=None, shell=True, line=True):
    check_name(name)
    if os.environ.get("CLAUDE_CONFIG_DIR"):
        die("CLAUDE_CONFIG_DIR is set in this shell. claude-sub treats ~/.claude as the main account;\n"
            "run setup from a shell without CLAUDE_CONFIG_DIR (or unset it first).")
    if not claude_bin():
        die("the `claude` command is not on PATH. Install Claude Code first.")
    v = claude_version()
    if v and v < (2, 0, 0):
        die(f"Claude Code {'.'.join(map(str, v))} is too old. Run `claude update` first.")
    if not ready("a"):
        die("your main account is not signed in (no login found in ~/.claude.json). Run `claude` and /login first.")
    sync(name)
    print(f"account dir: {shadow_dir(name)} (links back into ~/.claude; its login stays separate)")
    if shell:
        rc = rc or default_rc()
        if rc:
            install_rc(rc, accounts())
        else:
            print("shell: not zsh or bash, so no shell commands were added. Paste this into your shell config:\n")
            print(rc_block(accounts()))
    if line:
        install_statusline()
    if os.environ.get("ANTHROPIC_API_KEY"):
        print("note: ANTHROPIC_API_KEY is set in your shell. In -p (headless) mode Claude Code bills that key before your\n"
              "      plan. claude-a -p / claude-b -p / claude-sub exec drop it for that run so the plan pays.")
    print(f"\nNext (the only manual step): open a NEW terminal tab and run:  claude-sub login {name}")
    return 0


def uninstall(name, rc=None):
    check_name(name)
    d = shadow_dir(name)
    if not os.path.isfile(os.path.join(d, MARKER)):
        die(f"{d} is not a claude-sub account dir")
    if ready(name):
        logout(name)
    moved = f"{d}.removed-{time.strftime('%Y%m%d-%H%M%S')}"
    os.rename(d, moved)
    print(f"account dir moved to {moved} (delete it yourself when you are sure)")
    left = accounts()
    rc = rc or default_rc()
    if len(left) <= 1:
        remove_rc(rc)
        restore_statusline()
    elif rc and os.path.exists(rc) and RC_BEGIN in open(rc).read():
        install_rc(rc, left)
    if IS_MAC:
        print(f'if a saved login is left over, remove it with:  security delete-generic-password -s "{keychain_service(name)}"')
    return 0


def doctor(names):
    bad = 0

    def say(state, msg):
        nonlocal bad
        if state == "FAIL":
            bad += 1
        print(f"{state:<5} {msg}")

    b = claude_bin()
    v = claude_version()
    say("ok" if b else "FAIL", f"claude on PATH: {b or 'no'}" + (f" (v{'.'.join(map(str, v))})" if v else ""))
    a = read_json(MAIN_JSON, {}) or {}
    rc = default_rc()
    say("info", f"shell block in {rc}: {'yes' if rc and os.path.exists(rc) and RC_BEGIN in open(rc).read() else 'no'}")
    sl = str(((read_json(settings_path(), {}) or {}).get("statusLine") or {}).get("command", ""))
    say("info", f"status line records usage: {'yes' if 'claude-sub' in sl else 'no (claude-auto then has no usage data)'}")
    key = os.environ.get("ANTHROPIC_API_KEY") or ""
    main_email, _ = login_of("a")
    for n in names:
        print(f"-- account {n}")
        email, _ = login_of(n)
        say("ok" if email else "FAIL", f"signed in: {email or 'no'}")
        has_slot = slot_present(n)
        if email:
            say("ok" if has_slot else "FAIL", f"saved login slot: {'present' if has_slot else 'MISSING'}"
                + (f" ({keychain_service(n)})" if IS_MAC else ""))
        cj = read_json(config_json(n), {}) or {}
        if key:
            appr = (cj.get("customApiKeyResponses") or {}).get("approved") or []
            say("ok" if key[-20:] not in appr else "FAIL", "ANTHROPIC_API_KEY not approved (sessions bill the plan, not the API)")
        if n == "a":
            continue
        d = shadow_dir(n)
        say("ok" if os.path.isfile(os.path.join(d, MARKER)) else "FAIL", f"account dir {d}")
        for entry in sorted(os.listdir(MAIN)):
            p = os.path.join(d, entry)
            if shared(entry):
                if not (os.path.islink(p) and os.readlink(p) == os.path.join(MAIN, entry)):
                    say("FAIL", f"{entry} is not linked to ~/.claude (fix: claude-sub sync {n})")
            elif os.path.islink(p) and os.path.realpath(p).startswith(os.path.realpath(MAIN) + os.sep):
                say("FAIL", f"per-login item {entry} is linked to ~/.claude (a's login state leaks into {n})")
        say("ok", "shared links checked")
        mine, theirs = set(cj.get("mcpServers") or {}), set(a.get("mcpServers") or {})
        say("ok" if mine >= theirs else "FAIL", f"MCP servers carried over ({len(mine)} of {len(theirs)})")
        if email and main_email:
            say("ok" if email != main_email else "FAIL", "a different Claude account from a")
    print("doctor: all good" if not bad else f"doctor: {bad} problem(s)")
    return 1 if bad else 0


# ------------------------------------------------------------------ main

def main(argv):
    cmd = argv[1] if len(argv) > 1 else "status"
    rest = argv[2:]

    def opt(flag):
        if flag in rest:
            i = rest.index(flag)
            if i + 1 < len(rest):
                return rest[i + 1]
        return None

    if cmd == "status":
        status()
    elif cmd == "setup":
        if not rest or rest[0].startswith("-"):
            die("usage: claude-sub setup <name> [--rc FILE] [--no-shell] [--no-statusline]")
        return setup(rest[0], rc=opt("--rc"), shell="--no-shell" not in rest, line="--no-statusline" not in rest)
    elif cmd == "sync":
        if not rest:
            die("usage: claude-sub sync <name> [-q]")
        sync(rest[0], quiet="-q" in rest[1:])
    elif cmd == "login":
        if not rest:
            die("usage: claude-sub login <name> [--email you@example.com]")
        return login(rest[0], opt("--email"))
    elif cmd == "logout":
        if not rest:
            die("usage: claude-sub logout <name>")
        return logout(rest[0])
    elif cmd == "pick":
        print(pick())
    elif cmd == "ready":
        return 0 if ready(check_name(rest[0] if rest else "a", allow_main=True)) else 1
    elif cmd == "exec":
        if not rest:
            die("usage: claude-sub exec <name> [--] <claude args>")
        exec_account(rest[0], rest[1:])
    elif cmd == "statusline":
        return statusline(wrapped="--wrapped" in rest)
    elif cmd == "doctor":
        names = [check_name(rest[0], allow_main=True)] if rest else accounts()
        return doctor(names)
    elif cmd == "uninstall":
        if not rest:
            die("usage: claude-sub uninstall <name> [--rc FILE]")
        return uninstall(rest[0], rc=opt("--rc"))
    elif cmd in ("-h", "--help", "help"):
        print(__doc__)
    else:
        die(f"unknown command {cmd!r}; try: claude-sub help")
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv))
