#!/usr/bin/env python3
"""kvcs-disk-ops - EFC disk operations tool.

Provides disk status inspection, hot offline/online, device formatting,
and PanguCsDiskStatus marker management for bad-disk recovery. Does not
manage EFC lifecycle for the recover flow; offline/online require EFC
running.
"""

import fcntl
import json
import os
import re
import shutil
import stat
import subprocess
import sys
import time
import codecs


if hasattr(sys.stdout, "buffer"):
    sys.stdout = codecs.getwriter("utf-8")(sys.stdout.buffer)
if hasattr(sys.stderr, "buffer"):
    sys.stderr = codecs.getwriter("utf-8")(sys.stderr.buffer)

_print = print


def print(*args, **kwargs):
    kwargs.setdefault("flush", True)
    _print(*args, **kwargs)


SCRIPT_DIR = os.path.abspath(os.path.dirname(__file__))

# Build metadata written at packaging time: the deploy tarball ships COMMIT_ID next to this
# script, RPM/DEB installs it under /usr/share/kvcs.
_INSTALL_ROOT = os.environ.get("KVCS_ROOT", "") or SCRIPT_DIR
BUILD_INFO_PATHS = tuple(
    p
    for p in (
        os.environ.get("KVCS_BUILD_INFO", ""),  # lets release checks bypass installed packages
        os.path.join(_INSTALL_ROOT, "BUILD_INFO"),
        os.path.join(_INSTALL_ROOT, "COMMIT_ID"),
    )
    if p
)
# Read last: the host may carry an older RPM while this copy runs from a tarball or the repo,
# and then the metadata shipped with this copy is the truthful one.
SYSTEM_BUILD_INFO_PATHS = (
    os.environ.get("KVCS_SYSTEM_BUILD_INFO", "") or "/usr/share/kvcs/efc-build-info",
)


def read_build_info(paths: "tuple[str, ...]" = BUILD_INFO_PATHS) -> "dict[str, str]":
    for path in paths:
        try:
            with open(path, "r") as fh:
                text = fh.read()
        except OSError:
            continue
        info = {}
        for line in text.splitlines():
            key, sep, value = line.partition("=")
            if sep:
                info[key.strip()] = value.strip()
        if info.get("version"):
            return info
    return {}


def repo_build_info() -> "dict[str, str]":
    # Unpackaged (run straight from the repo): same version source as the compiled binaries,
    # build identity from git, and the file's own mtime as the closest thing to a build time.
    version_file = os.path.abspath(os.path.join(SCRIPT_DIR, "..", "..", "VERSION"))
    try:
        with open(version_file, "r") as fh:
            base = fh.read().strip()
    except OSError:
        return {}
    info = {"version": base}
    try:
        sha = subprocess.check_output(
            ["git", "-C", SCRIPT_DIR, "rev-parse", "HEAD"], stderr=subprocess.DEVNULL
        ).decode("utf-8", "replace").strip()
    except (subprocess.CalledProcessError, OSError):
        sha = ""
    if sha:
        info["commit"] = sha
        info["revision"] = sha[:8]
        try:
            dirty = subprocess.call(
                ["git", "-C", SCRIPT_DIR, "diff", "--quiet", "HEAD", "--"],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )
        except OSError:
            dirty = 0
        if dirty:
            info["revision"] += "-dirty"
    try:
        mtime = os.stat(os.path.abspath(__file__)).st_mtime
        info["build_time"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(mtime))
    except OSError:
        pass
    return info


def version_lines() -> "list[str]":
    info = (
        read_build_info()
        or repo_build_info()
        or read_build_info(SYSTEM_BUILD_INFO_PATHS)
    )
    # Drop the -dev.<sha8> artifact suffix and prefer `revision`, so this prints exactly what the
    # compiled binaries print for the same build.
    version = (info.get("version", "") or "unknown").split("-dev.")[0]
    revision = info.get("revision", "") or info.get("commit", "")[:8]
    arch = info.get("arch", "") or os.uname().machine
    return [
        "v" + version,
        "  commit:    " + (revision or "unknown"),
        "  built:     " + (info.get("build_time", "") or "unknown"),
        "  platform:  linux/" + arch,
    ]


def print_version() -> None:
    print("\n".join(version_lines()))


def find_efc_config():
    env = os.environ.get("EFC_CONFIG", "")
    if env and os.path.isfile(env):
        return env
    system = "/etc/kvcs/efc.yaml"
    if os.path.isfile(system):
        return system
    sibling = os.path.join(SCRIPT_DIR, "efc.yaml")
    if os.path.isfile(sibling):
        return sibling
    return ""


def parse_yaml_value(text, key):
    for line in text.splitlines():
        stripped = line.split("#")[0].strip()
        if stripped.startswith(key + ":"):
            return stripped[len(key) + 1:].strip()
    return ""


