from __future__ import annotations

import hashlib
import json
import os
import re
import shutil
import subprocess
import time
from dataclasses import asdict, dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Mapping


PROBE_VERSION = "1"
DEFAULT_CACHE_TTL_SEC = 12 * 60 * 60
DEFAULT_TIMEOUT_SEC = 20
DEFAULT_MAX_TARGETS_PER_NODE = 6
_PROBE_FILE_NAME = ".knoe-storage-probe"

_SYNOLOGY_RE = re.compile(r"^/synology/d(\d{3})$")
_DD_SUMMARY_RE = re.compile(
    r"(?P<bytes>\d+)\s+bytes.*?(?:,\s*|\sin\s+)(?P<seconds>[0-9.]+)\s*s"
)
_UNSAFE_MOUNTPOINTS = {
    "/",
    "/boot",
    "/boot/efi",
    "/etc",
    "/proc",
    "/sys",
    "/dev",
    "/run",
}


@dataclass(slots=True)
class StorageProbeTarget:
    node_name: str
    source_type: str
    mountpoint: str
    device: str | None
    filesystem: str | None
    storage_class: str | None
    label: str
    path_exists: bool | None
    writable: bool | None
    notes: list[str] = field(default_factory=list)


@dataclass(slots=True)
class StorageProbeMetrics:
    randread_iops: float | None = None
    randwrite_iops: float | None = None
    read_bw_bytes: int | None = None
    write_bw_bytes: int | None = None
    read_lat_ms_p95: float | None = None
    write_lat_ms_p95: float | None = None
    test_runtime_sec: float | None = None


@dataclass(slots=True)
class StorageProbeResult:
    target_id: str
    node_name: str
    label: str
    mountpoint: str
    source_type: str
    storage_class: str | None
    succeeded: bool
    skipped: bool
    error: str | None
    metrics: StorageProbeMetrics
    io_class: str | None
    timing_multiplier: float | None
    observed_at: str | None
    approximate: bool = False
    cached: bool = False
    notes: list[str] = field(default_factory=list)


@dataclass(slots=True)
class NodeStorageInventory:
    node_name: str
    items: list[StorageProbeResult] = field(default_factory=list)


def io_class_timing_multiplier(io_class: str | None) -> float | None:
    if io_class == "io_fast_local":
        return 1.0
    if io_class == "io_general":
        return 1.5
    if io_class == "io_slow_durable":
        return 3.0
    if io_class == "io_cold":
        return 5.0
    return None


def _target_id(target: StorageProbeTarget, *, quick: bool) -> str:
    payload = {
        "version": PROBE_VERSION,
        "node_name": target.node_name,
        "source_type": target.source_type,
        "mountpoint": target.mountpoint,
        "device": target.device,
        "filesystem": target.filesystem,
        "storage_class": target.storage_class,
        "mode": "quick" if quick else "full",
    }
    encoded = json.dumps(payload, separators=(",", ":"), sort_keys=True).encode("utf-8")
    return hashlib.sha256(encoded).hexdigest()[:16]


def _path_exists_writable(path: Path) -> tuple[bool, bool]:
    exists = path.exists() and path.is_dir()
    writable = exists and os.access(path, os.W_OK)
    return exists, writable


def _source_type_for_filesystem(mountpoint: str, fs: Mapping[str, Any]) -> tuple[str, str]:
    mp = mountpoint.strip()
    m = _SYNOLOGY_RE.match(mp)
    if m:
        return "synology_mount", f"Synology d{m.group(1)}"
    if mp == "/knoe/d004":
        return "pvc_base_dir", "Monitoring data"
    if mp.startswith("/var/lib/rancher"):
        return "rancher_storage", "k3s rancher storage"
    if mp.startswith("/var/lib/kubelet"):
        return "pvc_base_dir", "Kubernetes local PV base"
    if mp.startswith("/var/lib"):
        return "host_path", "General filesystem /var/lib"

    storage_class = str(fs.get("storageClass") or "").strip().lower()
    transport = str(fs.get("transport") or "").strip().lower()
    rotational = fs.get("rotational")
    if storage_class in {"nvme", "ssd", "local-ssd"} or (transport == "nvme" and rotational is False):
        return "filesystem_mount", "Local SSD mount"
    return "generic", f"General filesystem {mp}"


