#!/usr/bin/env python3
"""Kairo local coding agent CLI.

This is the standalone npm-packaged copy of the Kairo agent that was embedded
in AI Workshop's Tauri terminal. It talks to a local OpenAI-compatible chat
server and can act on the current workspace through a small ReAct-style tool
loop: read, ls, write, edit, and bash.
"""

import argparse
import atexit
import json
import os
import platform
import re
import shutil
import signal
import subprocess
import sys
import threading
import time
import urllib.request


A = "\033[38;2;108;142;245m"
D = "\033[2m"
B = "\033[1m"
G = "\033[38;2;0;214;143m"
R = "\033[38;2;248;113;113m"
Y = "\033[38;2;245;166;35m"
X = "\033[0m"

BANNER = r"""
   /\    /\    __ __       _
  //\\_//\\   / //_/__ _ (_)____ ___
  \_      _/ / ,<  /  ' \/ / __// _ \
  / * * \ /_/|_| \_,_//_/_/   \___/
   \  ^  /
"""

DEFAULT_PORTS = [18081, 18082, 33115, 18080]
WORKSPACE = os.path.realpath(os.getcwd())
MAX_OUT = 4000
BASH_TIMEOUT = 600

IS_WINDOWS = os.name == "nt"

ROOT = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
TINYQ4_BIN = os.path.join(ROOT, "bin", "tinyq4.exe" if IS_WINDOWS else "tinyq4")
MODELS_DIR = os.path.join(ROOT, "models")
SERVER_LOG = os.path.join(MODELS_DIR, "tinyq4-server.log")
SPIN = "|/-\\"

# Where users can grab larger models. Set this to your site to surface it in
# `kairolite model` and the in-chat /model command. Leave empty to hide it.
BIGGER_MODELS_URL = ""

# The GPU build is a separate, much faster package (NVIDIA/CUDA). `kairolite gpu`
# (and the in-chat /gpu) scans for a capable GPU and points the user to it.
GPU_BUILD_INSTALL = "npm install -g agentkairo"
GPU_BUILD_URL = ""  # optional: your site, shown alongside the install command
GPU_MIN_VRAM_MIB = 10000  # the GPU build's 14B model needs ~10 GB

try:
    with open(os.path.join(ROOT, "package.json"), "r", encoding="utf-8") as _pkg:
        VERSION = json.load(_pkg).get("version", "")
except Exception:
    VERSION = ""

# kairolite ships a small, CPU-friendly model. Larger models are an opt-in
# download (eventually from BIGGER_MODELS_URL).
DEFAULT_MODEL = {
    "id": "qwen2.5-coder-3b-q4_k_m",
    "name": "Qwen2.5-Coder-3B-Instruct GGUF Q4_K_M",
    "repo": "Qwen/Qwen2.5-Coder-3B-Instruct-GGUF",
    "quant": "Q4_K_M",
    "file": "qwen2.5-coder-3b-instruct-q4_k_m.gguf",
    "url": "https://huggingface.co/Qwen/Qwen2.5-Coder-3B-Instruct-GGUF/resolve/main/qwen2.5-coder-3b-instruct-q4_k_m.gguf",
    "license": "Apache-2.0",
    "size_gb": 1.96,
    "params": "3.09B",
    "port": 18081,
    "source": "https://huggingface.co/Qwen/Qwen2.5-Coder-3B-Instruct-GGUF",
}

MODEL_PATH = os.path.join(MODELS_DIR, DEFAULT_MODEL["file"])

SYSTEM = """You are Kairo, a local coding agent working in: __WS__
Host OS: __OS__; the bash tool runs in __SHELL__.

To act, reply with ONE fenced json block and nothing else, e.g.:
```json
{"name": "bash", "command": "ls"}
```

Tools:
- read  {"name":"read","path":"<file>"}
- ls    {"name":"ls","path":"<dir>"}
- write {"name":"write","path":"<file>","content":"<text>"}
- edit  {"name":"edit","path":"<file>","old":"<exact>","new":"<new>"}
- bash  {"name":"bash","command":"<cmd>"}

Rules: one tool per reply; paths are relative to the directory above; use edit for
existing files (make 'old' unique) and write for new files. When the task is done,
reply normally with NO json block.
"""

FENCE = re.compile(r"```(?:json)?\s*(\{.*?\})\s*```", re.S)


def parse_chat_args(argv):
    parser = argparse.ArgumentParser(
        prog="kairo",
        description="Kairo local coding agent CLI for OpenAI-compatible local model servers.",
    )
    parser.add_argument("--url", default=os.environ.get("KAIRO_BASE_URL"),
                        help="OpenAI-compatible base URL, for example http://127.0.0.1:18081")
    parser.add_argument("--port", type=int, default=_env_int("KAIRO_PORT"),
                        help="Local OpenAI-compatible server port")
    parser.add_argument("--workspace", default=os.getcwd(),
                        help="Workspace directory. Defaults to the current directory.")
    parser.add_argument("--yolo", action="store_true",
                        help="Run write/bash tools without confirmation.")
    parser.add_argument("--max-steps", type=int, default=_env_int("KAIRO_MAX_STEPS", 25),
                        help="Maximum tool-loop steps per prompt. Defaults to 25.")
    parser.add_argument("-p", "--print", dest="oneshot", default=None,
                        help="Run a single prompt non-interactively and exit. "
                             "Use '-' to read the prompt from stdin.")
    return parser.parse_args(argv)


def parse_setup_args(argv):
    parser = argparse.ArgumentParser(
        prog="kairo setup",
        description="Check local requirements and download the default Qwen Coder model.",
    )
    parser.add_argument("--start", action="store_true",
                        help="Start tinyq4 server after checks pass.")
    parser.add_argument("-y", "--yes", action="store_true",
                        help="Do not prompt before starting a first-run model download.")
    parser.add_argument("--port", type=int, default=_env_int("KAIRO_PORT", DEFAULT_MODEL["port"]),
                        help="Port for tinyq4 server when using --start.")
    return parser.parse_args(argv)