def load_data_dir():
    cfg = find_efc_config()
    if not cfg:
        sys.exit("ERROR: efc.yaml not found (set EFC_CONFIG or install to /etc/kvcs/efc.yaml)")
    try:
        with open(cfg, "r") as fh:
            text = fh.read()
    except Exception as exc:
        sys.exit("ERROR: failed to read efc.yaml: {} ({})".format(cfg, exc))
    data_dir = parse_yaml_value(text, "data_dir")
    if not data_dir:
        sys.exit("ERROR: efc.yaml missing data_dir: {}".format(cfg))
    return data_dir.rstrip("/")


# `-V` runs before load_data_dir(): a host without efc.yaml must still report a version.
if len(sys.argv) > 1 and sys.argv[1] in ("-V", "--version"):
    print_version()
    sys.exit(0)

DATA_DIR = load_data_dir()
DISK_CONF = DATA_DIR + "/chunk_file_system_conf.json"
DEFAULT_SOCKET = "/var/run/kvcs/efc-grpc.sock"


def log(message):
    print("[{}] {}".format(time.strftime("%H:%M:%S"), message))


def warn(message):
    log("WARN: {}".format(message))


def die(message):
    log("ERROR: {}".format(message))
    sys.exit(1)


def run(cmd, stdout=None, stderr=None):
    try:
        return subprocess.run(cmd, stdout=stdout, stderr=stderr)
    except FileNotFoundError:
        return subprocess.CompletedProcess(cmd, 127)


def check_output(cmd):
    try:
        return subprocess.check_output(cmd, stderr=subprocess.DEVNULL).decode("utf-8", "replace")
    except FileNotFoundError as exc:
        raise subprocess.CalledProcessError(127, cmd) from exc


def find_tool(name):
    found = shutil.which(name)
    if found:
        return found

    search_dirs = [
        SCRIPT_DIR,
        os.path.abspath(os.path.join(SCRIPT_DIR, "..")),
        os.path.abspath(os.path.join(SCRIPT_DIR, "..", "tools")),
        os.environ.get("KVCS_ROOT", ""),
    ]
    for directory in search_dirs:
        if not directory:
            continue
        candidate = os.path.join(directory, name)
        if os.path.isfile(candidate) and os.access(candidate, os.X_OK):
            return candidate
    return None


def parse_disk_conf():
    if not os.path.isfile(DISK_CONF):
        die("disk config not found: {}".format(DISK_CONF))

    try:
        with open(DISK_CONF, "r") as fh:
            conf = json.load(fh)
    except Exception as exc:
        die("disk config parse failed: {} ({})".format(DISK_CONF, exc))

    spdk = conf.get("SpdkFileSystem")
    if not isinstance(spdk, dict):
        die("disk config missing SpdkFileSystem: {}".format(DISK_CONF))

    disks = []
    for name, dev in spdk.items():
        if isinstance(name, str) and name.startswith("disk") and isinstance(dev, str):
            disks.append((name, dev))
    return disks


def get_device(disk_name):
    for name, dev in parse_disk_conf():
        if name == disk_name:
            return dev
    return ""


def marker_path(disk_name):
    return DATA_DIR + "/PanguCsDiskStatus" + disk_name


def is_bad_disk(disk_name):
    return os.path.isfile(marker_path(disk_name))


def list_bad_disks():
    bad = []
    for name, _dev in parse_disk_conf():
        if is_bad_disk(name):
            bad.append(name)
    return bad


def is_lvm(dev):
    return run(["lvs", dev], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL).returncode == 0


def pgrep_efc():
    try:
        output = check_output(["pgrep", "-x", "kvcs_efc_server"])
    except subprocess.CalledProcessError:
        return ""
    return output.strip()


def efc_is_running():
    return bool(pgrep_efc())


def read_nul_file(path):
    try:
        with open(path, "rb") as fh:
            data = fh.read()
    except OSError:
        return []
    return [part.decode("utf-8", "replace") for part in data.split(b"\0") if part]


def get_efc_socket():
    pid = pgrep_efc()
    if not pid:
        return DEFAULT_SOCKET
    pid = pid.splitlines()[0]

    for arg in read_nul_file("/proc/{}/cmdline".format(pid)):
        if arg.startswith("--grpc-socket="):
            return arg[len("--grpc-socket="):]

    for env in read_nul_file("/proc/{}/environ".format(pid)):
        if env.startswith("KVCS_GRPC_SOCKET="):
            return env[len("KVCS_GRPC_SOCKET="):]

    return DEFAULT_SOCKET


def is_block_device(path):
    try:
        return stat.S_ISBLK(os.stat(path).st_mode)
    except OSError:
        return False


