#!/usr/bin/env python3
"""
JSON engine for claude-mode (POSIX port).

Bash cannot safely read or rewrite JSON, and jq is not installed often enough to
depend on. Everything that touches structured data goes through here, so the
shell script only ever handles flat lines of text.

Subcommands:
  apply <settings> <state> <mode> [preset] [helper]  rewrite the managed keys
  summary <preset>                                   human lines for the UI
  models <preset>                                    "tier<TAB>model" lines
  set-tier <preset> <tier> <model>                   edit one tier in place
  scaffold <provider>                                print a blank preset
  get <file> <dotted.path>                           print one value
  presets <dir>                                      "name<TAB>provider<TAB>desc"
"""

import json
import os
import re
import sys

# Matches both the qualified gateway id (anthropic/claude-opus-5) and the bare
# internal id Claude Code persists (claude-opus-5, claude-haiku-4-5-20251001).
ANTHROPIC_MODEL_RE = re.compile(r"(^|/)claude[-.]|^anthropic/")


def is_anthropic_model(model_id):
    return bool(model_id) and bool(ANTHROPIC_MODEL_RE.search(str(model_id)))

# Must stay in lockstep with the Windows build's $script:BaseManagedEnvKeys.
BASE_MANAGED = [
    "ANTHROPIC_BASE_URL",
    "ANTHROPIC_AUTH_TOKEN",
    "ANTHROPIC_API_KEY",
    "ANTHROPIC_DEFAULT_OPUS_MODEL",
    "ANTHROPIC_DEFAULT_SONNET_MODEL",
    "ANTHROPIC_DEFAULT_HAIKU_MODEL",
    "ANTHROPIC_DEFAULT_FABLE_MODEL",
    # Both pin a concrete model outside the tier mapping. A stale value in either
    # survives a switch, and the small/fast slot is used by background
    # summarisation - a known cause of "works normally, dies on compaction".
    "ANTHROPIC_MODEL",
    "ANTHROPIC_SMALL_FAST_MODEL",
    "CLAUDE_CODE_SUBAGENT_MODEL",
    "CLAUDE_CODE_DISABLE_1M_CONTEXT",
    "CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY",
    "CLAUDE_CODE_ATTRIBUTION_HEADER",
    "CLAUDE_CODE_AUTO_COMPACT_WINDOW",
    "CLAUDE_CODE_MAX_CONTEXT_TOKENS",
    "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC",
    "API_TIMEOUT_MS",
]

TIERS = ["opus", "sonnet", "haiku", "fable"]


def load(path, default=None):
    if not os.path.exists(path):
        return default if default is not None else {}
    with open(path, "r", encoding="utf-8") as fh:
        text = fh.read().strip()
    if not text:
        return default if default is not None else {}
    return json.loads(text)


def save(path, data):
    os.makedirs(os.path.dirname(path), exist_ok=True)
    tmp = path + ".tmp"
    with open(tmp, "w", encoding="utf-8") as fh:
        json.dump(data, fh, indent=2)
        fh.write("\n")
    os.replace(tmp, path)


