"""SMART health monitoring module."""

import re
from dataclasses import dataclass
from typing import Optional, Dict, Tuple

# from p20260130_12_kilo01.shell_calls import (
from jusfltuls.shell_calls import (
    build_sudo_smartctl_x,
    build_sudo_smartctl_l_selftest,
    build_sudo_smartctl_t_short,
    build_sudo_smartctl_t_long,
    build_nvme_smart_log,
    build_nvme_self_test_log,
    build_nvme_device_self_test_status,
    build_nvme_device_self_test_short,
    build_nvme_device_self_test_long,
    build_nvme_start_self_test,
    run_command,
)
from jusfltuls.nvme_cli import (
    is_nvme_device,
    parse_nvme_smart_log,
    parse_nvme_self_test_log,
    parse_nvme_device_self_test_status,
    NvmeSelfTestLogResult,
)


@dataclass
class DiskHealth:
    """Represents the health status of a disk."""

    device: str
    overall_health: str  # "PASSED", "FAILED", or "UNKNOWN"
    power_on_hours: Optional[int]
    model_name: Optional[str]
    # Short test info
    short_test_in_progress: bool = False
    short_test_remaining_percent: Optional[int] = None
    last_short_test_hours_ago: Optional[int] = None
    last_short_test_failed: bool = False
    # Long test info
    long_test_in_progress: bool = False
    long_test_remaining_percent: Optional[int] = None
    last_long_test_hours_ago: Optional[int] = None
    last_long_test_failed: bool = False

    def __str__(self) -> str:
        health_str = f"{self.overall_health}"
        if (
            self.short_test_in_progress
            and self.short_test_remaining_percent is not None
        ):
            health_str += f" | Short: {self.short_test_remaining_percent}%"
        elif self.last_short_test_hours_ago is not None:
            suffix = "!" if self.last_short_test_failed else "h"
            health_str += f" | Short: {self.last_short_test_hours_ago}{suffix}"
        if self.long_test_in_progress and self.long_test_remaining_percent is not None:
            health_str += f" | Long: {self.long_test_remaining_percent}%"
        elif self.last_long_test_hours_ago is not None:
            suffix = "!" if self.last_long_test_failed else "h"
            health_str += f" | Long: {self.last_long_test_hours_ago}{suffix}"
        return health_str