def format_disk(dev, idempotent=False):
    """Full format: blkdiscard + spdk_fs_tool --format + verify.

    Matches kvcs-ctl reset-efc flow: blkdiscard (full-disk TRIM) then
    spdk_fs_tool --format --force then --check verify.

    If idempotent=True, check first and skip if already formatted.
    """
    spdk_tool = find_tool("spdk_fs_tool")
    if not spdk_tool:
        die("spdk_fs_tool not found")

    if idempotent:
        rc = run([spdk_tool, "--check", "--addr={}".format(dev),
                  "--access_mode=2", "--partition=0", "--enable_lvm_start"],
                 stdout=subprocess.PIPE, stderr=subprocess.PIPE).returncode
        if rc == 0:
            log("  {} already formatted, skip".format(dev))
            return

    log("  blkdiscard {}".format(dev))
    rc, _, err = run_capture(["blkdiscard", dev])
    if rc != 0:
        die("blkdiscard {} failed: {}".format(dev, err.strip()))

    log("  spdk_fs_tool --format {}".format(dev))
    rc, _, err = run_capture([
        spdk_tool,
        "--format",
        "--addr={}".format(dev),
        "--access_mode=2",
        "--alloc_block_size=256",
        "--partition=0",
        "--force",
        "--enable_lvm_start",
    ])
    if rc != 0:
        die("spdk_fs_tool --format {} failed: {}".format(dev, err.strip()))

    # Verify format
    rc, _, err = run_capture([spdk_tool, "--check", "--addr={}".format(dev),
              "--access_mode=2", "--partition=0", "--enable_lvm_start"])
    if rc != 0:
        die("spdk_fs_tool --check {} failed after format: {}".format(dev, err.strip()))

    log("  {} format complete".format(dev))


# ----------------------------------------------------------------
#  EFC / cslib internal helpers (user-facing output keeps cslib
#  concepts hidden; "online/offline" is the only status vocabulary).
# ----------------------------------------------------------------

def resolve_yaml_scalar(raw):
    """Resolve ${VAR:-default} / ${VAR} and strip YAML quotes."""
    raw = raw.split("#")[0].strip()
    raw = raw.strip('"').strip("'")
    m = re.match(r'^\$\{([A-Za-z_][A-Za-z0-9_]*)(:-(.*))?\}$', raw)
    if m:
        var = m.group(1)
        default = m.group(3) if m.group(2) is not None else ""
        raw = os.environ.get(var, default)
    return str(raw).strip('"').strip("'")


def parse_yaml_section(text, section):
    """Return a dict of key:value under a top-level section."""
    lines = text.splitlines()
    result = {}
    in_section = False
    section_indent = None
    for line in lines:
        stripped = line.lstrip()
        cur_indent = len(line) - len(stripped)
        if stripped.startswith(section + ":"):
            in_section = True
            section_indent = cur_indent
            continue
        if not in_section:
            continue
        if cur_indent <= section_indent and stripped and not stripped.startswith("#"):
            break
        if ":" in stripped:
            key, _, val = stripped.partition(":")
            result[key.strip()] = val.strip().split("#")[0].strip()
    return result


def get_efc_backend(text):
    return resolve_yaml_scalar(parse_yaml_value(text, "backend") or "")


def get_node_id_ip(text):
    node_id = resolve_yaml_scalar(parse_yaml_value(text, "node_id") or "")
    if ":" in node_id:
        return node_id.rsplit(":", 1)[0]
    return node_id or "127.0.0.1"


def get_cslib_port(text):
    extra = parse_yaml_section(text, "extra_pangu_flags")
    port_raw = extra.get("pangu_cslib_port", "10280")
    return str(int(resolve_yaml_scalar(port_raw)))


def cslib_addr(text):
    override = os.environ.get("KVCS_CSLIB_ADDR", "").strip()
    if override:
        return override
    return "tcp://{}:{}".format(get_node_id_ip(text), get_cslib_port(text))


def cslib_tool_path():
    tool = find_tool("pangu_cslib_tool")
    if not tool:
        die("disk management tool missing (pangu_cslib_tool), check deployment integrity")
    return tool


def run_capture(cmd, timeout=None):
    """Run cmd and return (rc, stdout, stderr)."""
    try:
        p = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, timeout=timeout)
    except subprocess.TimeoutExpired:
        return 1, "", "timeout"
    except FileNotFoundError:
        return 127, "", "command not found"
    return p.returncode, p.stdout.decode("utf-8", "replace"), p.stderr.decode("utf-8", "replace")


def cslib_get_status_raw(diskid, tool, addr):
    """Run getstatus; return (rc, stdout, stderr)."""
    cmd = [tool, "--mode=disk", "--addr={}".format(addr),
           "--op=getstatus", "--diskid={}".format(diskid)]
    return run_capture(cmd, timeout=5)


#  Disk status enum. pangu_cslib_tool diskStatus values:
# DISK_OK / DISK_ERROR / DISK_READONLY / DISK_SHUTDOWN。
ONLINE_STATUSES = ("DISK_OK",)
OFFLINE_STATUSES = ("DISK_ERROR",)


def is_online_status(status):
    return status in ONLINE_STATUSES