def cmd_apply(argv):
    settings_path, state_path, mode = argv[0], argv[1], argv[2]
    preset_path = argv[3] if len(argv) > 3 and argv[3] else None
    helper = argv[4] if len(argv) > 4 and argv[4] else None

    settings = load(settings_path)
    state = load(state_path)

    # Clear the baseline keys plus whatever the previous switch actually wrote,
    # so a preset's custom extraEnv key cannot outlive the preset that added it.
    doomed = set(BASE_MANAGED) | set(state.get("writtenEnvKeys") or [])
    env = settings.get("env")
    if isinstance(env, dict):
        for k in doomed:
            env.pop(k, None)
        if not env:
            settings.pop("env", None)
    settings.pop("apiKeyHelper", None)

    written = []

    if mode != "anthropic":
        if not preset_path:
            raise SystemExit("apply: a gateway mode needs a preset")
        preset = load(preset_path)
        if preset.get("provider") != mode:
            raise SystemExit(
                "preset declares provider '%s', not '%s'" % (preset.get("provider"), mode)
            )

        models = preset.get("models") or {}

        # Cost guard. Gateways resell Anthropic models at full list price with no
        # subscription discount, so routing a tier there is almost never intended.
        # Opt in per preset with "allowAnthropicModels": true.
        if not preset.get("allowAnthropicModels"):
            offenders = []
            for tier in TIERS:
                if is_anthropic_model(models.get(tier)):
                    offenders.append("%s -> %s" % (tier, models[tier]))
            if is_anthropic_model(preset.get("subagentModel")):
                offenders.append("subagent -> %s" % preset["subagentModel"])
            if offenders:
                raise SystemExit(
                    "refusing to switch: preset routes a tier at an Anthropic model "
                    "through '%s' (%s). Gateways bill these at full price. Add "
                    '"allowAnthropicModels": true to the preset if deliberate.'
                    % (mode, "; ".join(offenders))
                )

        block = {}
        block["ANTHROPIC_BASE_URL"] = str(preset["baseUrl"])
        # Explicitly empty, not absent: a cached Anthropic login can otherwise
        # override the gateway config and surface as a model-not-found error.
        # Removed entirely when switching back to anthropic.
        block["ANTHROPIC_API_KEY"] = ""
        for tier in TIERS:
            val = models.get(tier)
            if val:
                block["ANTHROPIC_DEFAULT_%s_MODEL" % tier.upper()] = str(val)

        if preset.get("subagentModel"):
            block["CLAUDE_CODE_SUBAGENT_MODEL"] = str(preset["subagentModel"])
        if preset.get("gatewayModelDiscovery"):
            block["CLAUDE_CODE_ENABLE_GATEWAY_MODEL_DISCOVERY"] = "1"

        # Behind a custom base URL Claude Code cannot resolve a third-party model
        # id to a context length, so it guesses low and auto-compacts early.
        ctx = preset.get("contextTokens")
        if ctx:
            block["CLAUDE_CODE_MAX_CONTEXT_TOKENS"] = str(int(ctx))
            block["CLAUDE_CODE_AUTO_COMPACT_WINDOW"] = str(int(ctx))

        for k, v in (preset.get("extraEnv") or {}).items():
            block[k] = str(v)

        auth = preset.get("auth") or {"mode": "vault", "keyRef": "openrouter"}
        if auth.get("mode") == "vault":
            if not helper:
                raise SystemExit("apply: vault auth needs the key-helper path")
            settings["apiKeyHelper"] = helper
        else:
            block["ANTHROPIC_AUTH_TOKEN"] = str(auth.get("token") or "lmstudio")

        env = settings.get("env")
        if isinstance(env, dict):
            env.update(block)
        else:
            settings["env"] = block
        written = list(block.keys())

    save(settings_path, settings)

    state["mode"] = mode
    state["preset"] = os.path.splitext(os.path.basename(preset_path))[0] if preset_path else ""
    state["writtenEnvKeys"] = written
    save(state_path, state)

    for k in written:
        print(k)


def cmd_summary(argv):
    p = load(argv[0])
    if p.get("description"):
        print(p["description"])
    models = p.get("models") or {}
    bits = ["%s=%s" % (t, models[t]) for t in TIERS if models.get(t)]
    if bits:
        print("   ".join(bits))
    ctx = p.get("contextTokens")
    if ctx:
        print("context {:,} tokens   base {}".format(int(ctx), p.get("baseUrl", "")))
    else:
        print("base %s" % p.get("baseUrl", ""))


def cmd_models(argv):
    p = load(argv[0])
    models = p.get("models") or {}
    for t in TIERS:
        print("%s\t%s" % (t, models.get(t, "")))
    print("subagent\t%s" % (p.get("subagentModel") or ""))


def cmd_set_tier(argv):
    path, tier, model = argv[0], argv[1], argv[2]
    p = load(path)
    if tier == "subagent":
        p["subagentModel"] = model
    elif tier in TIERS:
        p.setdefault("models", {})[tier] = model
    else:
        raise SystemExit("unknown tier '%s'" % tier)
    save(path, p)


def cmd_scaffold(argv):
    provider = argv[0]
    base = {"provider": provider, "description": "new preset"}
    if provider == "openrouter":
        base["baseUrl"] = "https://openrouter.ai/api"
        base["auth"] = {"mode": "vault", "keyRef": "openrouter"}
    elif provider == "zai":
        base["baseUrl"] = "https://api.z.ai/api/anthropic"
        base["auth"] = {"mode": "vault", "keyRef": "zai"}
    elif provider == "lmstudio":
        base["baseUrl"] = "http://127.0.0.1:1234"
        base["auth"] = {"mode": "literal", "token": "lmstudio"}
    else:
        raise SystemExit("unknown provider '%s'" % provider)

    base["models"] = {t: "" for t in TIERS}
    base["subagentModel"] = "inherit"
    base["gatewayModelDiscovery"] = provider == "openrouter"
    base["contextTokens"] = 262144 if provider == "lmstudio" else 1000000
    if provider == "lmstudio":
        base["extraEnv"] = {"CLAUDE_CODE_ATTRIBUTION_HEADER": "0"}
    if provider == "zai":
        base["extraEnv"] = {
            "API_TIMEOUT_MS": "3000000",
            "CLAUDE_CODE_DISABLE_NONESSENTIAL_TRAFFIC": "1",
        }
    print(json.dumps(base, indent=2))