def parse_smartctl_output(output: str) -> DiskHealth:
    """Parse smartctl -x output to extract health information.

    Args:
        output: The raw output from smartctl -x command

    Returns:
        DiskHealth object with parsed information
    """
    device = "/dev/unknown"
    overall_health = "UNKNOWN"
    power_on_hours = None
    model_name = None

    last_short_test_hours_ago = None
    last_long_test_hours_ago = None
    short_test_in_progress = False
    short_test_remaining_pct = None  # % remaining (not done)
    long_test_in_progress = False
    long_test_remaining_pct = None  # % remaining (not done)

    lines = output.split("\n")

    # Parse device from first line if present
    if lines and "/dev/" in lines[0]:
        match = re.search(r"(/dev/\w+)", lines[0])
        if match:
            device = match.group(1)

    # Parse overall health
    for line in lines:
        if "SMART overall-health self-assessment test result:" in line:
            if "PASSED" in line:
                overall_health = "PASSED"
            elif "FAILED" in line:
                overall_health = "FAILED"
            break

    # Parse model name
    for line in lines:
        if "Device Model:" in line or "Model Number:" in line:
            model_name = line.split(":", 1)[1].strip()
            break

    # Parse power on hours - handle both ATA and NVMe formats
    for line in lines:
        if "Power_On_Hours" in line and not line.startswith("#"):
            # ATA attribute table line, e.g.:
            #   9 Power_On_Hours  -O--CK  091  091  000  -  44834
            # The raw value is always the last whitespace-separated token.
            parts = line.split()
            if parts:
                try:
                    power_on_hours = int(parts[-1].replace(",", "").replace(".", ""))
                except (ValueError, IndexError):
                    pass
            break
        elif "Power On Hours:" in line:
            # NVMe / newer smartctl key-value line, e.g.:
            #   Power On Hours:  6,474    or    Power On Hours:  6.474
            # Strip both comma and period thousand-separators.
            match = re.search(r"Power On Hours:\s+([\d,. ]+)", line)
            if match:
                try:
                    raw = match.group(1).strip()
                    power_on_hours = int(
                        raw.replace(",", "").replace(".", "").replace(" ", "")
                    )
                except (ValueError, IndexError):
                    pass
            break

    # Detect in-progress test.  Three output formats are possible depending on
    # the smartctl version / device type:
    #
    # Format A — SMART Attribute table (classic ATA, older smartctl):
    #   1  0  0  Self_test_in_progress [90% left] (119254528-119320063)
    #   The test type is NOT available here; it is resolved in get_disk_health()
    #   via the -l selftest log.  We use the short slot as a temporary sentinel.
    #
    # Format B — Self-test status section (newer smartctl / some NVMe/SATA):
    #   Self-test status: Short self-test in progress (58% completed)
    #   Self-test status: Extended self-test in progress (42% completed)
    #   The test type IS available here; we set the correct flag directly.
    #
    # Format C — Self-test execution status (ATA, this smartctl version):
    #   Self-test execution status:      ( 249)	Self-test routine in progress...
    #                                   90% of test remaining.
    #   The percentage is on the NEXT line.  Type resolved via -l selftest.
    for i, line in enumerate(lines):
        # Format B — "Self-test status:" line
        if (
            line.strip().startswith("Self-test status:")
            and "in progress" in line.lower()
            and "no self-test" not in line.lower()
        ):
            # Extract % completed (convert to % remaining for consistency)
            m_pct = re.search(r"(\d+)%\s+completed", line, re.IGNORECASE)
            completed = int(m_pct.group(1)) if m_pct else None
            remaining = (100 - completed) if completed is not None else None
            # Determine test type from the line text
            line_lower = line.lower()
            if "short" in line_lower:
                short_test_in_progress = True
                short_test_remaining_pct = remaining
            elif "extended" in line_lower or "long" in line_lower:
                long_test_in_progress = True
                long_test_remaining_pct = remaining
            else:
                # Unknown type — fall back to short sentinel so get_disk_health
                # can resolve it via the -l selftest log
                short_test_in_progress = True
                short_test_remaining_pct = remaining
            break

        # Format C — "Self-test execution status:" two-line form
        if "Self-test execution status:" in line and "in progress" in line.lower():
            # Percentage is on the next line: "  90% of test remaining."
            remaining = None
            if i + 1 < len(lines):
                m_pct = re.search(
                    r"(\d+)%\s+of\s+test\s+remaining", lines[i + 1], re.IGNORECASE
                )
                if m_pct:
                    remaining = int(m_pct.group(1))
            # Type not stated here — use short sentinel; resolved via -l selftest
            short_test_in_progress = True
            short_test_remaining_pct = remaining
            break

        # Format A — SMART Attribute table
        if "Self_test_in_progress" in line:
            m = re.search(r"\[(\d+)%\s+left\]", line, re.IGNORECASE)
            if m:
                remaining = int(m.group(1))
                # Mark as "some test in progress"; type resolved later.
                # Use a sentinel: store remaining in short slot temporarily.
                short_test_in_progress = (
                    True  # placeholder — type fixed in get_disk_health
                )
                short_test_remaining_pct = remaining
            break

    return DiskHealth(
        device=device,
        overall_health=overall_health,
        power_on_hours=power_on_hours,
        model_name=model_name,
        short_test_in_progress=short_test_in_progress,
        short_test_remaining_percent=None
        if short_test_remaining_pct is None
        else (100 - short_test_remaining_pct),
        last_short_test_hours_ago=last_short_test_hours_ago,
        long_test_in_progress=long_test_in_progress,
        long_test_remaining_percent=None
        if long_test_remaining_pct is None
        else (100 - long_test_remaining_pct),
        last_long_test_hours_ago=last_long_test_hours_ago,
    )