def is_offline_status(status):
    return status in OFFLINE_STATUSES


def parse_cslib_status(output):
    """Conservatively parse getstatus output.

    Returns 'DISK_OK', 'DISK_ERROR', or None.
    cslib diskStatus enum is DISK_OK, DISK_ERROR, DISK_READONLY and
    DISK_SHUTDOWN. Only DISK_OK and DISK_ERROR are treated as
    first-class states here; DISK_READONLY/DISK_SHUTDOWN are returned
    as None (other) so callers that need strict offline semantics
    (e.g. wait_offline) do not confuse them.
    """
    out = output.upper()
    if re.search(r'\bDISK_READONLY\b', out):
        return None
    if re.search(r'\bDISK_SHUTDOWN\b', out):
        return None
    if re.search(r'\bDISK_ERROR\b', out):
        return "DISK_ERROR"
    if re.search(r'\bDISK_OK\b', out):
        return "DISK_OK"
    return None


def parse_disk_conf_full():
    if not os.path.isfile(DISK_CONF):
        die("disk config not found: {}".format(DISK_CONF))
    try:
        with open(DISK_CONF, "r") as fh:
            data = json.load(fh)
    except Exception as exc:
        die("disk config parse failed: {} ({})".format(DISK_CONF, exc))

    if not isinstance(data, dict):
        die("disk config format error: root must be an object")
    spdk = data.get("SpdkFileSystem")
    if not isinstance(spdk, dict):
        die("disk config missing SpdkFileSystem: {}".format(DISK_CONF))

    slots = []
    for name, dev in spdk.items():
        if not isinstance(name, str) or not name.startswith("disk"):
            die("disk config invalid slot name: {}".format(name))
        try:
            idx = int(name[4:])
        except ValueError:
            die("disk config invalid slot name: {}".format(name))
        if not isinstance(dev, str):
            die("disk config invalid device path: {}".format(name))
        slots.append((idx, name, dev))

    slots.sort(key=lambda x: x[0])
    basic = data.get("BasicFileSystem", {})
    return slots, basic


def acquire_lock():
    lock_path = os.path.join(DATA_DIR, ".disk-ops.lock")
    try:
        fd = os.open(lock_path, os.O_RDWR | os.O_CREAT, 0o644)
        fcntl.flock(fd, fcntl.LOCK_EX)
        return fd
    except OSError as exc:
        die("cannot acquire disk operation lock: {} ({})".format(lock_path, exc))


def release_lock(fd):
    try:
        fcntl.flock(fd, fcntl.LOCK_UN)
        os.close(fd)
    except OSError:
        pass


def is_mounted(dev):
    try:
        with open("/proc/mounts", "r") as fh:
            for line in fh:
                fields = line.split()
                if fields and fields[0] == dev:
                    return True
    except OSError:
        pass
    return False


def cslib_get_status(diskid, tool, addr):
    # Non-zero return code or ErrorCode != PANGU_OK means query failed.
    rc, out, _ = cslib_get_status_raw(diskid, tool, addr)
    if rc != 0:
        return None
    return parse_cslib_status(out)


def cslib_get_status_enum(diskid, tool, addr):
    """Return the full diskStatus enum reported by getstatus, or None.

    Unlike cslib_get_status this does not collapse DISK_READONLY /
    DISK_SHUTDOWN to None, so callers (status display) can show the raw
    enum. None means the query failed or the output was unparseable.
    """
    rc, out, _ = cslib_get_status_raw(diskid, tool, addr)
    if rc != 0:
        return None
    m = re.search(r'\bDISK_(?:OK|ERROR|READONLY|SHUTDOWN)\b', out.upper())
    return m.group(0) if m else None


def find_cslib_tool_quiet():
    """Like cslib_tool_path() but returns None instead of dying (for status)."""
    return find_tool("pangu_cslib_tool")


def read_efc_yaml():
    cfg = find_efc_config()
    if not cfg:
        return ""
    try:
        with open(cfg, "r") as fh:
            return fh.read()
    except Exception:
        return ""


def resolve_cslib_addr_safe(text):
    """Resolve cslib address; returns "" on failure instead of dying."""
    if not text:
        return ""
    try:
        return cslib_addr(text)
    except Exception:
        return ""


def cslib_set_status(diskid, status, tool, addr):
    return run([
        tool, "--mode=disk", "--addr={}".format(addr),
        "--op=setstatus", "--diskid={}".format(diskid), "--status={}".format(status)
    ], stdout=subprocess.PIPE, stderr=subprocess.PIPE).returncode


def cslib_notify_online(diskid, tool, addr):
    return run([
        tool, "--mode=disk", "--addr={}".format(addr),
        "--op=notifyonline", "--diskid={}".format(diskid)
    ], stdout=subprocess.PIPE, stderr=subprocess.PIPE).returncode