def _iter_filesystem_like(node: Any) -> list[Mapping[str, Any]]:
    raw = getattr(node, "filesystems", None)
    if raw is None:
        return []
    rows: list[Mapping[str, Any]] = []
    for item in raw:
        if isinstance(item, Mapping):
            rows.append(item)
            continue
        rows.append(
            {
                "device": getattr(item, "device", None),
                "mountpoint": getattr(item, "mountpoint", None),
                "fstype": getattr(item, "fstype", None),
                "storageClass": getattr(item, "storage_class", None),
                "transport": getattr(item, "transport", None),
                "rotational": getattr(item, "rotational", None),
            }
        )
    return rows


def discover_storage_probe_targets(node: Any, topology_context: Mapping[str, Any] | None = None) -> list[StorageProbeTarget]:
    del topology_context
    node_name = str(getattr(node, "name", "") or "").strip() or "unknown"
    filesystems = _iter_filesystem_like(node)
    targets: list[StorageProbeTarget] = []
    seen_mountpoints: set[str] = set()

    for fs in filesystems:
        mountpoint = str(fs.get("mountpoint") or "").strip()
        if not mountpoint or mountpoint in seen_mountpoints:
            continue
        if mountpoint in _UNSAFE_MOUNTPOINTS:
            continue
        source_type, label = _source_type_for_filesystem(mountpoint, fs)
        if source_type == "generic" and not mountpoint.startswith("/var/lib"):
            # Keep discovery conservative; prefer known Knoe-relevant roots.
            continue

        mount_path = Path(mountpoint)
        exists, writable = _path_exists_writable(mount_path)
        notes: list[str] = []
        if not exists:
            notes.append("path_missing")
        elif not writable:
            notes.append("not_writable")

        targets.append(
            StorageProbeTarget(
                node_name=node_name,
                source_type=source_type,
                mountpoint=mountpoint,
                device=(str(fs.get("device")).strip() or None) if fs.get("device") is not None else None,
                filesystem=(str(fs.get("fstype")).strip() or None) if fs.get("fstype") is not None else None,
                storage_class=(str(fs.get("storageClass")).strip() or None)
                if fs.get("storageClass") is not None
                else None,
                label=label,
                path_exists=exists,
                writable=writable,
                notes=notes,
            )
        )
        seen_mountpoints.add(mountpoint)

    for fallback_mp, fallback_label, fallback_type in (
        ("/synology/d004", "Synology d004", "synology_mount"),
        ("/knoe/d004", "Monitoring data", "pvc_base_dir"),
        ("/var/lib/rancher/k3s/storage", "k3s rancher storage", "rancher_storage"),
        ("/var/lib", "General filesystem /var/lib", "host_path"),
    ):
        if fallback_mp in seen_mountpoints:
            continue
        exists, writable = _path_exists_writable(Path(fallback_mp))
        if not exists:
            continue
        targets.append(
            StorageProbeTarget(
                node_name=node_name,
                source_type=fallback_type,
                mountpoint=fallback_mp,
                device=None,
                filesystem=None,
                storage_class=None,
                label=fallback_label,
                path_exists=exists,
                writable=writable,
                notes=[] if writable else ["not_writable"],
            )
        )
        seen_mountpoints.add(fallback_mp)

    return sorted(targets, key=lambda item: (item.source_type, item.mountpoint))


def parse_fio_json_metrics(payload: str | Mapping[str, Any]) -> StorageProbeMetrics:
    if isinstance(payload, str):
        try:
            body: Mapping[str, Any] = json.loads(payload)
        except json.JSONDecodeError:
            return StorageProbeMetrics()
    else:
        body = payload

    jobs = body.get("jobs") if isinstance(body, Mapping) else None
    if not isinstance(jobs, list):
        return StorageProbeMetrics()

    metrics = StorageProbeMetrics()
    for job in jobs:
        if not isinstance(job, Mapping):
            continue
        name = str(job.get("jobname") or "").lower()
        read = job.get("read") if isinstance(job.get("read"), Mapping) else {}
        write = job.get("write") if isinstance(job.get("write"), Mapping) else {}

        if "read" in name:
            metrics.randread_iops = _as_float(read.get("iops"))
            metrics.read_bw_bytes = _as_int(read.get("bw_bytes"))
            metrics.read_lat_ms_p95 = _fio_lat_p95_ms(read)
        if "write" in name:
            metrics.randwrite_iops = _as_float(write.get("iops"))
            metrics.write_bw_bytes = _as_int(write.get("bw_bytes"))
            metrics.write_lat_ms_p95 = _fio_lat_p95_ms(write)

    return metrics


