#!/usr/bin/env python3
"""
Мониторинг здоровья дисков с S.M.A.R.T (аналог CrystalDiskInfo)
Требует: sudo smartctl (smartmontools)
"""

import subprocess
import sys
from datetime import datetime

def check_smartctl():
    """Проверка наличия smartctl"""
    try:
        subprocess.run(['which', 'smartctl'], capture_output=True, check=True)
    except subprocess.CalledProcessError:
        print("❌ smartmontools не установлен!")
        print("Установите: sudo pacman -S smartmontools")
        sys.exit(1)

def get_disks():
    """Получить список дисков"""
    try:
        result = subprocess.run(['lsblk', '-d', '-n', '-o', 'NAME'],
                               capture_output=True, text=True, check=True)
        return [f"/dev/{disk}" for disk in result.stdout.strip().split('\n') if disk]
    except:
        return []

def get_disk_info(disk):
    """Получить информацию о диске"""
    try:
        model = subprocess.run(['lsblk', '-d', '-n', '-o', 'MODEL', disk],
                              capture_output=True, text=True).stdout.strip() or "Unknown"
        size = subprocess.run(['lsblk', '-d', '-n', '-o', 'SIZE', disk],
                             capture_output=True, text=True).stdout.strip() or "Unknown"
        return model, size
    except:
        return "Unknown", "Unknown"

def parse_smart_data(disk):
    """Парсить S.M.A.R.T данные"""
    try:
        result = subprocess.run(['sudo', 'smartctl', '-a', disk],
                               capture_output=True, text=True)
        
        output = result.stdout
        data = {
            'status': 'GOOD' if 'PASSED' in output else ('BAD' if 'FAILED' in output else 'UNKNOWN'),
            'temp': 'N/A',
            'power_hours': 'N/A',
            'power_cycles': 'N/A',
            'reallocated': 0,
            'pending': 0,
            'uncorrectable': 0,
            'attrs': {}
        }
        
        for line in output.split('\n'):
            parts = line.split()
            if len(parts) < 10:
                continue
            
            if '194' in parts[0] or 'Temperature_Celsius' in line:
                try:
                    data['temp'] = f"{parts[9]}°C"
                except:
                    pass
            
            if '9' in parts[0] or 'Power_On_Hours' in line:
                try:
                    hours = int(parts[9])
                    data['power_hours'] = f"{hours}h ({hours//24}d)"
                except:
                    pass
            
            if '12' in parts[0] or 'Power_Cycle_Count' in line:
                try:
                    data['power_cycles'] = parts[9]
                except:
                    pass
            
            if '5' in parts[0] or 'Reallocated_Sector_Ct' in line:
                try:
                    data['reallocated'] = int(parts[9])
                except:
                    pass
            
            if '197' in parts[0] or 'Current_Pending_Sect' in line:
                try:
                    data['pending'] = int(parts[9])
                except:
                    pass
            
            if '198' in parts[0] or 'Offline_Uncorrectable' in line:
                try:
                    data['uncorrectable'] = int(parts[9])
                except:
                    pass
            
            if parts and parts[0].isdigit() and len(parts) >= 10:
                try:
                    attr_id = parts[0]
                    attr_name = parts[1] if len(parts) > 1 else 'Unknown'
                    data['attrs'][attr_id] = {
                        'name': attr_name,
                        'value': parts[3],
                        'worst': parts[4],
                        'threshold': parts[5],
                        'raw': parts[9]
                    }
                except:
                    pass
        
        return data
    except:
        return None

