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


@dataclass
class DiskHealth:
    """Represents the health status of a disk."""
    device: str
    overall_health: str  # "PASSED", "FAILED", or "UNKNOWN"
    last_test_status: str
    last_test_hours_ago: Optional[int]  # Hours since last test
    power_on_hours: Optional[int]
    model_name: Optional[str]
    test_in_progress: bool = False
    test_remaining_percent: Optional[int] = None  # % remaining if test in progress

    def __str__(self) -> str:
        health_str = f"{self.overall_health}"
        if self.test_in_progress and self.test_remaining_percent is not None:
            health_str += f" | Test: {self.test_remaining_percent}% done"
        elif self.last_test_hours_ago is not None:
            health_str += f" | Test: {self.last_test_hours_ago}h ago"
        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"
    last_test_status = "No tests"
    last_test_hours_ago = None
    power_on_hours = None
    model_name = None
    test_in_progress = False
    test_remaining_percent = None

    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
    # ATA format: ID# ATTRIBUTE_NAME FLAGS VALUE WORST THRESH FAIL RAW_VALUE
    # Example: 9 Power_On_Hours -O--CK 091 091 000 - 44257
    # NVMe format: Power On Hours:                     6,576
    for line in lines:
        if 'Power_On_Hours' in line and not line.startswith('#'):
            parts = line.split()
            # Line format: 9 Power_On_Hours -O--CK 091 091 000 - 44257
            # Index:       0 1              2       3   4   5   6 7
            # RAW_VALUE is at index 7 (8th element)
            if len(parts) >= 8:
                try:
                    power_on_hours = int(parts[7])
                except (ValueError, IndexError):
                    pass
            break
        # NVMe format: "Power On Hours:                     6,576"
        elif 'Power On Hours:' in line:
            # Extract the number after the colon, handling comma separators
            match = re.search(r'Power On Hours:\s+([\d,]+)', line)
            if match:
                try:
                    # Remove commas and convert to int
                    power_on_hours = int(match.group(1).replace(',', ''))
                except (ValueError, IndexError):
                    pass
            break

    # Parse self-test execution status for in-progress tests
    for i, line in enumerate(lines):
        if 'Self-test execution status:' in line:
            # Check for in-progress status
            match = re.search(r'\(\s*(\d+)\s*\)', line)
            if match:
                status_code = int(match.group(1))
                # Status 0 = completed/no test, 240-249 = in progress
                if status_code >= 240 and status_code <= 249:
                    test_in_progress = True

            # Look for remaining percentage in this line or next few lines
            # Show percentage DONE (100 - remaining)
            for j in range(i, min(i + 3, len(lines))):
                remaining_match = re.search(r'(\d+)%\s+of\s+test\s+remaining', lines[j], re.IGNORECASE)
                if remaining_match:
                    remaining_percent = int(remaining_match.group(1))
                    test_remaining_percent = 100 - remaining_percent  # Convert to % done
                    break
            break

    # Also check for "Self-test status:" line (NVMe format)
    # Example: "Self-test status: Short self-test in progress (32% completed)"
    for line in lines:
        if 'Self-test status:' in line and 'in progress' in line.lower():
            test_in_progress = True
            # Extract percentage completed (already shows % done)
            completed_match = re.search(r'(\d+)%\s+completed', line, re.IGNORECASE)
            if completed_match:
                test_remaining_percent = int(completed_match.group(1))  # Already % done
            break

    # Parse self-test log for last completed test (only if no test in progress)
    if not test_in_progress:
        in_selftest_log = False
        for line in lines:
            if 'SMART Self-test log structure' in line:
                in_selftest_log = True
                continue

            if in_selftest_log:
                # Look for test entries like:
                # # 1  Short offline       Completed without error       00%     42191         -
                match = re.match(r'\s*#\s*1\s+(\S+)\s+(\S+)\s+(\S.*?)\s+(\d+)%\s+(\d+)', line)
                if match:
                    test_type = match.group(1)
                    test_status = match.group(3).strip()
                    lifetime_hours = int(match.group(5))

                    last_test_status = test_status

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

    return DiskHealth(
        device=device,
        overall_health=overall_health,
        last_test_status=last_test_status,
        last_test_hours_ago=last_test_hours_ago,
        power_on_hours=power_on_hours,
        model_name=model_name,
        test_in_progress=test_in_progress,
        test_remaining_percent=test_remaining_percent
    )


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

        # Get self-test information from smartctl -l selftest
        selftest_cmd = build_sudo_smartctl_l_selftest(device)
        selftest_returncode, selftest_output = run_command(selftest_cmd)

        # Get self-test information from smartctl -l selftest
        # Parse output even if return code is non-zero (some devices return warnings)
        test_status, hours_ago = parse_selftest_output(
            selftest_output,
            health.power_on_hours
        )
        # Update if we found test information
        if test_status != "No tests":
            health.last_test_status = test_status
            health.last_test_hours_ago = 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}"