def cslib_query_online(diskid, tool, addr):
    # Online success is determined solely by queryonline output:
    # 1) return code 0 (ErrorCode:PANGU_OK)
    # 2) onlineStatus: ONLINED
    # 3) diskStatus: DISK_OK
    rc, out, _ = run_capture([
        tool, "--mode=disk", "--addr={}".format(addr),
        "--op=queryonline", "--diskid={}".format(diskid)
    ])
    if rc != 0:
        return False
    # Exact match for "ONLINED"; avoid matching "NEED_ONLINE".
    if not re.search(r'\bonlineStatus\s*:\s*ONLINED\b', out, re.IGNORECASE):
        return False
    # \bdiskStatus with word boundary, won't mismatch "onlineStatus" or "onlinediskStatus".
    if not re.search(r'\bdiskStatus\s*:\s*DISK_OK\b', out, re.IGNORECASE):
        return False
    return True


def wait_online(diskid, tool, addr, timeout=60, interval=2):
    # New SOP: online success = queryonline return code 0 (ErrorCode:PANGU_OK),
    # OnlineStatus:ONLINED, DiskStatus:DISK_OK, all from queryonline output.
    deadline = time.time() + timeout
    while time.time() < deadline:
        if cslib_query_online(diskid, tool, addr):
            return True
        time.sleep(interval)
    return False


def wait_offline(diskid, tool, addr, timeout=10, interval=1):
    # New SOP: offline success = getstatus return code 0 and DiskStatus:DISK_ERROR.
    deadline = time.time() + timeout
    while time.time() < deadline:
        if is_offline_status(cslib_get_status(diskid, tool, addr)):
            return True
        time.sleep(interval)
    return False


def get_disk_identity(dev):
    """Return stable identity string (serial or WWN) for a block device."""
    # Try /sys/block/<base>/device/serial, then wwid, then by-id symlink
    base = os.path.basename(dev)
    # Handle /dev/mapper or LVM
    serial_paths = [
        "/sys/block/{}/device/serial".format(base),
        "/sys/block/{}/device/wwid".format(base),
    ]
    for p in serial_paths:
        try:
            with open(p, "r") as fh:
                val = fh.read().strip()
                if val:
                    return val
        except (OSError, IOError):
            pass
    # Fallback: look for by-id symlink
    by_id_dir = "/dev/disk/by-id"
    try:
        for link in os.listdir(by_id_dir):
            full = os.path.join(by_id_dir, link)
            if os.path.realpath(full) == os.path.realpath(dev):
                return link
    except OSError:
        pass
    return "-"



def show_disk_status():
    """Show disk status from cslib runtime state or marker-based fallback.

    When cslib (pangu_cslib_tool getstatus) is reachable:
      DISK_OK -> HEALTHY, DISK_ERROR -> BAD, other enums shown as-is.
    When cslib is unreachable (EFC not running / tool missing / timeout):
      bad-disk marker present -> BAD, no marker -> HEALTHY.

    Exit code: BAD (proactively offlined) does NOT trigger exit 1;
    only MISSING / INACCESSIBLE / DISK_READONLY / DISK_SHUTDOWN do.
    """
    slots, _ = parse_disk_conf_full()
    total = len(slots)
    offline_count = 0
    missing_count = 0
    unknown_count = 0
    degraded_count = 0
    cslib_reachable = True
    print("{:<10} {:<25} {:<22} {}".format("Slot", "Device", "Status", "Identity"))
    print("{:<10} {:<25} {:<22} {}".format("----", "------", "------", "--------"))

    tool = find_cslib_tool_quiet()
    if tool is None:
        cslib_reachable = False
    addr = resolve_cslib_addr_safe(read_efc_yaml())
    if not addr:
        cslib_reachable = False

    for diskid, name, dev in slots:
        identity = "-"

        if not os.path.exists(dev):
            status_text = "MISSING"
            missing_count += 1
        elif not is_block_device(dev):
            status_text = "INACCESSIBLE"
            missing_count += 1
        elif not os.access(dev, os.R_OK | os.W_OK):
            status_text = "INACCESSIBLE"
            missing_count += 1
        elif not cslib_reachable:
            # Cslib unreachable; fall back to marker-based status.
            identity = get_disk_identity(dev)
            if is_bad_disk(name):
                status_text = "BAD (NOT RUNNING)"
                offline_count += 1
            else:
                status_text = "HEALTHY (NOT RUNNING)"
        else:
            enum = cslib_get_status_enum(diskid, tool, addr)
            identity = get_disk_identity(dev)
            if enum == "DISK_OK":
                status_text = "HEALTHY"
            elif enum == "DISK_ERROR":
                status_text = "BAD"
                offline_count += 1
            elif enum is None:
                # Per-disk query failed (connection refused / timeout / bad output).
                # Mark cslib unreachable so remaining disks use marker-based fallback.
                cslib_reachable = False
                if is_bad_disk(name):
                    status_text = "BAD (NOT RUNNING)"
                    offline_count += 1
                else:
                    status_text = "HEALTHY (NOT RUNNING)"
            else:
                # DISK_READONLY / DISK_SHUTDOWN: strip DISK_ prefix for display.
                status_text = enum.replace("DISK_", "") if enum.startswith("DISK_") else enum
                degraded_count += 1

        print("{:<10} {:<25} {:<22} {}".format(name, dev, status_text, identity))

    print("")
    if not cslib_reachable:
        warn("EFC not reachable; disk status is based on local markers, not live EFC runtime state")
    # OFFLINE (DISK_ERROR, proactively offlined) is NOT a problem;
    # only MISSING / INACCESSIBLE / UNKNOWN / DISK_READONLY / DISK_SHUTDOWN
    # trigger exit 1.
    problem_count = missing_count + unknown_count + degraded_count
    if problem_count == 0:
        if offline_count:
            log("All disks healthy ({} total, {} offline)".format(total, offline_count))
        else:
            log("All disks healthy ({} total)".format(total))
    else:
        parts = []
        if missing_count:
            parts.append("{} missing/inaccessible".format(missing_count))
        if unknown_count:
            parts.append("{} unknown".format(unknown_count))
        if degraded_count:
            parts.append("{} degraded".format(degraded_count))
        tail = "({} total".format(total)
        if offline_count:
            tail += ", {} offline".format(offline_count)
        tail += ")"
        log("{} disk(s) with issues: {} {}".format(problem_count, ", ".join(parts), tail))
    return problem_count, cslib_reachable