@dataclass
class SelftestLogResult:
    """Parsed result from `smartctl -l selftest`."""

    last_short_hours_ago: Optional[int] = None
    last_short_failed: bool = False
    last_long_hours_ago: Optional[int] = None
    last_long_failed: bool = False
    short_in_progress: bool = False
    short_remaining_pct: Optional[int] = None  # % remaining (not done)
    long_in_progress: bool = False
    long_remaining_pct: Optional[int] = None  # % remaining (not done)


def parse_selftest_log(output: str, power_on_hours: Optional[int]) -> SelftestLogResult:
    """Parse `smartctl -l selftest` output.

    Handles two log formats:

    Format A — classic ATA (has leading '#', subtype word, and a '%' column):
      # 1  Short offline       Completed without error        00%  44832  -
      # 1  Extended offline    Self-test routine in progress  90%  44900  -

    Format B — NVMe columnar (no '#', no subtype, no '%' column):
      Num  Test_Description    Status                  Power_on_Hours  ...
        0  Short               Completed without error           6474  -  -  -  -

    Only "Completed without error" entries count toward last_*_hours_ago.
    In-progress entries set the *_in_progress / *_remaining_pct fields.

    Args:
        output: Raw output from `smartctl -l selftest`
        power_on_hours: Current power-on hours (used to compute hours-ago)

    Returns:
        SelftestLogResult with all fields populated.
    """
    # Format A regex — requires '#', a subtype token, a '%' column, and lifetime hours.
    ata_re = re.compile(
        r"#\s*\d+\s+"
        r"(Short|Long|Extended)\s+\S+"  # group 1: test type
        r"\s+(.+?)"  # group 2: status text (non-greedy)
        r"\s+(\d+)%\s+"  # group 3: remaining %
        r"(\d+)",  # group 4: lifetime hours
        re.IGNORECASE,
    )

    # Detect Format B by looking for the NVMe header line and recording the
    # column offset of Power_on_Hours so we can slice data lines reliably.
    poh_col = None
    for line in output.split("\n"):
        if re.search(r"Power_on_Hours", line, re.IGNORECASE):
            poh_col = line.index("Power_on_Hours")
            break

    result = SelftestLogResult()
    short_done = False
    long_done = False

    for line in output.split("\n"):
        # ── Format A ──────────────────────────────────────────────────────────
        m = ata_re.search(line)
        if m:
            test_type = m.group(1).lower()  # "short", "long", "extended"
            status = m.group(2).strip().lower()
            remaining = int(m.group(3))  # % remaining
            lifetime = int(m.group(4))
            is_long = test_type in ("long", "extended")

            if "in progress" in status:
                if test_type == "short":
                    result.short_in_progress = True
                    result.short_remaining_pct = remaining
                elif is_long:
                    result.long_in_progress = True
                    result.long_remaining_pct = remaining
                continue

            hours_ago = (
                (power_on_hours - lifetime) if power_on_hours is not None else None
            )
            failed = status != "completed without error"
            if test_type == "short" and not short_done:
                result.last_short_hours_ago = hours_ago
                result.last_short_failed = failed
                short_done = True
            elif is_long and not long_done:
                result.last_long_hours_ago = hours_ago
                result.last_long_failed = failed
                long_done = True

            if (
                short_done
                and long_done
                and not result.short_in_progress
                and not result.long_in_progress
            ):
                break
            continue

        # ── Format B (NVMe) ───────────────────────────────────────────────────
        # Only attempt if we found the Power_on_Hours column header.
        if poh_col is None:
            continue

        # Data lines start with an optional whitespace then a decimal index.
        # They must NOT start with 'Num' (header) or be empty.
        stripped = line.strip()
        if not stripped or not re.match(r"\d", stripped):
            continue

        # Extract lifetime hours from the Power_on_Hours column position.
        if len(line) <= poh_col:
            continue
        poh_substr = line[poh_col : poh_col + 20]
        poh_match = re.search(r"(\d[\d,. ]*)", poh_substr)
        if not poh_match:
            continue
        try:
            lifetime = int(
                poh_match.group(1).replace(",", "").replace(".", "").replace(" ", "")
            )
        except ValueError:
            continue

        # Extract test type and status from the beginning of the line.
        # Typical columns before Power_on_Hours: Num  Test_Description  Status
        prefix = line[:poh_col].strip()
        # prefix example: "0  Short               Completed without error"
        # Split off the leading index number then grab type and status.
        prefix_m = re.match(
            r"\d+\s+(Short|Long|Extended)\s+(.*)", prefix, re.IGNORECASE
        )
        if not prefix_m:
            continue
        test_type = prefix_m.group(1).lower()
        status = prefix_m.group(2).strip().lower()
        is_long = test_type in ("long", "extended")

        if "in progress" in status:
            if test_type == "short":
                result.short_in_progress = True
                result.short_remaining_pct = None  # NVMe format has no % column
            elif is_long:
                result.long_in_progress = True
                result.long_remaining_pct = None
            continue

        if "in progress" not in status:
            hours_ago = (
                (power_on_hours - lifetime) if power_on_hours is not None else None
            )
            failed = "completed without error" not in status
            if test_type == "short" and not short_done:
                result.last_short_hours_ago = hours_ago
                result.last_short_failed = failed
                short_done = True
            elif is_long and not long_done:
                result.last_long_hours_ago = hours_ago
                result.last_long_failed = failed
                long_done = True

        if (
            short_done
            and long_done
            and not result.short_in_progress
            and not result.long_in_progress
        ):
            break

    return result


