"""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
    
    # Short test tracking
    short_test_in_progress = False
    short_test_remaining_percent = None
    last_short_test_hours_ago = None
    
    # Long test tracking
    long_test_in_progress = False
    long_test_remaining_percent = None
    last_long_test_hours_ago = 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
    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

    # Parse self-test execution status for in-progress tests
    for i, line in enumerate(lines):
        if 'Self-test execution status:' in line:
            match = re.search(r'\(\s*(\d+)\s*\)', line)
            if match:
                status_code = int(match.group(1))
                if status_code >= 240 and status_code <= 249:
                    # Determine if short or long from context
                    for j in range(i, min(i + 3, len(lines))):
                        if 'short' in lines[j].lower():
                            short_test_in_progress = True
                            break
                        elif 'long' in lines[j].lower() or 'extended' in lines[j].lower():
                            long_test_in_progress = True
                            break
                    else:
                        # Default to short if we can't determine
                        short_test_in_progress = True

            # Look for remaining percentage
            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))
                    percent_done = 100 - remaining_percent
                    if short_test_in_progress:
                        short_test_remaining_percent = percent_done
                    elif long_test_in_progress:
                        long_test_remaining_percent = percent_done
                    break
            break

    # Check for "Self-test status:" line (NVMe format)
    for line in lines:
        if 'Self-test status:' in line and 'in progress' in line.lower():
            completed_match = re.search(r'(\d+)%\s+completed', line, re.IGNORECASE)
            if 'short' in line.lower():
                short_test_in_progress = True
                if completed_match:
                    short_test_remaining_percent = int(completed_match.group(1))
            elif 'long' in line.lower() or 'extended' in line.lower():
                long_test_in_progress = True
                if completed_match:
                    long_test_remaining_percent = int(completed_match.group(1))
            break

    # Parse self-test log for last completed tests
    in_selftest_log = False
    short_test_found = False
    long_test_found = False
    
    for line in lines:
        if 'SMART Self-test log structure' in line or 'SMART Extended Comprehensive Error Log' in line:
            in_selftest_log = True
            continue

        if in_selftest_log:
            # Look for test entries like:
            # # 1  Short offline       Completed without error       00%     42191         -
            # # 2  Extended offline    Completed without error       00%     42191         -
            # The format is: # NUM  TestDescription  Status  Percent  LifeTimeHours  LBA
            # Split by multiple spaces to handle variable spacing
            parts = [p.strip() for p in line.split('  ') if p.strip()]
            
            if len(parts) >= 5 and parts[0].startswith('#'):
                # Found a test entry line
                test_desc = parts[1].lower()  # e.g., "Short offline" or "Extended offline"
                
                # Find the lifetime hours (usually second-to-last field before LBA)
                try:
                    # Look for a numeric field that's likely the hours
                    lifetime_hours = None
                    for part in reversed(parts[:-1]):  # Skip last field (LBA)
                        if part.isdigit():
                            lifetime_hours = int(part)
                            break
                    
                    if lifetime_hours is not None and power_on_hours is not None:
                        hours_ago = power_on_hours - lifetime_hours
                    else:
                        hours_ago = None

                    if 'short' in test_desc and not short_test_found:
                        last_short_test_hours_ago = hours_ago
                        short_test_found = True
                    elif ('long' in test_desc or 'extended' in test_desc) and not long_test_found:
                        last_long_test_hours_ago = hours_ago
                        long_test_found = True
                    
                    # Stop if we found both
                    if short_test_found and long_test_found:
                        break
                except (ValueError, IndexError):
                    pass
            elif len(parts) >= 4 and parts[0].startswith('#'):
                # Alternative parsing - try regex for different formats
                # Format: # 1 Short offline Completed without error 00% 44257 -
                match = re.match(r'#\s*\d+\s+((?:Short|Long|Extended)(?:\s+\w+)?)', line, re.IGNORECASE)
                if match:
                    test_desc = match.group(1).lower()
                    # Try to find the hours field (numeric after %)
                    hours_match = re.search(r'(\d+)%?\s+(\d+)\s*[-]?\s*$', line)
                    if hours_match:
                        try:
                            lifetime_hours = int(hours_match.group(2))
                            if power_on_hours is not None:
                                hours_ago = power_on_hours - lifetime_hours
                            else:
                                hours_ago = None

                            if 'short' in test_desc and not short_test_found:
                                last_short_test_hours_ago = hours_ago
                                short_test_found = True
                            elif ('long' in test_desc or 'extended' in test_desc) and not long_test_found:
                                last_long_test_hours_ago = hours_ago
                                long_test_found = True
                            
                            if short_test_found and long_test_found:
                                break
                        except ValueError:
                            pass

    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=short_test_remaining_percent,
        last_short_test_hours_ago=last_short_test_hours_ago,
        long_test_in_progress=long_test_in_progress,
        long_test_remaining_percent=long_test_remaining_percent,
        last_long_test_hours_ago=last_long_test_hours_ago
    )


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

        # Self-test information is already parsed by parse_smartctl_output
        # from the -x output which includes the self-test log
        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}"