def find_kvcs_cli():
    """Find kvcs_cli binary: check build dirs first, then PATH."""
    build_dirs = [
        os.path.join(SCRIPT_DIR, "..", "build"),
        os.path.join(SCRIPT_DIR, "..", "build", "src", "cli"),
    ]
    for directory in build_dirs:
        candidate = os.path.join(directory, "kvcs_cli")
        if os.path.isfile(candidate) and os.access(candidate, os.X_OK):
            return candidate
    return find_tool("kvcs_cli")


def show_kvcs_cli_stat():
    """Run kvcs_cli stat and print output as-is."""
    print("")
    log("=== kvcs_cli stat ===")

    cli_bin = find_kvcs_cli()
    if not cli_bin:
        warn("kvcs_cli not found, skipping stat output")
        return

    socket = get_efc_socket()
    try:
        p = subprocess.run(
            [cli_bin, "stat", "--grpc-socket={}".format(socket)],
            stderr=subprocess.DEVNULL,
            timeout=5,
        )
        rc = p.returncode
    except subprocess.TimeoutExpired:
        warn("kvcs_cli stat timed out (EFC not running?)")
        return
    if rc != 0:
        warn("kvcs_cli stat failed (rc={})".format(rc))


def cmd_status():
    log("=== Disk Status ===")
    print("")
    problem_count, cslib_ok = show_disk_status()
    if cslib_ok:
        show_kvcs_cli_stat()
    if problem_count > 0:
        sys.exit(1)


def parse_offline_args(args):
    auto_yes = False
    devices = []
    for arg in args:
        if arg in ("--yes", "-y"):
            auto_yes = True
        elif arg.startswith("-"):
            die("unknown option: {}".format(arg))
        else:
            devices.append(arg)
    if len(devices) != 1:
        die("usage: kvcs-disk-ops offline <device> [--yes]")
    return auto_yes, devices


def parse_online_args(args):
    auto_yes = False
    devices = []
    for arg in args:
        if arg in ("--yes", "-y"):
            auto_yes = True
        elif arg.startswith("-"):
            die("unknown option: {}".format(arg))
        else:
            devices.append(arg)
    if len(devices) != 1:
        die("usage: kvcs-disk-ops online <device> [--yes]")
    return auto_yes, devices