def _fio_lat_p95_ms(section: Mapping[str, Any]) -> float | None:
    clat_ns = section.get("clat_ns")
    if not isinstance(clat_ns, Mapping):
        return None
    pct = clat_ns.get("percentile")
    if not isinstance(pct, Mapping):
        return None
    val = pct.get("95.000000")
    out = _as_float(val)
    if out is None:
        return None
    return out / 1_000_000.0


def _as_float(value: Any) -> float | None:
    if value is None:
        return None
    try:
        return float(value)
    except (TypeError, ValueError):
        return None


def _as_int(value: Any) -> int | None:
    if value is None:
        return None
    try:
        return int(value)
    except (TypeError, ValueError):
        return None


def _empty_result(target: StorageProbeTarget, *, quick: bool) -> StorageProbeResult:
    return StorageProbeResult(
        target_id=_target_id(target, quick=quick),
        node_name=target.node_name,
        label=target.label,
        mountpoint=target.mountpoint,
        source_type=target.source_type,
        storage_class=target.storage_class,
        succeeded=False,
        skipped=False,
        error=None,
        metrics=StorageProbeMetrics(),
        io_class=None,
        timing_multiplier=None,
        observed_at=None,
        approximate=False,
        cached=False,
        notes=list(target.notes),
    )


def classify_storage_io(result: StorageProbeResult) -> tuple[str | None, float | None]:
    if result.skipped or not result.succeeded:
        return "unknown", None

    rr = result.metrics.randread_iops
    rw = result.metrics.randwrite_iops
    rlat = result.metrics.read_lat_ms_p95
    wlat = result.metrics.write_lat_ms_p95

    if rr is None or rw is None:
        return "io_cold", io_class_timing_multiplier("io_cold")

    if rlat is None or wlat is None:
        max_p95 = float("inf")
    else:
        max_p95 = max(rlat, wlat)

    if rr > 8000 and rw > 4000 and max_p95 < 3:
        return "io_fast_local", io_class_timing_multiplier("io_fast_local")
    if rr > 2000 and rw > 1000 and max_p95 < 12:
        return "io_general", io_class_timing_multiplier("io_general")
    if rr > 300 and rw > 150 and max_p95 < 50:
        return "io_slow_durable", io_class_timing_multiplier("io_slow_durable")
    return "io_cold", io_class_timing_multiplier("io_cold")


def run_storage_probe(target: StorageProbeTarget, timeout_sec: int = DEFAULT_TIMEOUT_SEC, quick: bool = True) -> StorageProbeResult:
    started = time.monotonic()
    observed_at = datetime.now(timezone.utc).isoformat()
    result = _empty_result(target, quick=quick)
    result.observed_at = observed_at

    mount_path = Path(target.mountpoint)
    exists, writable = _path_exists_writable(mount_path)
    if not exists:
        result.skipped = True
        result.error = "mountpoint_not_found"
        result.notes.append("path_missing")
    elif not writable:
        result.skipped = True
        result.error = "mountpoint_not_writable"
        result.notes.append("not_writable")
    elif target.mountpoint in _UNSAFE_MOUNTPOINTS:
        result.skipped = True
        result.error = "unsafe_mountpoint"

    if result.skipped:
        result.metrics.test_runtime_sec = round(time.monotonic() - started, 3)
        result.io_class, result.timing_multiplier = classify_storage_io(result)
        return result

    test_file = mount_path / _PROBE_FILE_NAME
    fio_bin = shutil.which("fio")

    try:
        if fio_bin:
            read_metrics = _run_fio_job(fio_bin, test_file, rw="randread", timeout_sec=timeout_sec)
            write_metrics = _run_fio_job(fio_bin, test_file, rw="randwrite", timeout_sec=timeout_sec)
            result.metrics.randread_iops = read_metrics.randread_iops
            result.metrics.read_bw_bytes = read_metrics.read_bw_bytes
            result.metrics.read_lat_ms_p95 = read_metrics.read_lat_ms_p95
            result.metrics.randwrite_iops = write_metrics.randwrite_iops
            result.metrics.write_bw_bytes = write_metrics.write_bw_bytes
            result.metrics.write_lat_ms_p95 = write_metrics.write_lat_ms_p95
            result.succeeded = bool(
                result.metrics.randread_iops is not None
                and result.metrics.randwrite_iops is not None
            )
            if not result.succeeded:
                result.error = "fio_output_missing_metrics"
        else:
            result.notes.append("fio_unavailable")
            result.approximate = True
            _run_dd_fallback(test_file, result, timeout_sec=timeout_sec)
    except subprocess.TimeoutExpired:
        result.error = "probe_timeout"
        result.succeeded = False
    except Exception as exc:
        result.error = str(exc)
        result.succeeded = False
    finally:
        try:
            if test_file.exists():
                test_file.unlink()
        except Exception:
            pass

    result.metrics.test_runtime_sec = round(time.monotonic() - started, 3)
    result.io_class, result.timing_multiplier = classify_storage_io(result)
    return result