def parse_selftest_output(
    output: str, current_power_on_hours: Optional[int]
) -> Tuple[str, Optional[int]]:
    """Parse smartctl -l selftest output to extract last test information.

    Args:
        output: The raw output from smartctl -l selftest command
        current_power_on_hours: Current power-on hours for calculating hours ago

    Returns:
        Tuple of (test_status, hours_ago)
    """
    lines = output.split("\n")
    last_test_status = "No tests"
    last_test_hours_ago = None

    # Find the header line to locate Power_on_Hours column position
    power_on_hours_col_start = None
    status_col_start = None
    for line in lines:
        if "Power_on_Hours" in line or "LifeTime(hours)" in line:
            # Found header line, get position of the column
            power_on_hours_col_start = line.find("Power_on_Hours")
            if power_on_hours_col_start == -1:
                power_on_hours_col_start = line.find("LifeTime(hours)")
            # Also find Status column position for NVMe
            status_col_start = line.find("Status")
            break

    for line in lines:
        # Skip header lines and empty lines
        if (
            not line.strip()
            or line.startswith("Num")
            or line.startswith("SMART")
            or line.startswith("Self-test")
        ):
            continue

        # Check if this is a test entry line (starts with a number)
        if not re.match(r"\s*#?\s*\d+", line):
            continue

        # For NVMe: extract Power_on_Hours from the column position
        if power_on_hours_col_start is not None and power_on_hours_col_start > 0:
            # Extract substring from Power_on_Hours column position
            hours_substring = line[
                power_on_hours_col_start : power_on_hours_col_start + 15
            ]
            # Find first number in this substring
            hours_match = re.search(r"(\d+)", hours_substring)
            if hours_match:
                try:
                    lifetime_hours = int(hours_match.group(1))
                    # Extract test status from Status column
                    if status_col_start is not None and status_col_start > 0:
                        status_substring = line[
                            status_col_start:power_on_hours_col_start
                        ]
                        # Clean up the status string
                        last_test_status = " ".join(status_substring.split())
                    else:
                        last_test_status = "Unknown"

                    # Calculate hours ago if we have current power on hours
                    if current_power_on_hours is not None:
                        last_test_hours_ago = current_power_on_hours - lifetime_hours
                    break
                except (ValueError, IndexError):
                    pass
        else:
            # ATA format: # 1  Short offline       Completed without error       00%     4454         -
            match = re.match(
                r"\s*#?\s*(\d+)\s+(\S+)\s+(\S+)\s+(\S.*?)\s+(\d+)%?\s+(\d+)", line
            )
            if match:
                test_status = match.group(3).strip()
                lifetime_hours = int(match.group(6))

                last_test_status = test_status

                # Calculate hours ago if we have current power on hours
                if current_power_on_hours is not None:
                    last_test_hours_ago = current_power_on_hours - lifetime_hours
                break

    return last_test_status, last_test_hours_ago


