"""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,
    run_command,
)


@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
    # Long test info
    long_test_in_progress: bool = False
    long_test_remaining_percent: Optional[int] = None
    last_long_test_hours_ago: Optional[int] = None

    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:
            health_str += f" | Short: {self.last_short_test_hours_ago}h"
        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:
            health_str += f" | Long: {self.last_long_test_hours_ago}h"
        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('#'):
            parts = line.split()
            if len(parts) >= 8:
                try:
                    power_on_hours = int(parts[7])
                except (ValueError, IndexError):
                    pass
            break
        elif 'Power On Hours:' in line:
            match = re.search(r'Power On Hours:\s+([\d,]+)', line)
            if match:
                try:
                    power_on_hours = int(match.group(1).replace(',', ''))
                except (ValueError, IndexError):
                    pass
            break

    # Detect in-progress test from the SMART Attribute line:
    #   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 (first entry shows the running test type).
    for line in lines:
        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_long_hours_ago:    Optional[int] = None
    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 both completed entries and in-progress entries:
      # 1  Extended offline    Self-test routine in progress  90%  44900  -
      # 1  Short offline       Completed without error        00%  44832  -
      # 2  Extended offline    Aborted by host                90%  34922  -

    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.
    """
    entry_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,
    )

    result = SelftestLogResult()
    short_done = False
    long_done  = False

    for line in output.split('\n'):
        m = entry_re.search(line)
        if not m:
            continue

        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
            # in-progress entry is always #1 — no point reading further for this type
            continue

        if status == 'completed without error':
            hours_ago = (power_on_hours - lifetime) if power_on_hours is not None else None
            if test_type == 'short' and not short_done:
                result.last_short_hours_ago = hours_ago
                short_done = True
            elif is_long and not long_done:
                result.last_long_hours_ago = hours_ago
                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.

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

        Returns:
            DiskHealth object with health status
        """
        # 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:
            # `-x` detected a test in progress (Self_test_in_progress).
            # Use the `-l selftest` first entry to determine the actual type,
            # then move the remaining% to the right field.
            remaining_pct_done = health.short_test_remaining_percent  # stored temporarily in short slot
            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

            # 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_long_test_hours_ago  = st.last_long_hours_ago

        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.

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

        Returns:
            Tuple of (success, message)
        """
        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.

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

        Returns:
            Tuple of (success, message)
        """
        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}"