def parse_serve_args(argv):
    parser = argparse.ArgumentParser(
        prog="kairo serve",
        description="Start the bundled tinyq4 server with Kairo's default Qwen model.",
    )
    parser.add_argument("--port", type=int, default=_env_int("KAIRO_PORT", DEFAULT_MODEL["port"]),
                        help="Server port. Defaults to 18081.")
    parser.add_argument("-y", "--yes", action="store_true",
                        help="Do not prompt before a first-run model download.")
    return parser.parse_args(argv)


def _env_int(name, default=None):
    raw = os.environ.get(name)
    if raw is None or raw == "":
        return default
    try:
        return int(raw)
    except ValueError:
        return default


def disk_free_gb(path):
    try:
        usage = shutil.disk_usage(path)
        return usage.free / (1024 ** 3)
    except Exception:
        return None


MIN_RAM_GB = 3  # the 3B Q4 model is ~2 GB; needs headroom for KV cache + OS


def total_ram_gb():
    """Total physical RAM in GB, or None if it can't be determined."""
    try:
        if hasattr(os, "sysconf") and "SC_PHYS_PAGES" in os.sysconf_names:
            return os.sysconf("SC_PHYS_PAGES") * os.sysconf("SC_PAGE_SIZE") / (1024 ** 3)
    except Exception:
        pass
    if IS_WINDOWS:
        try:
            import ctypes

            class _MS(ctypes.Structure):
                _fields_ = [("dwLength", ctypes.c_ulong), ("dwMemoryLoad", ctypes.c_ulong),
                            ("ullTotalPhys", ctypes.c_ulonglong), ("ullAvailPhys", ctypes.c_ulonglong),
                            ("ullTotalPageFile", ctypes.c_ulonglong), ("ullAvailPageFile", ctypes.c_ulonglong),
                            ("ullTotalVirtual", ctypes.c_ulonglong), ("ullAvailVirtual", ctypes.c_ulonglong),
                            ("ullAvailExtendedVirtual", ctypes.c_ulonglong)]

            stat = _MS()
            stat.dwLength = ctypes.sizeof(_MS)
            ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(stat))
            return stat.ullTotalPhys / (1024 ** 3)
        except Exception:
            return None
    return None


def check_runtime():
    """Return a list of human-readable reasons the bundled engine can't run here.

    Empty list means the machine looks capable. kairolite runs on CPU, so the
    gate is just a supported OS (Linux/Windows x86-64), the engine binary, and
    enough RAM - no GPU required.
    """
    problems = []
    if sys.platform not in ("linux", "win32"):
        problems.append("kairolite ships engines for Linux and Windows (x86-64) only; this is %s" % sys.platform)
        return problems
    if not os.path.exists(TINYQ4_BIN):
        problems.append("the tinyq4 engine binary is missing at %s" % TINYQ4_BIN)
        return problems

    # On Linux a missing shared library (e.g. too-old glibc) stops it loading.
    if sys.platform == "linux":
        try:
            out = subprocess.run(["ldd", TINYQ4_BIN], capture_output=True, text=True, timeout=10).stdout
            missing = sorted({ln.split("=>")[0].strip() for ln in out.splitlines() if "not found" in ln})
            if missing:
                problems.append("missing system libraries: %s (needs glibc >= 2.34)" % ", ".join(missing))
        except Exception:
            pass

    ram = total_ram_gb()
    if ram is not None and ram < MIN_RAM_GB:
        problems.append("only ~%.1f GB RAM detected; about %d GB is recommended for the 3B model" % (ram, MIN_RAM_GB))
    return problems


def print_model_summary():
    print("Kairo model target")
    print("  model:   %s" % DEFAULT_MODEL["name"])
    print("  repo:    %s" % DEFAULT_MODEL["repo"])
    print("  quant:   %s" % DEFAULT_MODEL["quant"])
    print("  size:    about %.2f GB" % DEFAULT_MODEL["size_gb"])
    print("  params:  %s" % DEFAULT_MODEL["params"])
    print("  license: %s" % DEFAULT_MODEL["license"])
    print("  source:  %s" % DEFAULT_MODEL["source"])


def confirm_download():
    if os.path.exists(MODEL_PATH):
        return True
    try:
        answer = input(
            "This will download about %.2f GB to %s. Continue? [y/N] "
            % (DEFAULT_MODEL["size_gb"], MODELS_DIR)
        ).strip().lower()
    except (EOFError, KeyboardInterrupt):
        print()
        return False
    return answer in ("y", "yes")


def ensure_tinyq4_executable():
    """npm does not reliably preserve the +x bit on packaged files."""
    try:
        if os.path.exists(TINYQ4_BIN):
            mode = os.stat(TINYQ4_BIN).st_mode
            os.chmod(TINYQ4_BIN, mode | 0o111)
    except Exception:
        pass