class SmartHealthMonitor:
    """Main class for monitoring disk health via SMART."""

    def get_disk_health(self, device: str) -> DiskHealth:
        """Get health information for a disk using smartctl -x and -l selftest.

        For NVMe devices, prefers the nvme-cli utility which often provides
        better support than smartctl for certain NVMe disks.

        Args:
            device: Device path (e.g., /dev/sda)

        Returns:
            DiskHealth object with health status
        """
        # For NVMe devices, prefer nvme-cli utility
        if is_nvme_device(device):
            return self._get_nvme_disk_health(device)

        # For non-NVMe devices, use smartctl
        return self._get_smartctl_disk_health(device)

    def _get_smartctl_disk_health(self, device: str) -> DiskHealth:
        """Get health information using smartctl (for SATA/ATA disks)."""
        # Get basic health info from smartctl -x
        cmd = build_sudo_smartctl_x(device)
        returncode, output = run_command(cmd)

        # Parse output even if return code is non-zero
        # smartctl may return non-zero for USB devices but still provide valid data
        health = parse_smartctl_output(output)

        # If we successfully parsed a PASSED/FAILED status, use it
        # Otherwise, mark as ERROR only if we couldn't parse anything useful
        if health.overall_health == "UNKNOWN" and returncode != 0:
            health.overall_health = "ERROR"

        # Set the device path from the input
        health.device = device

        # Determine test type and history via `smartctl -l selftest`.
        selftest_cmd = build_sudo_smartctl_l_selftest(device)
        _, selftest_output = run_command(selftest_cmd)
        st = parse_selftest_log(selftest_output, health.power_on_hours)

        if health.short_test_in_progress or health.long_test_in_progress:
            if health.short_test_in_progress and not health.long_test_in_progress:
                # Format A sentinel path: type stored in short slot temporarily.
                # Use the `-l selftest` first entry to resolve the actual type.
                remaining_pct_done = health.short_test_remaining_percent
                if st.long_in_progress:
                    # The running test is a long/extended one
                    health.short_test_in_progress = False
                    health.short_test_remaining_percent = None
                    health.long_test_in_progress = True
                    health.long_test_remaining_percent = remaining_pct_done
                # else: it really is a short test — keep as-is
            # else: type was already determined by Format B — nothing to fix up

            # Do NOT fill in history hours while a test is running
            health.last_short_test_hours_ago = None
            health.last_long_test_hours_ago = None
        else:
            # No test in progress — populate history from -l selftest log
            health.last_short_test_hours_ago = st.last_short_hours_ago
            health.last_short_test_failed = st.last_short_failed
            health.last_long_test_hours_ago = st.last_long_hours_ago
            health.last_long_test_failed = st.last_long_failed

        return health

    def _get_nvme_disk_health(self, device: str) -> DiskHealth:
        """Get health information using nvme-cli (for NVMe disks).

        Uses the following NVMe commands:
        1. `nvme smart-log` - to extract critical warnings and power-on hours
        2. `nvme device-self-test` (no params) - to check current test progress
        3. `nvme self-test-log` - to get historical test results
        """
        # Get SMART log via nvme-cli to extract errors and critical warnings
        cmd = build_nvme_smart_log(device)
        returncode, output = run_command(cmd)

        if returncode != 0:
            # nvme-cli failed, return N/A health status
            return DiskHealth(
                device=device,
                overall_health="N/A",
                power_on_hours=None,
                model_name=None,
            )

        # Parse nvme smart-log output - now returns critical_warning too
        overall_health, power_on_hours, critical_warning = parse_nvme_smart_log(output)

        # Get self-test log first - it contains both current and historical test info
        # This tells us when the last test completed and if it passed/failed
        # It also shows if a test is currently in progress
        selftest_cmd = build_nvme_self_test_log(device)
        _, selftest_output = run_command(selftest_cmd)
        st = parse_nvme_self_test_log(selftest_output, power_on_hours)

        # Get current device self-test status using device-self-test -s 0h
        # This gives us the current progress percentage if a test is running
        status_cmd = build_nvme_device_self_test_status(device)
        returncode_status, status_output = run_command(status_cmd)

        # Debug output
        from jusfltuls.shell_calls import get_debug_mode
        import sys
        if get_debug_mode():
            print(f"\n[DEBUG] NVMe self-test status command returned: {returncode_status}", file=sys.stderr, flush=True)
            print(f"[DEBUG] Output: {repr(status_output)}", file=sys.stderr, flush=True)

        current_status = parse_nvme_device_self_test_status(status_output)

        # Merge current status with self-test-log data
        # Strategy: Use self-test-log to determine test TYPE (short vs long)
        #          Use device-self-test to get the most up-to-date progress PERCENTAGE

        # Check if device-self-test detected any test in progress
        device_test_in_progress = (current_status.short_in_progress or
                                   current_status.long_in_progress)
        device_test_remaining_pct = (current_status.short_remaining_pct or
                                     current_status.long_remaining_pct)

        # Determine test type from self-test-log and combine with device-self-test progress
        if st.short_in_progress:
            # Self-test-log says short test is running
            short_in_progress = True
            # Use device-self-test progress if available, else use self-test-log progress
            short_remaining_pct = device_test_remaining_pct if device_test_in_progress else st.short_remaining_pct
            long_in_progress = False
            long_remaining_pct = None
        elif st.long_in_progress:
            # Self-test-log says long test is running
            short_in_progress = False
            short_remaining_pct = None
            long_in_progress = True
            # Use device-self-test progress if available, else use self-test-log progress
            long_remaining_pct = device_test_remaining_pct if device_test_in_progress else st.long_remaining_pct
        elif device_test_in_progress:
            # device-self-test shows progress but self-test-log doesn't show in-progress
            # This can happen with "progress 76%" output
            # Try to determine test type from most recent test in history
            if st.last_short_hours_ago is not None and st.last_long_hours_ago is not None:
                # Compare which test was run more recently
                if st.last_short_hours_ago < st.last_long_hours_ago:
                    # Short test was more recent, assume it's a short test
                    short_in_progress = True
                    short_remaining_pct = device_test_remaining_pct
                    long_in_progress = False
                    long_remaining_pct = None
                else:
                    # Long test was more recent, assume it's a long test
                    short_in_progress = False
                    short_remaining_pct = None
                    long_in_progress = True
                    long_remaining_pct = device_test_remaining_pct
            elif st.last_short_hours_ago is not None:
                # Only short test history, assume it's a short test
                short_in_progress = True
                short_remaining_pct = device_test_remaining_pct
                long_in_progress = False
                long_remaining_pct = None
            elif st.last_long_hours_ago is not None:
                # Only long test history, assume it's a long test
                short_in_progress = False
                short_remaining_pct = None
                long_in_progress = True
                long_remaining_pct = device_test_remaining_pct
            else:
                # No history, default to short test (most common)
                short_in_progress = True
                short_remaining_pct = device_test_remaining_pct
                long_in_progress = False
                long_remaining_pct = None
        else:
            # No test in progress
            short_in_progress = False
            short_remaining_pct = None
            long_in_progress = False
            long_remaining_pct = None

        # Use historical data from self-test-log for completed tests
        last_short_hours_ago = st.last_short_hours_ago if not short_in_progress else None
        last_short_failed = st.last_short_failed
        last_long_hours_ago = st.last_long_hours_ago if not long_in_progress else None
        last_long_failed = st.last_long_failed

        # Debug output
        if get_debug_mode():
            print(f"[DEBUG] Final merged results:", file=sys.stderr, flush=True)
            print(f"[DEBUG]   short_in_progress={short_in_progress}, short_remaining_pct={short_remaining_pct}", file=sys.stderr, flush=True)
            print(f"[DEBUG]   long_in_progress={long_in_progress}, long_remaining_pct={long_remaining_pct}", file=sys.stderr, flush=True)
            print(f"[DEBUG]   last_short_hours_ago={last_short_hours_ago}, last_short_failed={last_short_failed}", file=sys.stderr, flush=True)
            print(f"[DEBUG]   last_long_hours_ago={last_long_hours_ago}, last_long_failed={last_long_failed}", file=sys.stderr, flush=True)

        # Build DiskHealth from nvme-cli results
        health = DiskHealth(
            device=device,
            overall_health=overall_health,
            power_on_hours=power_on_hours,
            model_name=None,  # nvme smart-log doesn't provide model name
            short_test_in_progress=short_in_progress,
            short_test_remaining_percent=short_remaining_pct,
            last_short_test_hours_ago=last_short_hours_ago,
            last_short_test_failed=last_short_failed,
            long_test_in_progress=long_in_progress,
            long_test_remaining_percent=long_remaining_pct,
            last_long_test_hours_ago=last_long_hours_ago,
            last_long_test_failed=last_long_failed,
        )

        return health

    def get_all_disks_health(self, devices: list) -> Dict[str, DiskHealth]:
        """Get health information for multiple disks.

        Args:
            devices: List of device paths

        Returns:
            Dictionary mapping device paths to DiskHealth objects
        """
        result = {}
        for device in devices:
            result[device] = self.get_disk_health(device)
        return result

    def start_short_test(self, device: str) -> Tuple[bool, str]:
        """Start a short self-test on the given device.

        For NVMe devices, uses `nvme device-self-test -s 1h` command
        which provides better support than smartctl for certain NVMe disks.

        Args:
            device: Device path (e.g., /dev/sda)

        Returns:
            Tuple of (success, message)
        """
        # For NVMe devices, use nvme device-self-test -s 1h
        if is_nvme_device(device):
            cmd = build_nvme_device_self_test_short(device)
            returncode, output = run_command(cmd)

            if returncode == 0:
                return True, "Short test started"
            else:
                return False, f"Failed to start test: {output}"

        # For non-NVMe devices, use smartctl
        cmd = build_sudo_smartctl_t_short(device)
        returncode, output = run_command(cmd)

        if returncode == 0:
            return True, "Short test started"
        else:
            return False, f"Failed to start test: {output}"

    def start_long_test(self, device: str) -> Tuple[bool, str]:
        """Start a long self-test on the given device.

        For NVMe devices, uses `nvme device-self-test -s 2h` command
        which provides better support than smartctl for certain NVMe disks.

        Args:
            device: Device path (e.g., /dev/sda)

        Returns:
            Tuple of (success, message)
        """
        # For NVMe devices, use nvme device-self-test -s 2h
        if is_nvme_device(device):
            cmd = build_nvme_device_self_test_long(device)
            returncode, output = run_command(cmd)

            if returncode == 0:
                return True, "Long test started"
            else:
                return False, f"Failed to start test: {output}"

        # For non-NVMe devices, use smartctl
        cmd = build_sudo_smartctl_t_long(device)
        returncode, output = run_command(cmd)

        if returncode == 0:
            return True, "Long test started"
        else:
            return False, f"Failed to start test: {output}"