def _run_fio_job(fio_bin: str, test_file: Path, *, rw: str, timeout_sec: int) -> StorageProbeMetrics:
    cmd = [
        fio_bin,
        "--name",
        rw,
        "--filename",
        str(test_file),
        "--rw",
        rw,
        "--bs",
        "4k",
        "--ioengine",
        "libaio",
        "--iodepth",
        "1",
        "--size",
        "256M",
        "--direct",
        "1",
        "--time_based",
        "1",
        "--runtime",
        "8",
        "--ramp_time",
        "1",
        "--group_reporting",
        "1",
        "--output-format=json",
    ]
    res = subprocess.run(cmd, capture_output=True, text=True, timeout=max(3, timeout_sec))
    if res.returncode != 0:
        stderr = (res.stderr or "").strip() or "fio_failed"
        raise RuntimeError(stderr)
    return parse_fio_json_metrics(res.stdout or "{}")


def _run_dd_fallback(test_file: Path, result: StorageProbeResult, *, timeout_sec: int) -> None:
    write = subprocess.run(
        [
            "dd",
            "if=/dev/zero",
            f"of={test_file}",
            "bs=4M",
            "count=64",
            "oflag=direct",
            "conv=fsync",
        ],
        capture_output=True,
        text=True,
        timeout=max(3, timeout_sec),
    )
    if write.returncode != 0:
        raise RuntimeError((write.stderr or "").strip() or "dd_write_failed")
    write_bytes, write_seconds = _parse_dd_summary(write.stderr or "")

    read = subprocess.run(
        [
            "dd",
            f"if={test_file}",
            "of=/dev/null",
            "bs=4M",
            "count=64",
            "iflag=direct",
        ],
        capture_output=True,
        text=True,
        timeout=max(3, timeout_sec),
    )
    if read.returncode != 0:
        raise RuntimeError((read.stderr or "").strip() or "dd_read_failed")
    read_bytes, read_seconds = _parse_dd_summary(read.stderr or "")

    if write_bytes and write_seconds and write_seconds > 0:
        write_bw = int(write_bytes / write_seconds)
        result.metrics.write_bw_bytes = write_bw
        result.metrics.randwrite_iops = float(write_bw / 4096.0)
    if read_bytes and read_seconds and read_seconds > 0:
        read_bw = int(read_bytes / read_seconds)
        result.metrics.read_bw_bytes = read_bw
        result.metrics.randread_iops = float(read_bw / 4096.0)

    result.succeeded = bool(result.metrics.randread_iops and result.metrics.randwrite_iops)
    if not result.succeeded:
        result.error = "dd_output_missing_metrics"
    result.notes.append("approximate_dd_metrics")


def _parse_dd_summary(stderr: str) -> tuple[int | None, float | None]:
    for line in reversed((stderr or "").splitlines()):
        m = _DD_SUMMARY_RE.search(line)
        if not m:
            continue
        try:
            return int(m.group("bytes")), float(m.group("seconds"))
        except (TypeError, ValueError):
            return None, None
    return None, None


def _cache_read(path: Path) -> dict[str, Any]:
    if not path.exists():
        return {}
    try:
        data = json.loads(path.read_text(encoding="utf-8"))
    except Exception:
        return {}
    return data if isinstance(data, dict) else {}


def _cache_write(path: Path, body: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(body, separators=(",", ":"), sort_keys=True), encoding="utf-8")


def _result_to_cache_value(result: StorageProbeResult) -> dict[str, Any]:
    out = asdict(result)
    out["cached"] = False
    return out