def cmd_get(argv):
    data = load(argv[0])
    cur = data
    for part in argv[1].split("."):
        if isinstance(cur, dict) and part in cur:
            cur = cur[part]
        else:
            return
    if isinstance(cur, (dict, list)):
        print(json.dumps(cur))
    elif isinstance(cur, bool):
        print("true" if cur else "false")
    elif cur is not None:
        print(cur)


def cmd_presets(argv):
    d = argv[0]
    if not os.path.isdir(d):
        return
    for name in sorted(os.listdir(d)):
        if not name.endswith(".json"):
            continue
        try:
            p = load(os.path.join(d, name))
        except Exception:
            continue
        print("%s\t%s\t%s" % (name[:-5], p.get("provider", "openrouter"), p.get("description", "")))


def cmd_managed(argv):
    state = load(argv[0]) if argv else {}
    for k in sorted(set(BASE_MANAGED) | set(state.get("writtenEnvKeys") or [])):
        print(k)


def cmd_settings_env(argv):
    """Print 'KEY=VALUE' for every managed key currently in settings.json."""
    settings = load(argv[0])
    state = load(argv[1]) if len(argv) > 1 else {}
    keys = set(BASE_MANAGED) | set(state.get("writtenEnvKeys") or [])
    if settings.get("apiKeyHelper"):
        print("apiKeyHelper=%s" % settings["apiKeyHelper"])
    env = settings.get("env") or {}
    for k in sorted(keys):
        if k in env:
            print("%s=%s" % (k, env[k]))


def cmd_or_models(argv):
    """OpenRouter /v1/models on stdin -> 'id<TAB>ctx<TAB>$in<TAB>$out' lines."""
    data = json.loads(sys.stdin.read())
    for m in sorted(data.get("data") or [], key=lambda x: x.get("id", "")):
        pricing = m.get("pricing") or {}

        def per_m(key):
            try:
                return round(float(pricing.get(key)) * 1e6, 3)
            except (TypeError, ValueError):
                return ""

        print("%s\t%s\t%s\t%s" % (m.get("id", ""), m.get("context_length", ""), per_m("prompt"), per_m("completion")))


def cmd_lms_models(argv):
    """LM Studio /api/v0/models on stdin -> 'id<TAB>state<TAB>ctx' lines."""
    data = json.loads(sys.stdin.read())
    for m in sorted(data.get("data") or [], key=lambda x: x.get("id", "")):
        print("%s\t%s\t%s" % (m.get("id", ""), m.get("state", "unknown"), m.get("max_context_length", "")))


# A model id carrying a bracket suffix - e.g. `claude-fable-5[1m]` - is Claude
# Code's extended-context marker. It belongs to Anthropic's 1M models and no
# gateway recognises it. A session that had it can carry the tag onto a new id
# after a switch, producing something like
# `~deepseek/deepseek-v4-flash-latest[1m]` that only fails at compaction, because
# compaction re-resolves the model from session state.
TAGGED_MODEL_RE = re.compile(r"\[[0-9]+[a-zA-Z]\]$")


def walk_model_ids(node, out):
    """Collect model-ish strings anywhere in the config.

    Claude Code caches resolved models under more than one key - both
    clientDataCacheSlots and additionalModelOptionsCache have been observed - so
    scan rather than reach for a fixed path.
    """
    if isinstance(node, dict):
        for k, v in node.items():
            if isinstance(v, str):
                if k in ("model", "value") and v and (v[0].isalnum() or v[0] == "~"):
                    out.add(v)
            else:
                walk_model_ids(v, out)
    elif isinstance(node, list):
        for item in node:
            walk_model_ids(item, out)


def cmd_stale_models(argv):
    """Cached model ids worth flagging, as 'kind<TAB>id' lines.

    kind is 'anthropic' (would bill at gateway list price) or 'tagged' (carries a
    [1m]-style marker that breaks compaction on a gateway).
    """
    ids = set()
    walk_model_ids(load(argv[0], {}), ids)
    for m in sorted(i for i in ids if is_anthropic_model(i)):
        print("anthropic\t" + m)
    for m in sorted(i for i in ids if TAGGED_MODEL_RE.search(i)):
        print("tagged\t" + m)


