"""Storage device probing for distributed mesh JEPA optimizer.

Probes storage devices for sequential/random IO throughput, queue depth
scaling, and fsync latency. Results feed into the JEPA-based system optimizer.
"""
from __future__ import annotations

import concurrent.futures
import json
import logging
import os
import random
import shutil
import subprocess
import tempfile
import time
from typing import Any, Dict, List, Optional

import numpy as np

logger = logging.getLogger(__name__)


class StorageProbe:
    """Probe storage devices for performance characteristics."""

    def detect_devices(self) -> List[Dict]:
        """Find block devices via lsblk --json.

        Returns list of dicts with name, type, size, mountpoint, model,
        scheduler for each device.
        """
        try:
            result = subprocess.run(
                ["lsblk", "--json", "-o",
                 "NAME,TYPE,SIZE,MOUNTPOINT,MODEL,SCHED"],
                capture_output=True, text=True, timeout=10,
            )
            if result.returncode != 0:
                logger.warning("lsblk failed: %s", result.stderr.strip())
                return []

            data = json.loads(result.stdout)
            devices = []
            for dev in data.get("blockdevices", []):
                devices.append({
                    "name": dev.get("name", ""),
                    "type": dev.get("type", ""),
                    "size": dev.get("size", ""),
                    "mountpoint": dev.get("mountpoint"),
                    "model": (dev.get("model") or "").strip() or None,
                    "scheduler": dev.get("sched"),
                })
                # Include partitions / children
                for child in dev.get("children", []):
                    devices.append({
                        "name": child.get("name", ""),
                        "type": child.get("type", ""),
                        "size": child.get("size", ""),
                        "mountpoint": child.get("mountpoint"),
                        "model": (child.get("model") or dev.get("model") or "").strip() or None,
                        "scheduler": child.get("sched") or dev.get("sched"),
                    })

            logger.info("Detected %d block devices/partitions", len(devices))
            return devices

        except FileNotFoundError:
            logger.warning("lsblk not found")
            return []
        except Exception as exc:
            logger.warning("Device detection failed: %s", exc)
            return []

    def probe_sequential(self, path: str, size_mb: int = 256) -> Dict:
        """Measure sequential read/write throughput using dd with direct IO.

        Args:
            path: Directory to create test files in.
            size_mb: Size of test file in megabytes.

        Returns:
            Dict with read_mbps, write_mbps.
        """
        os.makedirs(path, exist_ok=True)
        tmpfile = os.path.join(path, ".storage_probe_seq")
        result = {"read_mbps": 0.0, "write_mbps": 0.0}

        try:
            # --- Write with dd ---
            write_cmd = [
                "dd", f"if=/dev/zero", f"of={tmpfile}",
                f"bs=1M", f"count={size_mb}",
                "conv=fdatasync", "oflag=direct",
            ]
            try:
                proc = subprocess.run(
                    write_cmd, capture_output=True, text=True, timeout=120,
                )
                # dd reports speed on stderr
                result["write_mbps"] = self._parse_dd_speed(proc.stderr)
            except Exception as exc:
                # Fallback: oflag=direct may not work on all filesystems
                logger.debug("dd direct write failed, retrying without oflag=direct: %s", exc)
                write_cmd_fallback = [
                    "dd", "if=/dev/zero", f"of={tmpfile}",
                    "bs=1M", f"count={size_mb}", "conv=fdatasync",
                ]
                try:
                    proc = subprocess.run(
                        write_cmd_fallback, capture_output=True, text=True, timeout=120,
                    )
                    result["write_mbps"] = self._parse_dd_speed(proc.stderr)
                except Exception as exc2:
                    logger.warning("Sequential write probe failed: %s", exc2)

            if not os.path.exists(tmpfile):
                return result

            # Drop caches (best-effort, needs root)
            try:
                subprocess.run(
                    ["sh", "-c", "echo 3 > /proc/sys/vm/drop_caches"],
                    capture_output=True, timeout=5,
                )
            except Exception:
                pass

            # --- Read with dd ---
            read_cmd = [
                "dd", f"if={tmpfile}", "of=/dev/null",
                "bs=1M", "iflag=direct",
            ]
            try:
                proc = subprocess.run(
                    read_cmd, capture_output=True, text=True, timeout=120,
                )
                result["read_mbps"] = self._parse_dd_speed(proc.stderr)
            except Exception:
                # Fallback without iflag=direct
                read_cmd_fallback = [
                    "dd", f"if={tmpfile}", "of=/dev/null", "bs=1M",
                ]
                try:
                    proc = subprocess.run(
                        read_cmd_fallback, capture_output=True, text=True, timeout=120,
                    )
                    result["read_mbps"] = self._parse_dd_speed(proc.stderr)
                except Exception as exc:
                    logger.warning("Sequential read probe failed: %s", exc)

        finally:
            try:
                os.unlink(tmpfile)
            except OSError:
                pass

        logger.info(
            "Sequential %d MB -- write %.1f MB/s, read %.1f MB/s",
            size_mb, result["write_mbps"], result["read_mbps"],
        )
        return result

    @staticmethod
    def _parse_dd_speed(stderr: str) -> float:
        """Extract MB/s or GB/s from dd stderr output."""
        # dd output examples:
        #   "268435456 bytes (268 MB, 256 MiB) copied, 0.245s, 1095 MB/s"
        #   "... 1.1 GB/s"
        import re
        # Try GB/s first
        m = re.search(r"([\d.]+)\s*GB/s", stderr)
        if m:
            return round(float(m.group(1)) * 1024, 1)
        m = re.search(r"([\d.]+)\s*MB/s", stderr)
        if m:
            return round(float(m.group(1)), 1)
        # Try kB/s
        m = re.search(r"([\d.]+)\s*kB/s", stderr)
        if m:
            return round(float(m.group(1)) / 1024, 1)
        return 0.0

    def probe_random(
        self,
        path: str,
        size_mb: int = 64,
        block_size: int = 4096,
        count: int = 10000,
    ) -> Dict:
        """Measure random 4K read/write IOPS.

        Uses fio if available, falls back to manual random seeks with
        os.pread/os.pwrite.

        Args:
            path: Directory for test files.
            size_mb: Size of test file in MB.
            block_size: Block size for random IO (bytes).
            count: Number of IO operations.

        Returns:
            Dict with read_iops, write_iops, read_latency_us, write_latency_us.
        """
        result = {
            "read_iops": 0.0,
            "write_iops": 0.0,
            "read_latency_us": 0.0,
            "write_latency_us": 0.0,
        }

        # Try fio first
        if shutil.which("fio"):
            try:
                return self._probe_random_fio(path, size_mb, block_size, count)
            except Exception as exc:
                logger.debug("fio random probe failed, falling back to manual: %s", exc)

        # Manual fallback using os.pread / os.pwrite
        return self._probe_random_manual(path, size_mb, block_size, count)

    def _probe_random_fio(
        self, path: str, size_mb: int, block_size: int, count: int,
    ) -> Dict:
        """Random IO probe using fio."""
        os.makedirs(path, exist_ok=True)
        result = {
            "read_iops": 0.0, "write_iops": 0.0,
            "read_latency_us": 0.0, "write_latency_us": 0.0,
        }

        for rw, iops_key, lat_key in [
            ("randread", "read_iops", "read_latency_us"),
            ("randwrite", "write_iops", "write_latency_us"),
        ]:
            fio_cmd = [
                "fio",
                f"--name=probe_{rw}",
                f"--directory={path}",
                f"--rw={rw}",
                f"--bs={block_size}",
                f"--size={size_mb}m",
                f"--numjobs=1",
                f"--iodepth=1",
                f"--io_size={count * block_size}",
                "--direct=1",
                "--output-format=json",
                "--group_reporting",
            ]
            try:
                proc = subprocess.run(
                    fio_cmd, capture_output=True, text=True, timeout=120,
                )
                data = json.loads(proc.stdout)
                job = data["jobs"][0]
                key = "read" if "read" in rw else "write"
                result[iops_key] = round(job[key]["iops"], 1)
                # fio reports latency in nanoseconds
                result[lat_key] = round(job[key]["lat_ns"]["mean"] / 1000, 1)
            except Exception as exc:
                logger.warning("fio %s failed: %s", rw, exc)

        # Clean up fio files
        for f in os.listdir(path):
            if f.startswith("probe_rand"):
                try:
                    os.unlink(os.path.join(path, f))
                except OSError:
                    pass

        logger.info(
            "Random IO (fio) -- read %.0f IOPS (%.1f us), write %.0f IOPS (%.1f us)",
            result["read_iops"], result["read_latency_us"],
            result["write_iops"], result["write_latency_us"],
        )
        return result

    def _probe_random_manual(
        self, path: str, size_mb: int, block_size: int, count: int,
    ) -> Dict:
        """Random IO probe using os.pread/os.pwrite."""
        os.makedirs(path, exist_ok=True)
        file_size = size_mb * 1024 * 1024
        max_offset = file_size - block_size
        tmpfile = os.path.join(path, ".storage_probe_rand")
        result = {
            "read_iops": 0.0, "write_iops": 0.0,
            "read_latency_us": 0.0, "write_latency_us": 0.0,
        }

        if max_offset <= 0:
            logger.warning("File too small for random IO probe")
            return result

        try:
            # Create test file
            with open(tmpfile, "wb") as f:
                chunk = os.urandom(min(1024 * 1024, file_size))
                written = 0
                while written < file_size:
                    to_write = min(len(chunk), file_size - written)
                    f.write(chunk[:to_write])
                    written += to_write
                f.flush()
                os.fsync(f.fileno())

            # Generate random aligned offsets
            offsets = [(random.randint(0, max_offset // block_size) * block_size)
                       for _ in range(count)]
            data_block = os.urandom(block_size)

            # --- Random Read ---
            fd = os.open(tmpfile, os.O_RDONLY)
            try:
                latencies = []
                t0_total = time.perf_counter()
                for off in offsets:
                    t0 = time.perf_counter()
                    os.pread(fd, block_size, off)
                    latencies.append(time.perf_counter() - t0)
                elapsed = time.perf_counter() - t0_total
                if elapsed > 0:
                    result["read_iops"] = round(count / elapsed, 1)
                if latencies:
                    result["read_latency_us"] = round(
                        float(np.mean(latencies)) * 1e6, 1
                    )
            finally:
                os.close(fd)

            # --- Random Write ---
            fd = os.open(tmpfile, os.O_WRONLY)
            try:
                latencies = []
                t0_total = time.perf_counter()
                for off in offsets:
                    t0 = time.perf_counter()
                    os.pwrite(fd, data_block, off)
                    latencies.append(time.perf_counter() - t0)
                elapsed = time.perf_counter() - t0_total
                if elapsed > 0:
                    result["write_iops"] = round(count / elapsed, 1)
                if latencies:
                    result["write_latency_us"] = round(
                        float(np.mean(latencies)) * 1e6, 1
                    )
            finally:
                os.close(fd)

        except Exception as exc:
            logger.warning("Manual random IO probe failed: %s", exc)
        finally:
            try:
                os.unlink(tmpfile)
            except OSError:
                pass

        logger.info(
            "Random IO (manual) -- read %.0f IOPS (%.1f us), write %.0f IOPS (%.1f us)",
            result["read_iops"], result["read_latency_us"],
            result["write_iops"], result["write_latency_us"],
        )
        return result

    def probe_queue_depth(
        self,
        path: str,
        depths: Optional[List[int]] = None,
    ) -> Dict:
        """Test throughput at different IO queue depths using threads.

        Uses a ThreadPoolExecutor with N threads doing random reads to
        simulate various queue depths.

        Args:
            path: Directory for test files.
            depths: Queue depths to test. Defaults to [1, 2, 4, 8, 16, 32].

        Returns:
            Dict with optimal_depth (int) and throughput_by_depth (dict
            mapping depth -> IOPS).
        """
        if depths is None:
            depths = [1, 2, 4, 8, 16, 32]

        os.makedirs(path, exist_ok=True)
        file_size = 256 * 1024 * 1024  # 256 MB
        block_size = 4096
        test_duration = 2.0  # seconds per depth
        max_offset = file_size - block_size
        tmpfile = os.path.join(path, ".storage_probe_qd")

        result: Dict[str, Any] = {"optimal_depth": 1, "throughput_by_depth": {}}

        try:
            # Create test file
            with open(tmpfile, "wb") as f:
                chunk = os.urandom(1024 * 1024)
                for _ in range(256):
                    f.write(chunk)
                f.flush()
                os.fsync(f.fileno())
        except Exception as exc:
            logger.warning("Could not create queue depth test file: %s", exc)
            return result

        def _worker(stop_time: float) -> int:
            """Random-read worker, returns IO count."""
            ios = 0
            try:
                fd = os.open(tmpfile, os.O_RDONLY)
                try:
                    while time.perf_counter() < stop_time:
                        off = random.randint(0, max_offset // block_size) * block_size
                        os.pread(fd, block_size, off)
                        ios += 1
                finally:
                    os.close(fd)
            except Exception:
                pass
            return ios

        best_iops = 0.0
        try:
            for depth in depths:
                try:
                    stop_time = time.perf_counter() + test_duration
                    with concurrent.futures.ThreadPoolExecutor(max_workers=depth) as pool:
                        futures = [pool.submit(_worker, stop_time) for _ in range(depth)]
                        total_ios = sum(f.result() for f in futures)
                    iops = round(total_ios / test_duration, 1)
                    result["throughput_by_depth"][depth] = iops
                    if iops > best_iops:
                        best_iops = iops
                        result["optimal_depth"] = depth
                    logger.info("QD %3d -- %.0f IOPS", depth, iops)
                except Exception as exc:
                    logger.warning("Queue depth probe failed at depth=%d: %s", depth, exc)
                    result["throughput_by_depth"][depth] = 0.0
        finally:
            try:
                os.unlink(tmpfile)
            except OSError:
                pass

        logger.info("Optimal queue depth: %d (%.0f IOPS)", result["optimal_depth"], best_iops)
        return result

    def probe_fsync(self, path: str, count: int = 1000) -> Dict:
        """Measure fsync latency by writing small blocks and fsyncing.

        Args:
            path: Directory to create test file in.
            count: Number of write+fsync cycles.

        Returns:
            Dict with mean_us, p99_us, p999_us.
        """
        os.makedirs(path, exist_ok=True)
        tmpfile = os.path.join(path, ".storage_probe_fsync")
        data = b"x" * 512
        latencies_us: List[float] = []

        try:
            fd = os.open(tmpfile, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o644)
            try:
                for _ in range(count):
                    os.write(fd, data)
                    t0 = time.perf_counter()
                    os.fsync(fd)
                    t1 = time.perf_counter()
                    latencies_us.append((t1 - t0) * 1e6)
            finally:
                os.close(fd)
        except Exception as exc:
            logger.warning("Fsync latency probe failed: %s", exc)
            return {"mean_us": 0.0, "p99_us": 0.0, "p999_us": 0.0}
        finally:
            try:
                os.unlink(tmpfile)
            except OSError:
                pass

        arr = np.array(latencies_us)
        result = {
            "mean_us": round(float(np.mean(arr)), 1),
            "p99_us": round(float(np.percentile(arr, 99)), 1),
            "p999_us": round(float(np.percentile(arr, 99.9)), 1),
        }
        logger.info(
            "Fsync latency -- mean %.0f us, p99 %.0f us, p99.9 %.0f us",
            result["mean_us"], result["p99_us"], result["p999_us"],
        )
        return result

    def run_all(self, paths: Optional[List[str]] = None) -> Dict:
        """Run all probes, optionally auto-detecting mountpoints.

        Args:
            paths: List of directory paths to probe. If None, auto-detect
                   mounted filesystems from detect_devices().

        Returns:
            Dict with devices list and per-path probe results.
        """
        devices = self.detect_devices()
        result: Dict[str, Any] = {
            "probe_type": "storage",
            "timestamp": time.time(),
            "devices": devices,
            "results_by_path": {},
        }

        if paths is None:
            # Auto-detect: pick unique mountpoints from detected devices
            seen: set = set()
            paths = []
            for dev in devices:
                mp = dev.get("mountpoint")
                if mp and mp not in seen and os.path.isdir(mp):
                    seen.add(mp)
                    paths.append(mp)
            if not paths:
                # Fallback to /tmp
                paths = ["/tmp"]
            logger.info("Auto-detected probe paths: %s", paths)

        t0_total = time.perf_counter()

        for path in paths:
            logger.info("=== Probing %s ===", path)
            path_result: Dict[str, Any] = {}

            # Sequential
            try:
                path_result["sequential"] = self.probe_sequential(path)
            except Exception as exc:
                logger.warning("Sequential probe failed on %s: %s", path, exc)
                path_result["sequential"] = {"read_mbps": 0.0, "write_mbps": 0.0}

            # Random
            try:
                path_result["random"] = self.probe_random(path)
            except Exception as exc:
                logger.warning("Random probe failed on %s: %s", path, exc)
                path_result["random"] = {
                    "read_iops": 0.0, "write_iops": 0.0,
                    "read_latency_us": 0.0, "write_latency_us": 0.0,
                }

            # Queue depth
            try:
                path_result["queue_depth"] = self.probe_queue_depth(path)
            except Exception as exc:
                logger.warning("Queue depth probe failed on %s: %s", path, exc)
                path_result["queue_depth"] = {
                    "optimal_depth": 1, "throughput_by_depth": {},
                }

            # Fsync
            try:
                path_result["fsync"] = self.probe_fsync(path)
            except Exception as exc:
                logger.warning("Fsync probe failed on %s: %s", path, exc)
                path_result["fsync"] = {"mean_us": 0.0, "p99_us": 0.0, "p999_us": 0.0}

            result["results_by_path"][path] = path_result

        result["duration_s"] = round(time.perf_counter() - t0_total, 1)
        logger.info("Storage probe suite complete in %.1f s", result["duration_s"])
        return result