def cmd_offline(args):
    auto_yes, devices = parse_offline_args(args)
    device = devices[0]

    cfg = find_efc_config()
    if not cfg:
        die("efc.yaml not found")
    try:
        with open(cfg, "r") as fh:
            text = fh.read()
    except Exception as exc:
        die("failed to read efc.yaml: {} ({})".format(cfg, exc))

    if get_efc_backend(text) != "disk":
        die("offline only supports backend=disk")

    if not efc_is_running():
        die("EFC is not running, cannot hot-offline")

    tool = cslib_tool_path()
    addr = cslib_addr(text)

    slots, _ = parse_disk_conf_full()
    slot_info = None
    for diskid, name, dev in slots:
        if dev == device:
            slot_info = (diskid, name)
            break
    if slot_info is None:
        die("device {} not in disk config ({})".format(device, DISK_CONF))
    diskid, _ = slot_info

    # Idempotent: if the disk is already offline, exit without repeating.
    if is_offline_status(cslib_get_status(diskid, tool, addr)):
        log("device {} is already offline".format(device))
        sys.exit(0)

    online_count = 0
    for did, _, _ in slots:
        if is_online_status(cslib_get_status(did, tool, addr)):
            online_count += 1
    if online_count <= 1:
        warn("this is the last online disk; EFC will have no usable disks after offline")

    print("")
    log("about to offline device {}".format(device))
    log("disk cache data will become unavailable, EFC capacity will decrease")
    if not auto_yes:
        try:
            confirm = input("Confirm? [y/N] ")
        except EOFError:
            confirm = ""
        if confirm not in ("y", "Y"):
            log("Cancelled")
            sys.exit(0)

    lock_fd = acquire_lock()
    try:
        # New SOP: offline uses --status=DISK_ERROR.
        if cslib_set_status(diskid, "DISK_ERROR", tool, addr) != 0:
            die("failed to notify EFC offline; ensure EFC disk management channel is working")
        
        # Poll getstatus after setstatus for secondary confirmation; status propagation may have brief delay.
        if not wait_offline(diskid, tool, addr):
            rc, out, err = cslib_get_status_raw(diskid, tool, addr)
            status = parse_cslib_status(out) if rc == 0 else None
            if rc != 0:
                warn("status confirmation failed: getstatus rc={}, stderr: {}".format(
                    rc, err.strip() or "(empty)"))
            elif is_online_status(status):
                warn("status confirmation failed: getstatus still reports online ({}), disk may not have offlined".format(
                    status))
            else:
                warn("status confirmation failed: getstatus output unrecognized ({}), use status to verify".format(
                    out.strip() or "(empty)"))

        log("Device {} offlined".format(device))
    finally:
        release_lock(lock_fd)


def cmd_online(args):
    auto_yes, devices = parse_online_args(args)
    device = devices[0]

    cfg = find_efc_config()
    if not cfg:
        die("efc.yaml not found")
    try:
        with open(cfg, "r") as fh:
            text = fh.read()
    except Exception as exc:
        die("failed to read efc.yaml: {} ({})".format(cfg, exc))

    if get_efc_backend(text) != "disk":
        die("online only supports backend=disk")

    if not efc_is_running():
        die("EFC is not running, cannot hot-online")

    if not is_block_device(device):
        die("{} is not a block device".format(device))
    if is_mounted(device):
        die("{} is mounted on filesystem, operation refused".format(device))

    tool = cslib_tool_path()
    addr = cslib_addr(text)

    slots, _ = parse_disk_conf_full()

    slot_info = None
    for diskid, name, dev in slots:
        if dev == device:
            slot_info = (diskid, name)
            break
    if slot_info is None:
        die("device {} not in disk config ({}); online only supports restoring the same offlined disk".format(device, DISK_CONF))
    diskid, slot_name = slot_info

    print("")
    log("about to online device {}".format(device))
    log("will perform full-disk TRIM and format, ALL data on this disk will be erased")
    if not auto_yes:
        try:
            confirm = input("Confirm? [y/N] ")
        except EOFError:
            confirm = ""
        if confirm not in ("y", "Y"):
            log("Cancelled")
            sys.exit(0)

    lock_fd = acquire_lock()
    try:
        # Re-check status under lock to avoid TOCTOU race
        status = cslib_get_status(diskid, tool, addr)
        if is_online_status(status):
            die("device {} (slot {}) is already online, no action needed".format(device, slot_name))
        if not is_offline_status(status):
            die("device {} (slot {}) is not in offline state, cannot online; online only supports restoring the same offlined disk".format(device, slot_name))

        format_disk(device)

        if cslib_notify_online(diskid, tool, addr) != 0:
            die("failed to notify online; ensure EFC disk management channel is working")

        # wait_online polls with interval=2s, giving cslib time for state propagation
        if not wait_online(diskid, tool, addr):
            warn("online confirmation timeout; online notified, use status to verify")

        marker = marker_path(slot_name)
        if os.path.isfile(marker):
            try:
                os.remove(marker)
                log("cleaned up stale bad-disk marker")
            except OSError as exc:
                warn("failed to clean stale bad-disk marker: {}".format(exc))

        log("Device {} onlined".format(device))
    finally:
        release_lock(lock_fd)


def parse_recover_args(args):
    auto_yes = False
    disks = []
    for arg in args:
        if arg in ("--yes", "-y"):
            auto_yes = True
        elif arg.startswith("-"):
            die("unknown option: {}".format(arg))
        else:
            disks.append(arg)
    if not disks:
        die("usage: kvcs-disk-ops recover <disk>... [--yes]")
    return auto_yes, disks


