"""Add --json alias to every register() that uses --format/dest='fmt'.

Uses muse's own semantic tooling:
  - muse code query   → find every register() in cli/commands
  - muse code cat     → read the exact source of each one
  - muse code patch   → surgically replace only the functions that need updating

No regex, no AST manipulation, no risk of touching surrounding code.
"""
from __future__ import annotations

import json
import re
import subprocess
import sys
import tempfile
import pathlib


# Match any var.add_argument("--format" ...) call that includes dest="fmt"
# but is NOT already followed by a --json counterpart for that same var.
_FORMAT_RE = re.compile(
    r'(?m)^(?P<indent>[ \t]*)(?P<var>\w+)\.add_argument\('
    r'[^)]*"--format"[^)]*dest="fmt"[^)]*\)',
)

_JSON_ALIAS_RE = re.compile(r'"--json"[^)]*dest="fmt"')


def _muse(*args: str) -> str:
    """Run a muse command and return stdout. Raises on non-zero exit."""
    result = subprocess.run(
        ["muse", *args],
        capture_output=True, text=True,
    )
    if result.returncode not in (0, 32):  # 32 = muse query success code
        raise RuntimeError(
            f"muse {' '.join(args)} failed ({result.returncode}):\n{result.stderr}"
        )
    return result.stdout


def _vars_with_json(source: str) -> set[str]:
    """Return parser variable names that already have a --json add_argument call.

    Handles both single-line and multi-line calls correctly by tracking paren depth.
    """
    found: set[str] = set()
    lines = source.splitlines(keepends=True)
    i = 0
    while i < len(lines):
        m = re.match(r'^[ \t]*(?P<var>\w+)\.add_argument\(', lines[i])
        if m:
            var = m.group("var")
            # Collect entire call body.
            call_lines = [lines[i]]
            depth = lines[i].count("(") - lines[i].count(")")
            j = i + 1
            while depth > 0 and j < len(lines):
                call_lines.append(lines[j])
                depth += lines[j].count("(") - lines[j].count(")")
                j += 1
            call_text = "".join(call_lines)
            if '"--json"' in call_text:
                found.add(var)
            i = j
            continue
        i += 1
    return found


def _add_json_aliases(source: str) -> tuple[str, int]:
    """Return (patched_source, number_of_aliases_added)."""
    count = 0
    lines = source.splitlines(keepends=True)
    insertions: list[tuple[int, str]] = []  # (line_idx_after, text_to_insert)

    # Pre-compute which parser variables already have --json defined anywhere
    # in the function body (regardless of dest). Argparse will reject duplicates.
    already_has_json: set[str] = _vars_with_json(source)

    i = 0
    while i < len(lines):
        line = lines[i]

        # Does this line start a .add_argument( call?
        m = re.match(
            r'^(?P<indent>[ \t]*)(?P<var>\w+)\.add_argument\(', line
        )
        if m:
            indent = m.group("indent")
            var = m.group("var")

            # Always collect the full call body (handles single- and multi-line).
            call_lines = [line]
            depth = line.count("(") - line.count(")")

            j = i + 1
            while depth > 0 and j < len(lines):
                call_lines.append(lines[j])
                depth += lines[j].count("(") - lines[j].count(")")
                j += 1

            call_end = j - 1  # last line of the call (0-based)
            call_text = "".join(call_lines)

            # Only act when this is a --format/dest="fmt" call and the parser
            # variable doesn't already have a --json argument registered.
            if ('"--format"' in call_text
                    and 'dest="fmt"' in call_text
                    and var not in already_has_json):
                alias = (
                    f'{indent}{var}.add_argument(\n'
                    f'{indent}    "--json", action="store_const", const="json", dest="fmt",\n'
                    f'{indent}    help="Shorthand for --format json."\n'
                    f'{indent})\n'
                )
                insertions.append((call_end, alias))
                already_has_json.add(var)
                count += 1

            i = j
            continue

        i += 1

    # Apply insertions in reverse order so line numbers stay valid.
    for line_idx, text in sorted(insertions, reverse=True):
        lines.insert(line_idx + 1, text)

    return "".join(lines), count


def main() -> None:
    print("🔍 Querying symbol graph for register() functions in cli/commands…")
    raw = _muse(
        "code", "query",
        "kind=function", "name=register", "file~=cli/commands",
        "--json",
    )
    data = json.loads(raw)
    symbols: list[dict[str, object]] = data["results"]
    print(f"   Found {len(symbols)} register() functions.\n")

    patched = 0
    skipped = 0

    for sym in sorted(symbols, key=lambda s: str(s["address"])):
        address = str(sym["address"])

        # Get the exact current source of this function.
        source = _muse("code", "cat", address)

        # Strip the header comment line muse prepends ("# file::symbol  L…").
        source_lines = source.splitlines(keepends=True)
        if source_lines and source_lines[0].startswith("#"):
            source_lines = source_lines[1:]
        source = "".join(source_lines)

        new_source, added = _add_json_aliases(source)

        if added == 0:
            skipped += 1
            continue

        # Write new body to a temp file and patch via muse.
        with tempfile.NamedTemporaryFile(
            mode="w", suffix=".py", delete=False, encoding="utf-8"
        ) as tmp:
            tmp.write(new_source)
            tmp_path = tmp.name

        result = subprocess.run(
            ["muse", "code", "patch", address, "--body", tmp_path, "--json"],
            capture_output=True, text=True,
        )
        pathlib.Path(tmp_path).unlink(missing_ok=True)

        if result.returncode == 0:
            patch_data = json.loads(result.stdout)
            print(
                f"  ✅ {address}\n"
                f"     +{added} alias(es) added  "
                f"({patch_data.get('lines_replaced')}→{patch_data.get('new_lines')} lines)"
            )
            patched += 1
        else:
            print(f"  ❌ {address}\n     {result.stderr.strip()}", file=sys.stderr)

    print(f"\n{patched} function(s) patched, {skipped} already correct.")


if __name__ == "__main__":
    main()