def calculate_health(smart_data):
    """Правильный расчёт здоровья диска"""
    if not smart_data:
        return 50, []
    
    if smart_data['status'] == 'BAD':
        return 5, ["🔴 S.M.A.R.T статус: BAD"]
    
    health = 100
    warnings = []
    
    reallocated = smart_data['reallocated']
    pending = smart_data['pending']
    uncorrectable = smart_data['uncorrectable']
    
    # === ПЕРЕНАЗНАЧЕННЫЕ СЕКТОРА ===
    if reallocated > 0:
        if reallocated > 500:
            penalty = min(80, reallocated * 0.5)
            health -= penalty
            warnings.append(f"🔴 КРИТИЧНО: {reallocated} переназначенных секторов! Диск может отказать!")
        elif reallocated > 100:
            penalty = min(70, reallocated * 0.3)
            health -= penalty
            warnings.append(f"🟠 ВНИМАНИЕ: {reallocated} переназначенных секторов. Начните резервное копирование!")
        elif reallocated > 10:
            penalty = reallocated * 0.2
            health -= penalty
            warnings.append(f"🟡 ВНИМАНИЕ: {reallocated} переназначенных секторов")
        else:
            health -= reallocated * 0.1
    
    # === ОЖИДАЮЩИЕ СЕКТОРА ===
    if pending > 0:
        health -= min(70, pending * 2)
        warnings.append(f"🔴 КРИТИЧНО: {pending} ожидающих секторов!")
    
    # === НЕИСПРАВИМЫЕ ОШИБКИ ===
    if uncorrectable > 0:
        health -= min(80, uncorrectable * 5)
        warnings.append(f"🔴 КРИТИЧНО: {uncorrectable} неисправимых ошибок!")
    
    return max(5, int(health)), warnings

def get_color(health):
    """Цвет в зависимости от процента"""
    if health >= 85:
        return '\033[92m'  # Зеленый
    elif health >= 60:
        return '\033[93m'  # Желтый
    elif health >= 30:
        return '\033[91m'  # Красный
    else:
        return '\033[41m'  # Красный фон

def get_icon(health):
    """Иконка здоровья"""
    if health >= 85:
        return "✓"
    elif health >= 60:
        return "⚠"
    elif health >= 30:
        return "✗"
    else:
        return "💥"

def print_disk(disk, model, size, health, warnings, smart_data):
    """Вывести информацию о диске"""
    color = get_color(health)
    icon = get_icon(health)
    reset = '\033[0m'
    
    print(f"\n{'='*75}")
    print(f"Диск: {disk:15} Модель: {model:25} Размер: {size}")
    
    bar_len = 40
    filled = int(bar_len * health / 100)
    bar = '█' * filled + '░' * (bar_len - filled)
    print(f"Здоровье: {color}{icon} {bar} {health}%{reset}")
    
    if warnings:
        print("\n⚠️  ПРЕДУПРЕЖДЕНИЯ:")
        for warning in warnings:
            print(f"   {warning}")
    
    if not smart_data:
        print("S.M.A.R.T: Не поддерживается")
        return
    
    print(f"\nСтатус: {smart_data['status']:8} | Температура: {smart_data['temp']:8} | "
          f"Часов: {smart_data['power_hours']:15} | Циклов: {smart_data['power_cycles']}")
    
    print(f"\nКритические атрибуты:")
    print(f"  • Переназначенные сектора: {smart_data['reallocated']}")
    print(f"  • Ожидающие сектора: {smart_data['pending']}")
    print(f"  • Неисправимые ошибки: {smart_data['uncorrectable']}")

def main():
    """Главная функция"""
    check_smartctl()
    
    print("\n" + "="*75)
    print(f"  МОНИТОРИНГ ЗДОРОВЬЯ ДИСКОВ (S.M.A.R.T)")
    print(f"  {datetime.now().strftime('%d.%m.%Y %H:%M:%S')}")
    print("="*75)
    
    disks = get_disks()
    
    if not disks:
        print("❌ Диски не найдены")
        return
    
    for disk in disks:
        model, size = get_disk_info(disk)
        smart_data = parse_smart_data(disk)
        health, warnings = calculate_health(smart_data)
        print_disk(disk, model, size, health, warnings, smart_data)
    
    print("\n" + "="*75 + "\n")

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n\nПрограмма прервана")
    except Exception as e:
        print(f"❌ Ошибка: {e}")