"""Manual, single-model maintenance launcher. Never starts/stops Contextra."""
import argparse
import fcntl
import json
import os
from pathlib import Path
import signal
import subprocess
import sys
import time

ROOT = Path(__file__).resolve().parent
STATE = Path.home() / ".local/state/contextra-llm"


def read_status():
    containers = subprocess.run(["docker", "ps", "--format", "{{.Names}}"],
        check=True, capture_output=True, text=True, timeout=10).stdout.splitlines()
    active = [name for name in containers if name.startswith("contextra-icerti-")]
    gpu = subprocess.run(["nvidia-smi", "--query-gpu=memory.used", "--format=csv,noheader,nounits"],
        check=True, capture_output=True, text=True, timeout=10)
    return dict(contextra_running=active, gpu_used_mib=[int(x) for x in gpu.stdout.splitlines()])


def pid_identity(pid):
    try:
        stat = (Path("/proc") / str(pid) / "stat").read_text()
        started = stat.rsplit(")", 1)[1].split()[19]
        executable = str((Path("/proc") / str(pid) / "exe").resolve(strict=True))
        return dict(pid=pid, started=started, executable=executable)
    except (OSError, IndexError):
        return None


def stop_owned():
    receipt = STATE / "process.json"
    if not receipt.is_file():
        print("등록된 LLM 실행이 없습니다.")
        return
    record = json.loads(receipt.read_text())
    actual = pid_identity(record["pid"])
    if actual is None:
        print("등록된 LLM은 이미 종료됐습니다.")
        return
    if actual != record or not Path(actual["executable"]).is_relative_to(ROOT):
        raise RuntimeError("PID ownership mismatch; no signal sent")
    os.kill(record["pid"], signal.SIGTERM)
    deadline = time.monotonic() + 30
    while time.monotonic() < deadline:
        if pid_identity(record["pid"]) != record:
            print("LLM 정상 종료를 확인했습니다.")
            return
        time.sleep(.2)
    raise RuntimeError("LLM did not exit; no forced kill or Contextra restart performed")


def main():
    parser = argparse.ArgumentParser(description="오프라인 LLM: chat/serve/cpu/status/stop")
    parser.add_argument("mode", choices=("chat", "serve", "cpu", "status", "stop"))
    parser.add_argument("--context", type=int, choices=(8192, 16384))
    args = parser.parse_args()
    if os.geteuid() == 0:
        raise RuntimeError("si 계정에서 실행하세요. sudo는 필요 없습니다.")
    os.umask(0o077)
    STATE.mkdir(parents=True, exist_ok=True)
    if args.mode == "stop":
        stop_owned()
        return
    status = read_status()
    if args.mode == "status":
        print(json.dumps(status, ensure_ascii=False, indent=2))
        return
    config = json.loads((ROOT / "config.json").read_text())
    if status["contextra_running"]:
        raise RuntimeError("Contextra 컨테이너가 실행 중입니다. 유지보수 정지를 먼저 확인하세요.")
    if args.mode != "cpu" and any(n > config["gpu_idle_max_mib"] for n in status["gpu_used_mib"]):
        raise RuntimeError("다른 GPU 작업 또는 잔여 메모리가 있습니다. 임의 종료하지 않습니다.")
    with (STATE / "run.lock").open("a") as lock:
        fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
        variant = "cpu" if args.mode == "cpu" else "gpu"
        binary_dir = ROOT / variant / "llama-b11146"
        command = [str(binary_dir / ("llama-server" if args.mode == "serve" else "llama-cli")),
            "-m", config["model"], "--offline", "-c", str(args.context or config["context"]),
            "-t", str(config["threads"]), "-ngl", "0" if variant == "cpu" else "all",
            "--reasoning", "off",
            "--temp", "0.7", "--top-p", "0.8", "--top-k", "20", "--min-p", "0",
            "--presence-penalty", "1.5", "--no-mmproj"]
        if args.mode == "serve":
            command += ["--host", "127.0.0.1", "--port", str(config["port"]),
                "--parallel", "1", "--alias", config["alias"],
                "--api-key-file", str(ROOT / "api-key"), "--timeout", "180", "-n", "4096"]
            print("LLM 수동 실행: localhost 전용, 동시 추론1개. 종료는 Ctrl+C 또는 ct-llm stop.", flush=True)
        else:
            command += ["--system-prompt-file", str(ROOT / "system.ko.txt"), "-n", "4096"]
        env = dict(os.environ, LANG="C.UTF-8", LC_ALL="C.UTF-8")
        libraries = [str(binary_dir)]
        if variant == "gpu":
            libraries.append(str(ROOT / "cuda-runtime/cudart-llama-b11146-bin-ubuntu-cuda-12.8-x64"))
        env["LD_LIBRARY_PATH"] = ":".join(libraries)
        for key in list(env):
            if key.startswith("LLAMA_ARG_") or key in {"HF_TOKEN", "HUGGING_FACE_HUB_TOKEN"}:
                env.pop(key)
        child = subprocess.Popen(command, env=env)
        receipt = STATE / "process.json"
        for _ in range(50):
            identity = pid_identity(child.pid)
            if identity and Path(identity["executable"]).is_relative_to(ROOT):
                receipt.write_text(json.dumps(identity))
                break
            if child.poll() is not None:
                raise RuntimeError("LLM exited before ownership registration")
            time.sleep(.02)
        else:
            child.terminate()
            child.wait(timeout=15)
            raise RuntimeError("Could not register LLM ownership")
        try:
            code = child.wait()
        except KeyboardInterrupt:
            child.terminate()
            code = child.wait(timeout=30)
        finally:
            if child.poll() is not None:
                receipt.unlink(missing_ok=True)
        raise SystemExit(code)


if __name__ == "__main__":
    main()