def cmd_recover(args):
    auto_yes, disks = parse_recover_args(args)

    log("=== Current Disk Status ===")
    print("")
    show_disk_status()
    print("")

    for name, _dev in parse_disk_conf():
        if not is_bad_disk(name):
            continue
        if name not in disks:
            warn("{} is also a bad disk but not included in this recovery".format(name))

    log("=== Validating Target Disks ===")
    disk_devs = []
    for disk_name in disks:
        dev = get_device(disk_name)
        if not dev:
            die("{} not in disk config ({})".format(disk_name, DISK_CONF))
        if not is_bad_disk(disk_name):
            die("{} has no bad-disk marker, no recovery needed".format(disk_name))
        if not is_block_device(dev):
            die("{} is not a block device (disk not replaced yet?)".format(dev))
        disk_devs.append(dev)
        log("  {} -> {} [bad-disk marker present, device accessible]".format(disk_name, dev))

    if efc_is_running():
        die("EFC is still running, stop EFC before recover")

    if not auto_yes:
        print("")
        print("The following operations will be performed:")
        for disk_name, dev in zip(disks, disk_devs):
            print("  - Format {} ({})".format(disk_name, dev))
        print("")
        print("Prerequisite: EFC stopped ✓")
        print("WARNING: formatting will erase ALL data on the device(s)!")
        try:
            confirm = input("Confirm? [y/N] ")
        except EOFError:
            confirm = ""
        if confirm not in ("y", "Y"):
            log("Cancelled")
            sys.exit(0)

    log("=== Formatting ===")
    for disk_name, dev in zip(disks, disk_devs):
        log("formatting {} ({}) ...".format(disk_name, dev))
        format_disk(dev, idempotent=True)

    log("=== Removing Bad-Disk Markers ===")
    for disk_name in disks:
        path = marker_path(disk_name)
        try:
            os.remove(path)
        except FileNotFoundError:
            pass
        log("  deleted PanguCsDiskStatus{}".format(disk_name))

    log("=== Verification ===")
    new_bad = list_bad_disks()
    if new_bad:
        die("new bad-disk markers found: {} (new disk may be faulty)".format(" ".join(new_bad)))
    log("no new bad-disk markers")

    print("")
    log("=== Recovery Complete ===")
    for disk_name, dev in zip(disks, disk_devs):
        log("  {} ({}): formatted, bad-disk marker removed".format(disk_name, dev))
    print("")
    log("Next steps:")
    log("  1. Start EFC service")
    log("  2. Run kvcs-disk-ops status to verify recovery")


def usage():
    print("""kvcs-disk-ops - EFC disk operations tool

Usage:
  kvcs-disk-ops status                  Show all disk status and EFC runtime status
  kvcs-disk-ops offline <device> [-y]   Hot-offline a device (runtime only, no efc.yaml change)
  kvcs-disk-ops online  <device> [-y]   Hot-online a device (only the same offlined disk)
  kvcs-disk-ops recover <disk>... [-y]  Format bad disk(s) and remove bad-disk markers

Options:
  --yes, -y      Skip confirmation prompt (for automation)
  --help, -h     Show this help
  --version, -V  Show version and build info

Environment Variables:
  EFC_CONFIG      Path to efc.yaml (default /etc/kvcs/efc.yaml or sibling efc.yaml)
  KVCS_ROOT       EFC installation root, used to find spdk_fs_tool / kvcs_cli / pangu_cslib_tool

Hot-offline flow:
  1. Notify EFC to mark the disk offline (runtime only, no efc.yaml / disk config change)
  2. Verify status

Hot-online flow:
  1. Verify target device slot is in offline state (can only restore the same offlined disk)
  2. Perform full-disk TRIM and format
  3. Notify EFC to reload the disk
  4. Verify online status

Recovery flow:
  1. Confirm new disk is replaced (lsblk shows target block device)
  2. Stop EFC (systemctl stop kvcs-efc)
  3. kvcs-disk-ops recover <disk> --yes
  4. Start EFC (systemctl start kvcs-efc)
  5. kvcs-disk-ops status to verify

Full operation flow (with process management):

Recover after stopping EFC:
  kvcs-ctl stop kvcs-efc
  kvcs-disk-ops recover disk2 --yes
  kvcs-ctl start kvcs-efc

Hot-swap while EFC is running:
  kvcs-disk-ops offline /dev/nvme2n1 --yes
  # Replace physical disk
  kvcs-disk-ops online /dev/nvme2n1 --yes
  kvcs-disk-ops status

Examples:
  kvcs-disk-ops status
  kvcs-disk-ops offline /dev/nvme2n1
  kvcs-disk-ops online /dev/nvme2n1
  kvcs-disk-ops recover disk2
  kvcs-disk-ops recover disk1 disk2 --yes""")
    sys.exit(0)


def main(argv):
    command = argv[1] if len(argv) > 1 else ""
    if command == "status":
        cmd_status()
    elif command == "recover":
        cmd_recover(argv[2:])
    elif command == "offline":
        cmd_offline(argv[2:])
    elif command == "online":
        cmd_online(argv[2:])
    elif command in ("--help", "-h", ""):
        usage()
    elif command in ("--version", "-V"):
        print_version()
    else:
        die("unknown command: {} (use --help for usage)".format(command))


if __name__ == "__main__":
    main(sys.argv)
