#!/usr/bin/env python3
"""kubectl clickllm — size and plan against the cluster you are pointed at.

Installed by putting this file on PATH as `kubectl-clickllm`; kubectl discovers
any executable named `kubectl-*` and exposes it as a subcommand. No plugin
manifest, no registration, no compilation.

    kubectl clickllm nodes                       what the cluster can actually run
    kubectl clickllm fit  --context 32k          which models fit the best node
    kubectl clickllm plan -f workload.yaml       the Deployment an IW would produce

`plan` runs the same reconcile the controller runs, so what it prints is what
would be applied — a dry run that cannot drift from the real path because it is
the real path.
"""

from __future__ import annotations

import argparse
import json
import pathlib
import sys

sys.path.insert(0, __file__.rsplit("/", 1)[0] + "/src")

from clickllm.k8s.nodes import read_cluster  # noqa: E402
from clickllm.k8s.reconcile import reconcile, select_node  # noqa: E402


def _parse_size(s: str) -> int:
    s = str(s).strip().lower()
    if s.endswith("k"):
        return int(float(s[:-1]) * 1024)
    if s.endswith("m"):
        return int(float(s[:-1]) * 1024 * 1024)
    return int(s)


def cmd_nodes(args) -> int:
    nodes = read_cluster(args.context)
    if args.json:
        print(
            json.dumps(
                [
                    n.__dict__
                    if not hasattr(n, "__slots__")
                    else {f: getattr(n, f) for f in n.__slots__}
                    for n in nodes
                ],
                indent=2,
            )
        )
        return 0
    print(f"\n  {'node':<28}{'kind':<9}{'devices':>8}{'per-device':>12}  notes")
    print(f"  {'-' * 78}")
    for n in nodes:
        per = f"{n.device_bytes / 1024**3:.0f} GiB" if n.device_bytes else "—"
        note = n.unknown or n.note or ""
        print(f"  {n.name[:27]:<28}{n.kind:<9}{n.devices:>8}{per:>12}  {note[:34]}")
    unsized = [n for n in nodes if n.unknown]
    if unsized:
        print(f"\n  {len(unsized)} node(s) could not be sized. The first says:")
        print(f"    {unsized[0].unknown}")
    print()
    return 0


def cmd_fit(args) -> int:
    from clickllm import fit as fitmod

    nodes = read_cluster(args.context)
    node, why = select_node(nodes)
    if node is None:
        print(f"error: {why}", file=sys.stderr)
        return 1
    hw = node.to_hardware()
    print(f"\n  sizing against {node.name} — {why}")
    if node.note:
        print(f"  note: {node.note}")
    feasible, rejected = fitmod.rank(hw, _parse_size(args.context_len), args.concurrency)
    print(f"\n  {'model':<26}{'quant':>7}{'total':>10}{'free':>10}{'~tok/s':>9}")
    print(f"  {'-' * 62}")
    for f in feasible:
        tps = f"{f.tokens_per_sec:.0f}" if f.tokens_per_sec else "?"
        print(
            f"  {f.model.name[:25]:<26}{f.quant:>7}"
            f"{f.total_bytes / 1024**3:>9.1f}G{f.headroom_bytes / 1024**3:>9.1f}G{tps:>9}"
        )
    if rejected:
        print("\n  NOT FEASIBLE")
        for m, reason in rejected[:4]:
            print(f"  {m.name[:25]:<26}{reason[:44]}")
    print()
    return 0


def cmd_plan(args) -> int:
    text = sys.stdin.read() if args.file == "-" else pathlib.Path(args.file).read_text()
    # Minimal YAML: the CRD is small and flat, and a runtime dependency on
    # PyYAML for one file would break the zero-dependency promise. JSON is
    # accepted directly; YAML is parsed only for the shapes the CRD allows.
    try:
        obj = json.loads(text)
    except json.JSONDecodeError:
        obj = _tiny_yaml(text)

    nodes = read_cluster(args.context)
    r = reconcile(obj, nodes)

    if args.json:
        print(json.dumps({"objects": list(r.objects), "status": r.status}, indent=2))
        return 0 if r.ready else 1

    st = r.status
    print(f"\n  engine   {st.get('engine', '—')}")
    if st.get("engineReason"):
        print(f"           {st['engineReason']}")
    print(f"  node     {st.get('node', '—')}")
    if st.get("knobs"):
        print("\n  settings")
        for k in st["knobs"]:
            print(f"    {k['setting']:<18} {k['value']}")
            print(f"      {k['why'][:96]}")
    for label, key in (("gaps", "gaps"), ("warnings", "warnings")):
        if st.get(key):
            print(f"\n  {label}")
            for g in st[key]:
                print(f"    · {g[:100]}")
    for c in st.get("conditions", []):
        print(f"\n  {c['type']}={c['status']}  {c.get('reason', '')}: {c.get('message', '')[:80]}")
    print()
    return 0 if r.ready else 1


def _tiny_yaml(text: str) -> dict:
    """Parse the flat subset of YAML an InferenceWorkload actually uses.

    Deliberately small: two levels of mapping, scalars, and nothing else. A
    document this parser cannot handle is rejected loudly rather than
    misinterpreted — silently dropping a field would produce a plan for a
    workload nobody described.
    """
    out: dict = {}
    stack = [(0, out)]
    for raw in text.splitlines():
        if not raw.strip() or raw.lstrip().startswith("#") or raw.strip() == "---":
            continue
        indent = len(raw) - len(raw.lstrip())
        line = raw.strip()
        if ":" not in line:
            raise SystemExit(f"unsupported YAML line (expected 'key: value'): {line!r}")
        key, _, val = line.partition(":")
        key, val = key.strip(), val.strip().strip("\"'")
        while stack and indent < stack[-1][0]:
            stack.pop()
        parent = stack[-1][1]
        if not val:
            child: dict = {}
            parent[key] = child
            stack.append((indent + 2, child))
        else:
            if val.lower() in ("true", "false"):
                parent[key] = val.lower() == "true"
            else:
                try:
                    parent[key] = int(val)
                except ValueError:
                    try:
                        parent[key] = float(val)
                    except ValueError:
                        parent[key] = val
    return out


def main(argv=None) -> int:
    p = argparse.ArgumentParser(
        prog="kubectl clickllm", description="Size and plan inference against this cluster."
    )
    p.add_argument("--context", help="kubectl context to use")
    sub = p.add_subparsers(dest="cmd", required=True)

    n = sub.add_parser("nodes", help="what the cluster can actually run")
    n.add_argument("--json", action="store_true")
    n.set_defaults(fn=cmd_nodes)

    f = sub.add_parser("fit", help="which models fit the best node")
    f.add_argument("--context", dest="context_len", default="32k")
    f.add_argument("--concurrency", type=int, default=8)
    f.set_defaults(fn=cmd_fit)

    pl = sub.add_parser("plan", help="the Deployment an InferenceWorkload would produce")
    pl.add_argument(
        "-f",
        "--file",
        required=True,
        help="InferenceWorkload YAML or JSON, or - for stdin",
    )
    pl.add_argument("--json", action="store_true")
    pl.set_defaults(fn=cmd_plan)

    args = p.parse_args(argv)
    try:
        return args.fn(args)
    except (RuntimeError, OSError, ValueError) as e:
        print(f"error: {e}", file=sys.stderr)
        return 2


if __name__ == "__main__":
    raise SystemExit(main())