def download_model():
    if os.path.exists(MODEL_PATH):
        return True
    os.makedirs(MODELS_DIR, exist_ok=True)
    # Download to a temporary file and rename on success so an interrupted
    # download never leaves a half-written .gguf that looks complete.
    part = MODEL_PATH + ".part"
    print(f"Downloading model from {DEFAULT_MODEL['url']}...")
    print(f"Saving to {MODEL_PATH}")
    try:
        def progress(count, block_size, total_size):
            if total_size > 0:
                percent = min(100, count * block_size * 100 // total_size)
                done_gb = count * block_size / (1024 ** 3)
                sys.stdout.write(
                    "\r   Downloading: %3d%% (%.2f / %.2f GB)"
                    % (percent, done_gb, total_size / (1024 ** 3))
                )
                sys.stdout.flush()

        urllib.request.urlretrieve(DEFAULT_MODEL["url"], part, reporthook=progress)
        os.replace(part, MODEL_PATH)
        print("\nDownload complete.")
        return True
    except Exception as e:
        print(f"\nError downloading model: {e}")
        for stale in (part, MODEL_PATH):
            if os.path.exists(stale):
                try:
                    os.remove(stale)
                except Exception:
                    pass
        return False


def tinyq4_server_command(port):
    return [
        TINYQ4_BIN,
        MODEL_PATH,
        "--server", str(port),
    ]


def run_tinyq4_server(port):
    if not os.path.exists(MODEL_PATH):
        if not download_model():
            return 1
    ensure_tinyq4_executable()
    cmd = tinyq4_server_command(port)
    print("Starting tinyq4 server:")
    print("  " + " ".join(cmd))
    print()
    return subprocess.call(cmd)


def wait_for_health(base_url, timeout=240):
    """Poll until the server answers, returning its model id or None on timeout."""
    deadline = time.time() + timeout
    spin = "|/-\\"
    i = 0
    while time.time() < deadline:
        model = probe_server(base_url, require_health=True)
        if model:
            return model
        sys.stderr.write("\r   %s%s loading model...%s " % (D, spin[i % 4], X))
        sys.stderr.flush()
        i += 1
        time.sleep(1.0)
    return None


def spawn_managed_server(port):
    """Start tinyq4 as a child process we own. Returns the Popen or None."""
    if not os.path.exists(MODEL_PATH):
        print("Model not found - downloading once (about %.2f GB)..." % DEFAULT_MODEL["size_gb"])
        if not download_model():
            return None
    if not os.path.exists(TINYQ4_BIN):
        print("%serror:%s tinyq4 binary missing at %s" % (R, X, TINYQ4_BIN))
        return None
    ensure_tinyq4_executable()
    os.makedirs(MODELS_DIR, exist_ok=True)
    log = open(SERVER_LOG, "ab")
    kwargs = {"stdout": log, "stderr": subprocess.STDOUT}
    if hasattr(os, "setsid"):
        kwargs["start_new_session"] = True  # own process group so we can kill children
    try:
        return subprocess.Popen(tinyq4_server_command(port), **kwargs)
    except Exception as error:
        print("%serror starting server: %s%s" % (R, error, X))
        return None


def terminate_server(proc):
    if not proc or proc.poll() is not None:
        return
    try:
        if hasattr(os, "killpg"):
            os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
        else:
            proc.terminate()
    except Exception:
        try:
            proc.terminate()
        except Exception:
            return
    try:
        proc.wait(timeout=10)
    except Exception:
        try:
            if hasattr(os, "killpg"):
                os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
            else:
                proc.kill()
        except Exception:
            pass


def _tinyq4_pids(port=None):
    """PIDs of running tinyq4 engine processes (optionally on a given port).

    Linux only - reads /proc. On Windows kill_tinyq4 takes a separate path.
    """
    pids = []
    try:
        for entry in os.listdir("/proc"):
            if not entry.isdigit():
                continue
            try:
                with open("/proc/%s/cmdline" % entry, "rb") as handle:
                    parts = [a.decode("utf-8", "replace") for a in handle.read().split(b"\x00") if a]
            except Exception:
                continue
            if not parts or os.path.basename(parts[0]) != "tinyq4":
                continue
            if port is None:
                pids.append(int(entry))
            elif "--server" in parts:
                i = parts.index("--server")
                if i + 1 < len(parts) and parts[i + 1] == str(port):
                    pids.append(int(entry))
    except Exception:
        pass
    return pids


def _kill_tinyq4_windows():
    """Kill every tinyq4.exe via taskkill (port filtering isn't worth it here)."""
    try:
        listing = subprocess.run(["tasklist", "/FI", "IMAGENAME eq tinyq4.exe", "/NH"],
                                  capture_output=True, text=True, timeout=10).stdout
        running = "tinyq4.exe" in listing.lower()
        if not running:
            return 0
        subprocess.run(["taskkill", "/F", "/T", "/IM", "tinyq4.exe"],
                       capture_output=True, text=True, timeout=15)
        return 1
    except Exception:
        return 0


def kill_tinyq4(port=None):
    """Stop tinyq4 server processes, freeing their memory. Returns count stopped."""
    if IS_WINDOWS:
        return _kill_tinyq4_windows()
    pids = _tinyq4_pids(port)
    for pid in pids:
        try:
            os.kill(pid, signal.SIGTERM)
        except Exception:
            pass
    if pids:
        deadline = time.time() + 8
        while time.time() < deadline and _tinyq4_pids(port):
            time.sleep(0.3)
        for pid in _tinyq4_pids(port):
            try:
                os.kill(pid, signal.SIGKILL)
            except Exception:
                pass
    return len(pids)


def _is_local_url(url):
    return bool(url) and ("127.0.0.1" in url or "localhost" in url)


def _port_of_url(url):
    try:
        return int(url.rsplit(":", 1)[1].split("/")[0])
    except Exception:
        return None


def ensure_local_server(args):
    """Resolve a server for the chat session.

    Returns (base_url, model, proc). `proc` is a child process the caller must
    terminate on exit, or None when we attached to something already running.
    """
    # An explicit --url means "use this server"; never manage it.
    if args.url:
        base_url, model = find_server(args.url, args.port)
        return base_url, model, None

    # Reuse a server that is already up (e.g. one left running, or AI Workshop).
    base_url, model = find_server(None, args.port)
    if base_url:
        return base_url, model, None

    # Before spawning, make sure this machine can actually run the engine -
    # otherwise tinyq4 would crash with a cryptic missing-library error.
    problems = check_runtime()
    if problems:
        print("%scannot start the local model server on this machine:%s" % (R, X), file=sys.stderr)
        for problem in problems:
            print("  - %s" % problem, file=sys.stderr)
        print("kairolite runs on CPU and needs Linux or Windows (x86-64) with a few GB of RAM.",
              file=sys.stderr)
        print("Or point it at an existing OpenAI-compatible server with --url.", file=sys.stderr)
        return None, None, None

    # Otherwise start our own and wait for it to load.
    port = args.port or DEFAULT_MODEL["port"]
    base_url = "http://127.0.0.1:%d" % port
    print("   %s*%s starting local model server on %s" % (A, X, base_url), file=sys.stderr)
    proc = spawn_managed_server(port)
    if not proc:
        return None, None, None
    try:
        model = wait_for_health(base_url, timeout=240)
    except KeyboardInterrupt:
        terminate_server(proc)
        raise
    _clear_line(sys.stderr)
    if not model:
        terminate_server(proc)
        print("%sserver did not become healthy in time - see %s%s" % (R, SERVER_LOG, X), file=sys.stderr)
        return None, None, None
    return base_url, model, proc


def cmd_postinstall(_argv):
    """Run during `npm install`. Best-effort: must NEVER fail the install."""
    print("kairolite: preparing local model...")
    ensure_tinyq4_executable()

    # Don't pull the model onto a machine that can't run the engine.
    problems = check_runtime()
    if problems:
        print("%skairolite: this machine cannot run the bundled model engine:%s" % (Y, X))
        for problem in problems:
            print("  - %s" % problem)
        print("Skipping the model download. Requirements: Linux or Windows (x86-64),")
        print("a few GB of RAM (no GPU needed; on Linux, glibc >= 2.34).")
        print("Once that is in place, run: %skairolite model%s" % (B, X))
        return 0

    if os.path.exists(MODEL_PATH):
        print("kairolite: model already present, nothing to download.")
        return 0
    print_model_summary()
    print()
    if download_model():
        print("kairolite: ready. Run `kairolite` to start.")
    else:
        print(
            "%skairolite: model download did not finish.%s\n"
            "  It will download automatically the first time you run `kairolite`,\n"
            "  or you can fetch it now with: %skairolite model%s"
            % (Y, X, B, X)
        )
    return 0  # always succeed so `npm install` does not break


def cmd_stop(_argv):
    """Stop any running tinyq4 server(s) and free their memory."""
    count = kill_tinyq4()
    if count:
        print("%s*%s stopped the local tinyq4 server; memory freed." % (G, X))
    else:
        print("%sno running tinyq4 server found.%s" % (D, X))
    return 0


def _print_bigger_models_hint():
    if BIGGER_MODELS_URL:
        print("%sLarger models: %s%s" % (D, BIGGER_MODELS_URL, X))


_GPU_SCAN = "unset"  # cache: None = no GPU, (name, vram) = found


def detect_nvidia_gpu():
    """Return (name, vram_mib_or_None) for the first NVIDIA GPU, or None. Cached."""
    global _GPU_SCAN
    if _GPU_SCAN != "unset":
        return _GPU_SCAN
    result = None
    try:
        out = subprocess.run(
            ["nvidia-smi", "--query-gpu=name,memory.total", "--format=csv,noheader,nounits"],
            capture_output=True, text=True, timeout=10,
        )
        if out.returncode == 0 and out.stdout.strip():
            first = out.stdout.strip().splitlines()[0]
            parts = [p.strip() for p in first.split(",")]
            name = parts[0] if parts else "NVIDIA GPU"
            vram = int(parts[1]) if len(parts) > 1 and parts[1].isdigit() else None
            result = (name, vram)
    except Exception:
        result = None
    _GPU_SCAN = result
    return result


def gpu_is_capable():
    """True if an NVIDIA GPU with enough VRAM for the GPU build is present."""
    gpu = detect_nvidia_gpu()
    if not gpu:
        return False
    _name, vram = gpu
    return vram is None or vram >= GPU_MIN_VRAM_MIB


def cmd_gpu(_argv):
    """Scan for an NVIDIA GPU and, if capable, point to the faster GPU build."""
    gpu = detect_nvidia_gpu()
    if not gpu:
        print("%s*%s No NVIDIA GPU detected. kairolite (CPU) is the right build for this machine." % (D, X))
        return 0
    name, vram = gpu
    vram_str = ("%d MiB VRAM" % vram) if vram is not None else "VRAM unknown"
    print("%s*%s NVIDIA GPU detected: %s%s%s %s(%s)%s" % (G, X, B, name, X, D, vram_str, X))
    if gpu_is_capable():
        print()
        print("   This machine can run the %sGPU build%s - roughly 20x+ faster than CPU:" % (B, X))
        print("     %s%s%s" % (B, GPU_BUILD_INSTALL, X))
        if GPU_BUILD_URL:
            print("     %s%s%s" % (D, GPU_BUILD_URL, X))
        print("   %s(It downloads a larger 14B model and needs the CUDA 12 runtime.)%s" % (D, X))
    else:
        print("   %sThis GPU has under %d MiB VRAM, so the GPU build's 14B model won't fit." % (Y, GPU_MIN_VRAM_MIB))
        print("   kairolite (CPU) is the better choice here.%s" % X)
    return 0


def cmd_model(argv):
    """Show model status and download it on demand (recovery path)."""
    parser = argparse.ArgumentParser(prog="kairo model")
    parser.add_argument("-y", "--yes", action="store_true",
                        help="Download without prompting if the model is missing.")
    parser.add_argument("--status", action="store_true",
                        help="Only print status; do not download.")
    args = parser.parse_args(argv)

    print_model_summary()
    print()
    if os.path.exists(MODEL_PATH):
        try:
            size_gb = os.path.getsize(MODEL_PATH) / (1024 ** 3)
        except OSError:
            size_gb = 0.0
        print("%s*%s model present: %s (%.2f GB)" % (G, X, MODEL_PATH, size_gb))
        _print_bigger_models_hint()
        return 0

    print("%s!%s model not downloaded yet." % (Y, X))
    _print_bigger_models_hint()
    if args.status:
        print("Run `kairolite model` to download it.")
        return 0
    if not args.yes and not confirm_download():
        print("Download skipped.")
        return 1
    return 0 if download_model() else 1


def cmd_setup(argv):
    args = parse_setup_args(argv)
    print_model_summary()
    print()

    free = disk_free_gb(os.path.expanduser("~"))
    if free is None:
        print("%s!%s Could not determine free disk space." % (Y, X))
    elif free < DEFAULT_MODEL["size_gb"] + 4:
        print("%s!%s Low disk space: %.1f GB free. Keep at least %.1f GB free." % (
            Y, X, free, DEFAULT_MODEL["size_gb"] + 4
        ))
    else:
        print("%s*%s Disk space looks OK: %.1f GB free." % (G, X, free))

    if not os.path.exists(TINYQ4_BIN):
        print("%s!%s tinyq4 binary was not found at %s." % (R, X, TINYQ4_BIN))
        return 1

    print("%s*%s tinyq4 binary found." % (G, X))
    
    if os.path.exists(MODEL_PATH):
        print("%s*%s Model found: %s" % (G, X, MODEL_PATH))
    else:
        print("%s!%s Model missing: %s" % (Y, X, MODEL_PATH))

    print()

    if args.start:
        if not args.yes and not confirm_download():
            print("Setup checked. Server was not started.")
            return 0
        return run_tinyq4_server(args.port)

    print("Setup checked. To download and serve the model, run:")
    print("  kairo serve")
    print()
    return 0


def cmd_serve(argv):
    args = parse_serve_args(argv)
    if not os.path.exists(TINYQ4_BIN):
        print("%serror:%s tinyq4 binary was not found at %s." % (R, X, TINYQ4_BIN))
        return 1
    print_model_summary()
    print()
    if not args.yes and not confirm_download():
        print("Server was not started.")
        return 0
    return run_tinyq4_server(args.port)


def cmd_doctor(_argv):
    print_model_summary()
    print()
    if os.path.exists(TINYQ4_BIN):
        print("%s*%s tinyq4: %s" % (G, X, TINYQ4_BIN))
    else:
        print("%s!%s tinyq4: missing at %s" % (R, X, TINYQ4_BIN))

    if os.path.exists(MODEL_PATH):
        print("%s*%s model: %s" % (G, X, MODEL_PATH))
    else:
        print("%s!%s model: missing at %s" % (R, X, MODEL_PATH))

    free = disk_free_gb(os.path.expanduser("~"))
    if free is not None:
        print("%s*%s home disk free: %.1f GB" % (G, X, free))

    print("%s*%s platform: %s" % (G, X, sys.platform))
    ram = total_ram_gb()
    if ram is not None:
        mark = G if ram >= MIN_RAM_GB else Y
        print("%s*%s system RAM: %.1f GB (CPU inference, no GPU needed)" % (mark, X, ram))

    problems = check_runtime()
    if problems:
        print("%s!%s this machine cannot run the bundled engine:" % (R, X))
        for problem in problems:
            print("    - %s" % problem)
    else:
        print("%s*%s runtime check: this machine can run the bundled engine" % (G, X))

    base_url, model = find_server(configured_port=DEFAULT_MODEL["port"])
    if base_url:
        print("%s*%s local server: %s (%s)" % (G, X, base_url, model))
    else:
        print("%s!%s local server: not detected on known ports" % (Y, X))
    return 0


def normalize_base_url(url):
    if not url:
        return None
    if "://" not in url:
        url = "http://" + url
    return url.rstrip("/")


def urlopen_json(url, timeout=0.5):
    with urllib.request.urlopen(url, timeout=timeout) as response:
        return json.load(response)


def probe_server(base_url, require_health):
    healthy = False
    try:
        response = urllib.request.urlopen(base_url + "/health", timeout=0.4)
        healthy = 200 <= response.status < 500
    except Exception:
        if require_health:
            return None

    try:
        data = urlopen_json(base_url + "/v1/models", timeout=0.5)
        model = data.get("data", [{}])[0].get("id") or "local model"
        return model
    except Exception:
        if healthy or not require_health:
            return "configured server"
    return None


def discover_ports(preferred_port=None):
    found = []
    if preferred_port:
        found.append(preferred_port)

    # AI Workshop spawns tinyq4 on a dynamically allocated free port. On Linux,
    # scan /proc for `tinyq4 ... --server <port>` and fall back to known ports.
    try:
        for pid in os.listdir("/proc"):
            if not pid.isdigit():
                continue
            try:
                with open("/proc/%s/cmdline" % pid, "rb") as handle:
                    args = [a.decode("utf-8", "replace") for a in handle.read().split(b"\x00") if a]
            except Exception:
                continue
            if not args or not any("tinyq4" in a for a in args):
                continue
            for index, item in enumerate(args):
                if item == "--server" and index + 1 < len(args):
                    try:
                        found.append(int(args[index + 1]))
                    except ValueError:
                        pass
    except Exception:
        pass

    seen = set()
    ordered = []
    for port in found + DEFAULT_PORTS:
        if port not in seen:
            seen.add(port)
            ordered.append(port)
    return ordered


def find_server(configured_url=None, configured_port=None):
    base_url = normalize_base_url(configured_url)
    if base_url:
        model = probe_server(base_url, require_health=False)
        return base_url, model or "configured server"

    for port in discover_ports(configured_port):
        candidate = "http://127.0.0.1:%d" % port
        model = probe_server(candidate, require_health=True)
        if model:
            return candidate, model
    return None, None


def _clear_line(handle=sys.stdout):
    handle.write("\r" + " " * 52 + "\r")
    handle.flush()


def stream(base_url, messages, render=True):
    """Stream a completion. Returns a dict with the text plus timing/usage.

    Ctrl-C during generation is caught here: it stops the stream and returns the
    partial text with ``interrupted`` set, so the caller can drop back to the
    prompt instead of crashing.
    """
    body = json.dumps({
        "messages": messages,
        "stream": True,
        "max_tokens": 2048,
        "temperature": 0.3,
    }).encode("utf-8")
    request = urllib.request.Request(
        base_url + "/v1/chat/completions",
        data=body,
        headers={"Content-Type": "application/json"},
    )
    output = ""
    in_think = False
    on_prefill = False
    interrupted = False
    started = time.time()
    ttft = None
    usage_tokens = None
    spin = 0
    try:
        with urllib.request.urlopen(request) as response:
            for raw in response:
                line = raw.decode("utf-8", "replace").strip()
                if not line.startswith("data:"):
                    continue
                payload = line[5:].strip()
                if payload == "[DONE]":
                    break
                try:
                    chunk = json.loads(payload)
                except Exception:
                    continue
                usage = chunk.get("usage")
                if isinstance(usage, dict) and usage.get("completion_tokens") is not None:
                    usage_tokens = usage["completion_tokens"]
                try:
                    delta = chunk["choices"][0]["delta"]
                except Exception:
                    continue

                # tinyq4 emits prefill_progress events before the first token.
                progress = delta.get("prefill_progress")
                if progress and ttft is None:
                    if render:
                        done = progress.get("done", 0)
                        total = progress.get("total", 0) or 0
                        bar = (" %d/%d" % (done, total)) if total else ""
                        sys.stderr.write("\r   %s%s reading prompt%s%s " % (D, SPIN[spin % 4], bar, X))
                        sys.stderr.flush()
                        spin += 1
                        on_prefill = True
                    continue

                think = delta.get("reasoning_content")
                content = delta.get("content")
                if think and render:
                    if on_prefill:
                        _clear_line(sys.stderr)
                        on_prefill = False
                    if not in_think:
                        sys.stdout.write(D)
                        in_think = True
                    sys.stdout.write(think)
                    sys.stdout.flush()
                if content:
                    if ttft is None:
                        ttft = time.time() - started
                    if render:
                        if on_prefill:
                            _clear_line(sys.stderr)
                            on_prefill = False
                        if in_think:
                            sys.stdout.write(X)
                            in_think = False
                        sys.stdout.write(content)
                        sys.stdout.flush()
                    output += content
    except KeyboardInterrupt:
        interrupted = True
    if on_prefill and render:
        _clear_line(sys.stderr)
    if in_think and render:
        sys.stdout.write(X)
    elapsed = time.time() - started
    tokens = usage_tokens if usage_tokens is not None else max(1, round(len(output) / 4))
    return {
        "text": output,
        "elapsed": elapsed,
        "ttft": ttft,
        "tokens": tokens,
        "approx": usage_tokens is None,
        "interrupted": interrupted,
    }


def safe_path(path):
    full = os.path.realpath(os.path.join(WORKSPACE, path))
    if full != WORKSPACE and not full.startswith(WORKSPACE + os.sep):
        raise ValueError("path escapes workspace: %s" % path)
    return full


def trunc(text):
    if len(text) <= MAX_OUT:
        return text
    return text[:MAX_OUT] + ("\n...[truncated %d chars]" % (len(text) - MAX_OUT))


def extract_tool(text):
    candidate = None
    match = FENCE.search(text)
    if match:
        candidate = match.group(1)
    else:
        start = text.find('{"name"')
        if start == -1:
            start = text.find('{ "name"')
        if start != -1:
            depth = 0
            for index in range(start, len(text)):
                if text[index] == "{":
                    depth += 1
                elif text[index] == "}":
                    depth -= 1
                    if depth == 0:
                        candidate = text[start:index + 1]
                        break
    if not candidate:
        return None
    try:
        obj = json.loads(candidate)
    except Exception:
        return None
    return obj if isinstance(obj, dict) and "name" in obj else None


def confirm(summary, yolo):
    if yolo:
        return True
    try:
        answer = input("   %sallow%s %s %s[y/N]%s " % (Y, X, summary, D, X)).strip().lower()
    except (EOFError, KeyboardInterrupt):
        print()
        return False
    return answer in ("y", "yes")


def run_tool(obj, yolo):
    name = obj.get("name", "")
    try:
        if name == "read":
            with open(safe_path(obj["path"]), "r", encoding="utf-8", errors="replace") as handle:
                return "read " + obj["path"], trunc(handle.read()), True, False

        if name == "ls":
            rel = obj.get("path", ".")
            path = safe_path(rel)
            items = sorted(os.listdir(path))
            output = "\n".join((item + "/") if os.path.isdir(os.path.join(path, item)) else item for item in items)
            return "ls " + rel, (output or "(empty)"), True, False

        if name == "write":
            rel = obj["path"]
            content = obj.get("content", "")
            if not confirm("write %s%s%s (%d bytes)" % (B, rel, X, len(content)), yolo):
                return "write " + rel, "denied by user", False, False
            path = safe_path(rel)
            os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
            with open(path, "w", encoding="utf-8") as handle:
                handle.write(content)
            return "write " + rel, "wrote %d bytes" % len(content), True, False

        if name == "edit":
            rel = obj["path"]
            old = obj.get("old", "")
            new = obj.get("new", "")
            if not confirm("edit %s%s%s" % (B, rel, X), yolo):
                return "edit " + rel, "denied by user", False, False
            path = safe_path(rel)
            with open(path, "r", encoding="utf-8", errors="replace") as handle:
                source = handle.read()
            count = source.count(old)
            if old == "" or count == 0:
                return "edit " + rel, "old text not found; adjust and retry", False, False
            if count > 1:
                return "edit " + rel, "old text appears %d times; add more context" % count, False, False
            with open(path, "w", encoding="utf-8") as handle:
                handle.write(source.replace(old, new, 1))
            return "edit " + rel, "replaced 1 occurrence", True, False

        if name == "bash":
            command = obj.get("command", "")
            if not confirm("run: %s%s%s" % (B, command, X), yolo):
                return "bash", "denied by user", False, False
            print("\n   %s*%s %sbash %s%s" % (A, X, B, command, X))
            kwargs = {
                "cwd": WORKSPACE,
                "stdout": subprocess.PIPE,
                "stderr": subprocess.STDOUT,
                "text": True,
                "bufsize": 1,
            }
            if IS_WINDOWS:
                # cmd.exe (what shell=True uses on Windows) chokes on the Unix-style
                # commands the model emits (e.g. `mkdir ~/Desktop/test1`). PowerShell
                # understands `~`, forward slashes, and most common Unix verbs.
                popen_target = ["powershell", "-NoProfile", "-NonInteractive", "-Command", command]
            else:
                kwargs["shell"] = True
                if hasattr(os, "setsid"):
                    kwargs["start_new_session"] = True
                popen_target = command
            process = subprocess.Popen(popen_target, **kwargs)
            killed = {"value": False}

            def kill_process():
                killed["value"] = True
                try:
                    if hasattr(os, "killpg"):
                        os.killpg(os.getpgid(process.pid), signal.SIGKILL)
                    else:
                        process.kill()
                except Exception:
                    try:
                        process.kill()
                    except Exception:
                        pass

            timer = threading.Timer(BASH_TIMEOUT, kill_process)
            timer.start()
            capture = []
            try:
                for line in process.stdout:
                    sys.stdout.write("   %s|%s %s" % (D, X, line))
                    sys.stdout.flush()
                    capture.append(line)
                    if len(capture) > 400:
                        capture.pop(0)
            finally:
                process.wait()
                timer.cancel()
            rc = process.returncode if process.returncode is not None else -1
            tail = "killed: exceeded %ds" % BASH_TIMEOUT if killed["value"] else "exit %d" % rc
            ok = rc == 0 and not killed["value"]
            print("   %s| [%s]%s" % ((G if ok else R), tail, X))
            result = ("".join(capture).strip() + "\n[%s]" % tail).strip()
            return "bash " + command, trunc(result), ok, True

        return name, "unknown tool", False, False
    except Exception as error:
        return name, "error: %s" % error, False, False


def show_result(text):
    lines = text.splitlines() or [""]
    for line in lines[:40]:
        print("   %s|%s %s" % (D, X, line))
    if len(lines) > 40:
        print("   %s| ...(%d more lines)%s" % (D, len(lines) - 40, X))


def stats_line(result):
    """A dim one-liner like '142 tok · 1.8s · 79 tok/s · first 0.3s'."""
    tps = result["tokens"] / result["elapsed"] if result["elapsed"] > 0 else 0.0
    approx = "~" if result["approx"] else ""
    first = (" %s·%s first %.1fs" % (D, X, result["ttft"])) if result["ttft"] else ""
    return "   %s%s%d tok %s·%s %.1fs %s·%s %.0f tok/s%s%s" % (
        D, approx, result["tokens"], D, X, result["elapsed"], D, X, tps, X, first,
    )


def transcript_markdown(messages):
    out = ["# Kairo session", "", "_%s_" % time.strftime("%Y-%m-%d %H:%M:%S"), ""]
    for message in messages:
        role = message.get("role")
        content = message.get("content", "")
        if role == "system":
            continue
        if role == "user" and content.startswith("[tool result"):
            out.append("```\n%s\n```" % content.strip())
            out.append("")
            continue
        out.append("**%s:**" % ("You" if role == "user" else "Kairo"))
        out.append("")
        out.append(content)
        out.append("")
    return "\n".join(out)


def header(base_url, model, yolo):
    os.system("cls" if IS_WINDOWS else "clear")
    print(A + BANNER + X)
    tag = (" %sv%s%s" % (D, VERSION, X)) if VERSION else ""
    print("   %slocal coding agent · CPU · lite%s%s" % (D, X, tag))
    print()
    if base_url:
        print("   %s*%s %s %s- %s%s" % (G, X, model, D, base_url, X))
    else:
        print("   %s*%s no local model server %s- run `kairolite model`, then restart, or pass --url%s" % (R, X, D, X))
    mode = (Y + "yolo" + X) if yolo else (D + "confirm" + X)
    print("   %stools%s read ls write edit bash %s·%s %s" % (D, X, D, X, mode))
    print("   %s/yolo /tools /model /gpu /save /cwd /stats /clear /help /exit%s" % (D, X))
    print("   %scwd %s%s" % (D, WORKSPACE, X))
    if gpu_is_capable():
        print("   %s* NVIDIA GPU found - run /gpu for the ~20x faster GPU build%s" % (G, X))
    print("   %sAS IS · no warranty · autonomous agent that runs commands · see DISCLAIMER.md%s\n" % (D, X))


def main(argv=None):
    global WORKSPACE

    if argv is None:
        argv = sys.argv[1:]

    if argv:
        command = argv[0]
        if command == "setup":
            return cmd_setup(argv[1:])
        if command == "serve":
            return cmd_serve(argv[1:])
        if command == "doctor":
            return cmd_doctor(argv[1:])
        if command == "postinstall":
            return cmd_postinstall(argv[1:])
        if command in ("model", "models"):
            return cmd_model(argv[1:])
        if command == "stop":
            return cmd_stop(argv[1:])
        if command == "gpu":
            return cmd_gpu(argv[1:])

    args = parse_chat_args(argv)
    WORKSPACE = os.path.realpath(args.workspace)
    if not os.path.isdir(WORKSPACE):
        print("%serror: workspace does not exist: %s%s" % (R, WORKSPACE, X), file=sys.stderr)
        return 2

    try:
        base_url, model, server_proc = ensure_local_server(args)
    except KeyboardInterrupt:
        print()
        return 0

    is_oneshot = args.oneshot is not None
    # On exit, free memory by stopping the local server we were using - whether we
    # spawned it or attached to one already running. We never touch a server the
    # user explicitly pointed us at with --url, and for a quick one-shot we only
    # stop a server we started ourselves (so back-to-back -p calls can reuse it).
    sweep_port = None
    if base_url and not args.url and _is_local_url(base_url) and not is_oneshot:
        sweep_port = _port_of_url(base_url)

    def teardown():
        if server_proc is not None:
            terminate_server(server_proc)
        if sweep_port is not None:
            kill_tinyq4(sweep_port)

    atexit.register(teardown)

    yolo = args.yolo
    max_steps = args.max_steps if args.max_steps and args.max_steps > 0 else 25
    os_name = "Windows" if IS_WINDOWS else ("macOS" if sys.platform == "darwin" else "Linux")
    shell_name = "PowerShell" if IS_WINDOWS else "a POSIX shell (bash/sh)"
    sys_content = (SYSTEM
                   .replace("__WS__", WORKSPACE)
                   .replace("__OS__", os_name)
                   .replace("__SHELL__", shell_name))
    sys_msg = {"role": "system", "content": sys_content}

    try:
        if is_oneshot:
            return run_oneshot(args, base_url, sys_msg, yolo, max_steps)
        header(base_url, model, yolo)
        chat_loop(args, base_url, model, yolo, max_steps, sys_msg, [sys_msg])
        return 0
    finally:
        teardown()


def run_turn(args, base_url, messages, yolo, max_steps, stats_on):
    """Stream the model and run its tool calls until it stops.

    Returns the base_url to keep using (it can change if we reconnect to a
    different server), or None if the connection was lost.
    """
    for _ in range(max_steps):
        print("\n   %skairo%s  " % (A, X), end="", flush=True)
        try:
            result = stream(base_url, messages)
        except Exception as first_error:
            next_url, _model = find_server(args.url, args.port)
            if next_url:
                base_url = next_url
                try:
                    result = stream(base_url, messages)
                except Exception as second_error:
                    print("%serror: %s%s" % (R, second_error, X))
                    return base_url
            else:
                print("%sconnection lost: %s%s" % (R, first_error, X))
                return None

        text = result["text"]
        if result["interrupted"]:
            if text:
                messages.append({"role": "assistant", "content": text})
            print("\n   %s(interrupted)%s" % (Y, X))
            return base_url
        if stats_on and text:
            print("\n" + stats_line(result))

        messages.append({"role": "assistant", "content": text})
        tool = extract_tool(text)
        if not tool:
            break

        label, output, ok, shown = run_tool(tool, yolo)
        if not shown:
            print("\n   %s*%s %s%s%s" % ((G if ok else R), X, B, label, X))
            show_result(output)
        messages.append({"role": "user", "content": "[tool result - %s]\n%s" % (label, output)})
    else:
        print("\n   %sreached step limit (%d). Type to continue.%s" % (Y, max_steps, X))
    print()
    return base_url


def run_oneshot(args, base_url, sys_msg, yolo, max_steps):
    prompt = args.oneshot
    if prompt == "-":
        prompt = sys.stdin.read().strip()
    if not prompt:
        print("%serror: empty prompt%s" % (R, X), file=sys.stderr)
        return 2
    if not base_url:
        print("%serror: no model server available%s" % (R, X), file=sys.stderr)
        return 1
    messages = [sys_msg, {"role": "user", "content": prompt}]
    run_turn(args, base_url, messages, yolo, max_steps, stats_on=False)
    return 0


def chat_loop(args, base_url, model, yolo, max_steps, sys_msg, messages):
    stats_on = True
    while True:
        try:
            user = input("   %s>%s " % (A, X)).strip()
        except (EOFError, KeyboardInterrupt):
            print()
            break

        if not user:
            continue
        if user in ("/exit", "/quit", "/q"):
            break
        if user in ("/clear", "/reset"):
            messages = [sys_msg]
            header(base_url, model, yolo)
            continue
        if user == "/yolo":
            yolo = not yolo
            message = "ON - tools run without asking" if yolo else "OFF - confirm writes and bash"
            print("   %syolo %s%s\n" % (Y, message, X))
            continue
        if user == "/tools":
            print("   %sread · ls · write · edit · bash%s\n" % (D, X))
            continue
        if user == "/model":
            cmd_model([])
            print()
            continue
        if user == "/gpu":
            cmd_gpu([])
            print()
            continue
        if user == "/cwd":
            print("   %s%s%s\n" % (D, WORKSPACE, X))
            continue
        if user == "/stats":
            stats_on = not stats_on
            print("   %sstats %s%s\n" % (D, "on" if stats_on else "off", X))
            continue
        if user.split(None, 1)[0] == "/save":
            parts = user.split(None, 1)
            name = parts[1].strip() if len(parts) > 1 else time.strftime("kairo-%Y%m%d-%H%M%S.md")
            try:
                with open(safe_path(name), "w", encoding="utf-8") as handle:
                    handle.write(transcript_markdown(messages))
                print("   %ssaved transcript to %s%s\n" % (G, name, X))
            except Exception as error:
                print("   %scould not save: %s%s\n" % (R, error, X))
            continue
        if user == "/help":
            print("   %scommands%s /yolo /tools /model /gpu /save [file] /cwd /stats /clear /help /exit" % (D, X))
            print("   %s/gpu%s scans for an NVIDIA GPU and links the faster GPU build" % (D, X))
            print("   %sanything else is sent to the agent - it can read/write files and run bash here%s\n" % (D, X))
            continue

        if not base_url:
            base_url, model = find_server(args.url, args.port)
            if not base_url:
                print("   %sno server - restart `kairolite`, or run `/model` to (re)download%s\n" % (R, X))
                continue
            header(base_url, model, yolo)

        messages.append({"role": "user", "content": user})
        try:
            base_url = run_turn(args, base_url, messages, yolo, max_steps, stats_on)
        except KeyboardInterrupt:
            print("\n   %s(interrupted)%s\n" % (Y, X))

    return 0


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