def _result_from_cache_value(value: Mapping[str, Any]) -> StorageProbeResult | None:
    try:
        metrics_raw = value.get("metrics")
        if not isinstance(metrics_raw, Mapping):
            metrics_raw = {}
        metrics = StorageProbeMetrics(
            randread_iops=_as_float(metrics_raw.get("randread_iops")),
            randwrite_iops=_as_float(metrics_raw.get("randwrite_iops")),
            read_bw_bytes=_as_int(metrics_raw.get("read_bw_bytes")),
            write_bw_bytes=_as_int(metrics_raw.get("write_bw_bytes")),
            read_lat_ms_p95=_as_float(metrics_raw.get("read_lat_ms_p95")),
            write_lat_ms_p95=_as_float(metrics_raw.get("write_lat_ms_p95")),
            test_runtime_sec=_as_float(metrics_raw.get("test_runtime_sec")),
        )
        return StorageProbeResult(
            target_id=str(value.get("target_id") or ""),
            node_name=str(value.get("node_name") or ""),
            label=str(value.get("label") or ""),
            mountpoint=str(value.get("mountpoint") or ""),
            source_type=str(value.get("source_type") or "generic"),
            storage_class=(str(value.get("storage_class")).strip() or None)
            if value.get("storage_class") is not None
            else None,
            succeeded=bool(value.get("succeeded")),
            skipped=bool(value.get("skipped")),
            error=(str(value.get("error")).strip() or None) if value.get("error") else None,
            metrics=metrics,
            io_class=(str(value.get("io_class")).strip() or None) if value.get("io_class") else None,
            timing_multiplier=_as_float(value.get("timing_multiplier")),
            observed_at=(str(value.get("observed_at")).strip() or None) if value.get("observed_at") else None,
            approximate=bool(value.get("approximate")),
            cached=True,
            notes=[str(item) for item in (value.get("notes") or []) if str(item).strip()],
        )
    except Exception:
        return None


def _cache_get(cache_body: dict[str, Any], key: str, *, ttl_sec: int) -> StorageProbeResult | None:
    records = cache_body.get("records")
    if not isinstance(records, dict):
        return None
    item = records.get(key)
    if not isinstance(item, Mapping):
        return None
    observed_epoch = _as_float(item.get("observed_epoch"))
    if observed_epoch is None:
        return None
    if (time.time() - observed_epoch) > max(0, ttl_sec):
        return None
    result_raw = item.get("result")
    if not isinstance(result_raw, Mapping):
        return None
    return _result_from_cache_value(result_raw)


def _cache_put(cache_body: dict[str, Any], key: str, result: StorageProbeResult) -> None:
    records = cache_body.setdefault("records", {})
    if not isinstance(records, dict):
        records = {}
        cache_body["records"] = records
    records[key] = {
        "observed_epoch": time.time(),
        "result": _result_to_cache_value(result),
    }


def probe_node_storage(
    node: Any,
    *,
    topology_context: Mapping[str, Any] | None = None,
    timeout_sec: int = DEFAULT_TIMEOUT_SEC,
    quick: bool = True,
    skip: bool = False,
    force_refresh: bool = False,
    cache_path: Path | None = None,
    cache_ttl_sec: int = DEFAULT_CACHE_TTL_SEC,
    max_targets: int = DEFAULT_MAX_TARGETS_PER_NODE,
) -> NodeStorageInventory:
    node_name = str(getattr(node, "name", "") or "").strip() or "unknown"
    targets = discover_storage_probe_targets(node, topology_context)
    if max_targets > 0:
        targets = targets[:max_targets]

    if skip:
        out: list[StorageProbeResult] = []
        for target in targets:
            res = _empty_result(target, quick=quick)
            res.skipped = True
            res.error = "probe_skipped"
            res.observed_at = datetime.now(timezone.utc).isoformat()
            res.io_class, res.timing_multiplier = classify_storage_io(res)
            out.append(res)
        return NodeStorageInventory(node_name=node_name, items=out)

    cache_file = cache_path
    cache_body: dict[str, Any] = {}
    if cache_file is not None:
        cache_body = _cache_read(cache_file)

    out_items: list[StorageProbeResult] = []
    cache_changed = False
    for target in targets:
        key = _target_id(target, quick=quick)
        cached = None
        if not force_refresh and cache_file is not None:
            cached = _cache_get(cache_body, key, ttl_sec=cache_ttl_sec)
        if cached is not None:
            out_items.append(cached)
            continue

        probed = run_storage_probe(target, timeout_sec=timeout_sec, quick=quick)
        out_items.append(probed)
        if cache_file is not None:
            _cache_put(cache_body, key, probed)
            cache_changed = True

    if cache_file is not None and cache_changed:
        _cache_write(cache_file, cache_body)

    return NodeStorageInventory(node_name=node_name, items=out_items)


def storage_result_to_dict(result: StorageProbeResult) -> dict[str, Any]:
    return asdict(result)