def cmd_strip_tags(argv):
    """Remove [1m]-style suffixes from cached model ids.

    The tag is NOT junk everywhere: on an Anthropic id it is how Claude Code
    selects the 1M variant, so stripping it silently downgrades that choice to
    200k. On a gateway id the same tag is meaningless and breaks compaction.
    Default is gateway ids only; pass 'all' as argv[1] to include Anthropic ids.

    Prints 'strip<TAB>id' and 'keep<TAB>id' lines.
    """
    path = argv[0]
    strip_all = len(argv) > 1 and argv[1] == "all"

    raw = open(path, encoding="utf-8").read()
    tagged = sorted(set(re.findall(r'"([^"]*\[[0-9]+[a-zA-Z]\])"', raw)))
    if not tagged:
        return

    target = [t for t in tagged if strip_all or not is_anthropic_model(t)]
    kept = [t for t in tagged if t not in target]

    for k in kept:
        print("keep\t" + k)
    if not target:
        return

    fixed = raw
    for t in target:
        fixed = fixed.replace('"' + t + '"', '"' + re.sub(r"\[[0-9]+[a-zA-Z]\]$", "", t) + '"')
    json.loads(fixed)  # refuse to write anything that is not valid JSON
    with open(path, "w", encoding="utf-8") as fh:
        fh.write(fixed)
    for t in target:
        print("strip\t" + t)


def cmd_health(argv):
    """Write health.json - the machine-readable state a fleet reader consumes.

    usage: health <root> <claude.json> <mode> <preset> <guardrail> <keyBackend> <version>

    Contract, deliberately narrow:
      * NO key material. keysConfigured is names only.
      * Model lists are [{id, anthropic}] so a reader never re-derives the
        Anthropic matcher. A tagged *Anthropic* id is normal (that is how the 1M
        variant is selected) and must not render as a fault.
      * guardrailStatus is tri-state or null; never collapse it to a boolean.
    """
    root, cfg_path, mode, preset, guardrail, key_backend, version = (argv + [""] * 7)[:7]

    ids = set()
    walk_model_ids(load(cfg_path, {}), ids)
    try:
        raw = open(cfg_path, encoding="utf-8").read()
        ids.update(re.findall(r'"([^"]*\[[0-9]+[a-zA-Z]\])"', raw))
    except OSError:
        pass

    def entries(seq):
        return [{"id": i, "anthropic": bool(is_anthropic_model(i))} for i in sorted(seq)]

    h = {
        "schema": 1,
        "tool": "claude-mode",
        "version": version or "0.0.0",
        "updatedAt": __import__("datetime").datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%SZ"),
        "os": "posix",
        "mode": mode,
        "preset": "" if mode == "anthropic" else preset,
    }

    tagged = entries(i for i in ids if TAGGED_MODEL_RE.search(i))

    if mode == "anthropic":
        h["staleModelIds"] = []
        h["taggedModelIds"] = tagged
    else:
        p = load(os.path.join(root, "presets", preset + ".json"), {})
        h["provider"] = p.get("provider", "")
        h["baseUrl"] = p.get("baseUrl", "")
        h["models"] = {t: p.get("models", {}).get(t, "") for t in TIERS if p.get("models", {}).get(t)}
        h["subagentModel"] = p.get("subagentModel", "")
        h["contextTokens"] = int(p["contextTokens"]) if p.get("contextTokens") else None
        h["gatewayDiscovery"] = bool(p.get("gatewayModelDiscovery"))
        auth = p.get("auth") or {}
        if auth.get("mode", "vault") == "vault":
            h["keyBackend"] = key_backend or "unknown"
            h["keysConfigured"] = [auth.get("keyRef", "openrouter")]
        else:
            h["keyBackend"] = "inline"
            h["keysConfigured"] = []
        h["costGuardPassed"] = True
        h["guardrailStatus"] = guardrail or None
        h["staleModelIds"] = entries(i for i in ids if is_anthropic_model(i))
        h["taggedModelIds"] = tagged

    save(os.path.join(root, "health.json"), h)


COMMANDS = {
    "health": cmd_health,
    "stale-models": cmd_stale_models,
    "strip-tags": cmd_strip_tags,
    "or-models": cmd_or_models,
    "lms-models": cmd_lms_models,
    "apply": cmd_apply,
    "summary": cmd_summary,
    "models": cmd_models,
    "set-tier": cmd_set_tier,
    "scaffold": cmd_scaffold,
    "get": cmd_get,
    "presets": cmd_presets,
    "managed": cmd_managed,
    "settings-env": cmd_settings_env,
}

if __name__ == "__main__":
    if len(sys.argv) < 2 or sys.argv[1] not in COMMANDS:
        sys.exit("usage: cm-json.py <%s> ..." % "|".join(sorted(COMMANDS)))
    try:
        COMMANDS[sys.argv[1]](sys.argv[2:])
    except SystemExit:
        raise
    except Exception as exc:  # surface a clean message to the shell
        sys.exit("cm-json: %s" % exc)