#!/usr/bin/env python3
"""
DiskInfo class to retrieve comprehensive disk information.
Supports physical disks, SSDs, and virtual disks (VMware, VirtualBox, KVM).
No third-party libraries required, uses /sys, /proc, and system commands.
"""

import os
import re
import glob
import subprocess
from typing import Optional, Dict, List, Tuple, Any


class DiskInfo:
    """Class to retrieve detailed information about disks."""
    
    def __init__(self):
        self.disk_name: Optional[str] = None
        self._info_cache: Dict[str, Any] = {}
    
    def get_all_disks(self) -> List[str]:
        """
        Get list of all disk devices (excluding partitions).
        Returns disk names like ['sda', 'sdb', 'vda', 'nvme0n1'].
        """
        disks = []
        
        # Method 1: /sys/block/
        try:
            for path in glob.glob('/sys/block/*'):
                disk = os.path.basename(path)
                # Skip loop devices, ram disks, and removable drives if needed
                if not disk.startswith(('loop', 'ram', 'sr')):
                    # Check if it's a real disk (has device directory)
                    if os.path.exists(f'/sys/block/{disk}/device'):
                        disks.append(disk)
        except Exception:
            pass
        
        # Method 2: /proc/partitions (for additional disks)
        try:
            with open('/proc/partitions', 'r') as f:
                for line in f:
                    parts = line.strip().split()
                    if len(parts) >= 4:
                        device = parts[3]
                        # Major number 8 = SCSI/SATA, 9 = MD, 11 = SCSI CD-ROM, 259 = NVMe
                        major = int(parts[1]) if parts[1].isdigit() else 0
                        if (major in (8, 259) or device.startswith(('sd', 'hd', 'vd', 'nvme'))) and \
                           device not in disks and not device[-1].isdigit():
                            disks.append(device)
        except Exception:
            pass
        
        return sorted(set(disks))
    
    def set_disk(self, disk_name: str) -> 'DiskInfo':
        """Set specific disk to query."""
        self.disk_name = disk_name
        self._info_cache = {}
        return self
    
    def get_disk_info(self, disk_name: Optional[str] = None) -> Optional[Dict[str, Any]]:
        """
        Get comprehensive information for a disk.
        If disk_name is None, uses the currently set disk.
        Returns dictionary with all available information.
        """
        target_disk = disk_name or self.disk_name
        if not target_disk:
            # Return info for all disks
            all_info = {}
            for disk in self.get_all_disks():
                all_info[disk] = self.get_disk_info(disk)
            return all_info
        
        # Check cache
        if target_disk in self._info_cache:
            return self._info_cache[target_disk]
        
        # Gather all information
        info = {
            'name': target_disk,
            'model': '',
            'serial': '',
            'size_bytes': None,
            'size_human': None,
            'firmware': None,
            'interface': None,
            'type': None,  # HDD, SSD, NVMe, Virtual
            'rotational': None,  # True for HDD, False for SSD
            'rpm': None,  # Rotational speed for HDD
            'sector_size': None,
            'sector_size_logical': None,
            'sector_size_physical': None,
            'smart_supported': None,
            'smart_enabled': None,
            'temperature': None,
            'power_on_hours': None,
            'power_on_count': None,
            'bad_sectors': None,  # Reallocated sectors count
            'pending_sectors': None,  # Pending reallocation
            'uncorrectable_sectors': None,
            'health_status': None,  # PASSED, FAILED, UNKNOWN
            'vendor': '',
            'removable': None,
            'wwn': '',  # World Wide Name
            'device_path': None,
            'partitions': None,
            'is_virtual': False,
            'virtual_type': None,  # VMware, VirtualBox, KVM, etc.
        }
        
        # Try to get information from sysfs (most reliable for basic info)
        self._get_info_from_sysfs(target_disk, info)
        
        # Try to get SMART information (requires smartctl)
        self._get_smart_info(target_disk, info)
        
        # Try to get additional info from /proc/scsi/scsi
        self._get_scsi_info(target_disk, info)
        
        # Detect virtualization
        self._detect_virtualization(target_disk, info)
        
        # Get partition information
        info['partitions'] = self._get_partitions(target_disk)
        
        # Cache the results
        self._info_cache[target_disk] = info
        
        return info
    
    def _get_info_from_sysfs(self, disk: str, info: Dict[str, Any]):
        """Get disk information from sysfs (/sys/block/)."""
        base_path = f'/sys/block/{disk}'
        
        if not os.path.exists(base_path):
            return
        
        # Device path
        info['device_path'] = f'/dev/{disk}'
        
        # Size in sectors (each sector usually 512 bytes)
        size_path = f'{base_path}/size'
        if os.path.exists(size_path):
            try:
                with open(size_path, 'r') as f:
                    sectors = int(f.read().strip())
                    sector_size = 512  # Default standard
                    # Try to get actual sector size
                    if os.path.exists(f'{base_path}/queue/hw_sector_size'):
                        with open(f'{base_path}/queue/hw_sector_size', 'r') as f2:
                            sector_size = int(f2.read().strip())
                    info['size_bytes'] = sectors * sector_size
                    info['size_human'] = self._bytes_to_human(info['size_bytes'])
                    info['sector_size'] = sector_size
            except Exception:
                pass
        
        # Logical and physical sector size
        logical_path = f'{base_path}/queue/logical_block_size'
        physical_path = f'{base_path}/queue/physical_block_size'
        if os.path.exists(logical_path):
            try:
                with open(logical_path, 'r') as f:
                    info['sector_size_logical'] = int(f.read().strip())
            except Exception:
                pass
        if os.path.exists(physical_path):
            try:
                with open(physical_path, 'r') as f:
                    info['sector_size_physical'] = int(f.read().strip())
            except Exception:
                pass
        
        # Rotational (spinning disk vs SSD)
        rotational_path = f'{base_path}/queue/rotational'
        if os.path.exists(rotational_path):
            try:
                with open(rotational_path, 'r') as f:
                    rotational = int(f.read().strip())
                    info['rotational'] = rotational == 1
                    info['type'] = 'HDD' if rotational else 'SSD'
            except Exception:
                pass
        
        # Removable
        removable_path = f'{base_path}/removable'
        if os.path.exists(removable_path):
            try:
                with open(removable_path, 'r') as f:
                    info['removable'] = int(f.read().strip()) == 1
            except Exception:
                pass
        
        # Model from device directory
        device_path = f'{base_path}/device'
        if os.path.exists(device_path):
            # Try different model file locations
            for model_file in ['model', 'product', 'inquiry', 'device/model']:
                full_path = os.path.join(device_path, model_file)
                if os.path.exists(full_path):
                    try:
                        with open(full_path, 'r') as f:
                            model = f.read().strip()
                            if model and not model.isspace():
                                info['model'] = model
                                break
                    except Exception:
                        pass
            
            # Firmware version
            for fw_file in ['fwrev', 'firmware_rev', 'rev']:
                full_path = os.path.join(device_path, fw_file)
                if os.path.exists(full_path):
                    try:
                        with open(full_path, 'r') as f:
                            info['firmware'] = f.read().strip()
                            break
                    except Exception:
                        pass
            
            # Serial number
            for serial_file in ['serial', 'wwid']:
                full_path = os.path.join(device_path, serial_file)
                if os.path.exists(full_path):
                    try:
                        serial = f.read().strip() if os.path.isfile(full_path) else None
                        if serial and not serial.isspace():
                            info['serial'] = serial
                            break
                    except Exception:
                        pass
            
            # WWN (World Wide Name)
            wwn_path = f'{base_path}/wwid'
            if os.path.exists(wwn_path):
                try:
                    with open(wwn_path, 'r') as f:
                        wwn = f.read().strip()
                        if wwn.startswith('naa.'):
                            info['wwn'] = wwn
                except Exception:
                    pass
            
            # Interface type
            if os.path.exists(f'{base_path}/device/subsystem'):
                try:
                    subsystem = os.path.realpath(f'{base_path}/device/subsystem')
                    if 'nvme' in subsystem:
                        info['interface'] = 'NVMe'
                    elif 'ata' in subsystem or 'scsi' in subsystem:
                        # Determine SATA or SAS
                        if os.path.exists(f'{device_path}/sas_address'):
                            info['interface'] = 'SAS'
                        else:
                            info['interface'] = 'SATA'
                    elif 'virtio' in subsystem:
                        info['interface'] = 'VirtIO'
                except Exception:
                    pass
    
    def _get_scsi_info(self, disk: str, info: Dict[str, Any]):
        """Get SCSI disk information from /proc/scsi/scsi."""
        try:
            with open('/proc/scsi/scsi', 'r') as f:
                content = f.read()
            
            # Find matching disk entry
            patterns = [
                rf'Host:\s+scsi\d+\s+Channel:\s+\d+\s+Id:\s+\d+\s+Lun:\s+\d+\s*\n\s*Vendor:\s+(\S+)\s+Model:\s+(.+?)\s+Rev:\s+(\S+)\s*\n\s*Type:\s+\S+\s+ANSI\s+SCSI\s+revision:\s+\d+',
            ]
            
            for pattern in patterns:
                matches = re.findall(pattern, content, re.MULTILINE)
                for vendor, model, rev in matches:
                    if disk in model.lower() or disk in vendor.lower():
                        info['vendor'] = vendor.strip()
                        if not info['model']:
                            info['model'] = model.strip()
                        if not info['firmware']:
                            info['firmware'] = rev.strip()
                        break
        except Exception:
            pass
    
    def _get_smart_info(self, disk: str, info: Dict[str, Any]):
        """
        Get SMART information using smartctl command.
        This requires smartmontools package to be installed.
        """
        if not self._check_command_exists('smartctl'):
            info['smart_supported'] = False
            return
        
        try:
            # Check if SMART is supported
            result = subprocess.run(
                ['smartctl', '-i', f'/dev/{disk}'],
                capture_output=True,
                text=True,
                timeout=10
            )
            
            if 'SMART support is:' in result.stdout:
                smart_lines = result.stdout
                if 'SMART support is: Available' in smart_lines:
                    info['smart_supported'] = True
                if 'SMART support is: Enabled' in smart_lines:
                    info['smart_enabled'] = True
            
            # Get health information
            result = subprocess.run(
                ['smartctl', '-H', f'/dev/{disk}'],
                capture_output=True,
                text=True,
                timeout=10
            )
            
            if 'SMART overall-health self-assessment test result:' in result.stdout:
                health_line = re.search(
                    r'SMART overall-health self-assessment test result:\s+(.+)',
                    result.stdout
                )
                if health_line:
                    info['health_status'] = health_line.group(1).strip()
            
            # Get detailed SMART attributes
            result = subprocess.run(
                ['smartctl', '-A', f'/dev/{disk}'],
                capture_output=True,
                text=True,
                timeout=10
            )
            
            # Parse SMART attributes
            attributes = {
                'Reallocated_Sector_Ct': 'bad_sectors',
                'Current_Pending_Sector': 'pending_sectors',
                'Offline_Uncorrectable': 'uncorrectable_sectors',
                'Temperature_Celsius': 'temperature',
                'Power_On_Hours': 'power_on_hours',
                'Power_Cycle_Count': 'power_on_count',
                'Rotational_Rate': 'rpm'
            }
            
            for line in result.stdout.split('\n'):
                for attr, key in attributes.items():
                    if attr in line:
                        # Parse the attribute value (usually the 10th column)
                        parts = line.split()
                        if len(parts) >= 10:
                            try:
                                value = parts[9]  # RAW_VALUE column
                                if value.isdigit():
                                    if key == 'rpm' and int(value) > 0:
                                        info['rpm'] = int(value)
                                        if info['type'] == 'HDD':
                                            info['type'] = 'HDD'
                                    elif key != 'rpm':
                                        info[key] = int(value)
                            except (ValueError, IndexError):
                                pass
                        break
            
            # Try to get rotational rate for HDDs from other sources
            if info.get('rotational') and not info.get('rpm'):
                # Try to get from sysfs
                rpm_path = f'/sys/block/{disk}/device/rotational_rate'
                if os.path.exists(rpm_path):
                    try:
                        with open(rpm_path, 'r') as f:
                            rpm = int(f.read().strip())
                            if rpm > 0:
                                info['rpm'] = rpm
                    except Exception:
                        pass
        
        except subprocess.TimeoutExpired:
            pass
        except FileNotFoundError:
            info['smart_supported'] = False
        except Exception:
            pass
    
    def _detect_virtualization(self, disk: str, info: Dict[str, Any]):
        """Detect if disk is virtual (VMware, VirtualBox, KVM, etc.)."""
        info['is_virtual'] = False
        
        # Check disk model for virtualization indicators
        model = info.get('model', '').lower()
        vendor = info.get('vendor', '').lower()
        
        if 'vmware' in model or 'vmware' in vendor:
            info['is_virtual'] = True
            info['virtual_type'] = 'VMware'
            info['type'] = 'Virtual'
        elif 'vbox' in model or 'virtualbox' in model:
            info['is_virtual'] = True
            info['virtual_type'] = 'VirtualBox'
            info['type'] = 'Virtual'
        elif 'qemu' in model or 'kvm' in model:
            info['is_virtual'] = True
            info['virtual_type'] = 'KVM/QEMU'
            info['type'] = 'Virtual'
        elif 'virtio' in model:
            info['is_virtual'] = True
            info['virtual_type'] = 'VirtIO'
            info['type'] = 'Virtual'
        
        # Check system virtualization
        try:
            with open('/proc/cpuinfo', 'r') as f:
                if 'hypervisor' in f.read().lower():
                    if not info['is_virtual']:
                        info['is_virtual'] = True
                        info['virtual_type'] = 'Unknown Hypervisor'
        except Exception:
            pass
    
    def _get_partitions(self, disk: str) -> Optional[List[Dict[str, Any]]]:
        """Get partition information for the disk."""
        partitions = []
        
        try:
            # Read from /proc/partitions
            with open('/proc/partitions', 'r') as f:
                for line in f:
                    parts = line.strip().split()
                    if len(parts) >= 4:
                        device = parts[3]
                        if device.startswith(disk) and device != disk:
                            size_kb = int(parts[2]) if parts[2].isdigit() else 0
                            partitions.append({
                                'name': device,
                                'size_kb': size_kb,
                                'size_human': self._bytes_to_human(size_kb * 1024)
                            })
        except Exception:
            pass
        
        return partitions if partitions else None
    
    def get_disk_summary(self, disk_name: Optional[str] = None) -> str:
        """Get a human-readable summary for a disk."""
        info = self.get_disk_info(disk_name)
        
        if not info:
            return f"No information available for disk: {disk_name}"
        
        if isinstance(info, dict) and 'name' in info:
            # Single disk
            return self._format_disk_summary(info)
        elif isinstance(info, dict):
            # Multiple disks
            summaries = []
            for disk, disk_info in info.items():
                summaries.append(self._format_disk_summary(disk_info))
            return '\n\n'.join(summaries)
        
        return "No disk information available"
    
    def _format_disk_summary(self, info: Dict[str, Any]) -> str:
        """Format disk information as a human-readable string."""
        lines = []
        lines.append(f"Disk: {info['name']} ({info['device_path']})")
        lines.append(f"{'=' * 60}")
        
        if info.get('model'):
            lines.append(f"  Model: {info['model']}")
        if info.get('vendor'):
            lines.append(f"  Vendor: {info['vendor']}")
        if info.get('serial'):
            lines.append(f"  Serial Number: {info['serial']}")
        if info.get('firmware'):
            lines.append(f"  Firmware: {info['firmware']}")
        if info.get('wwn'):
            lines.append(f"  WWN: {info['wwn']}")
        
        lines.append(f"  Type: {info.get('type', 'Unknown')}")
        if info.get('rotational') is not None:
            lines.append(f"  Rotational: {'Yes' if info['rotational'] else 'No'}")
        if info.get('rpm'):
            lines.append(f"  RPM: {info['rpm']}")
        
        if info.get('size_human'):
            lines.append(f"  Size: {info['size_human']} ({info.get('size_bytes', 0):,} bytes)")
        
        if info.get('interface'):
            lines.append(f"  Interface: {info['interface']}")
        
        if info.get('sector_size'):
            lines.append(f"  Sector Size: {info['sector_size']} bytes")
            if info.get('sector_size_logical'):
                lines.append(f"  Logical Sector Size: {info['sector_size_logical']} bytes")
            if info.get('sector_size_physical'):
                lines.append(f"  Physical Sector Size: {info['sector_size_physical']} bytes")
        
        lines.append(f"  Removable: {'Yes' if info.get('removable') else 'No'}")
        lines.append(f"  Virtual: {'Yes' if info.get('is_virtual') else 'No'}")
        if info.get('virtual_type'):
            lines.append(f"  Virtual Type: {info['virtual_type']}")
        
        # SMART information
        if info.get('smart_supported'):
            lines.append("  SMART:")
            lines.append(f"    Supported: Yes")
            lines.append(f"    Enabled: {'Yes' if info.get('smart_enabled') else 'No'}")
            if info.get('health_status'):
                lines.append(f"    Health: {info['health_status']}")
            if info.get('temperature'):
                lines.append(f"    Temperature: {info['temperature']}°C")
            if info.get('power_on_hours'):
                days = info['power_on_hours'] // 24
                lines.append(f"    Power On Hours: {info['power_on_hours']} hours ({days} days)")
            if info.get('power_on_count'):
                lines.append(f"    Power Cycle Count: {info['power_on_count']}")
            
            # Bad sector information
            bad_info = []
            if info.get('bad_sectors'):
                bad_info.append(f"Reallocated: {info['bad_sectors']}")
            if info.get('pending_sectors'):
                bad_info.append(f"Pending: {info['pending_sectors']}")
            if info.get('uncorrectable_sectors'):
                bad_info.append(f"Uncorrectable: {info['uncorrectable_sectors']}")
            
            if bad_info:
                lines.append(f"    Bad Sectors: {', '.join(bad_info)}")
        else:
            lines.append("  SMART: Not supported or smartctl not installed")
        
        # Partitions
        if info.get('partitions'):
            lines.append(f"  Partitions: {len(info['partitions'])}")
            for part in info['partitions'][:5]:  # Show first 5
                lines.append(f"    - {part['name']}: {part['size_human']}")
            if len(info['partitions']) > 5:
                lines.append(f"    ... and {len(info['partitions']) - 5} more")
        
        return '\n'.join(lines)
    
    def check_disk_health(self, disk_name: Optional[str] = None) -> Dict[str, Any]:
        """
        Check disk health and return a summary.
        Returns dictionary with health status and recommendations.
        """
        info = self.get_disk_info(disk_name)
        
        if not info or (isinstance(info, dict) and 'name' not in info):
            return {'error': 'No disk information available'}
        
        if isinstance(info, dict) and 'name' in info:
            disks_info = {info['name']: info}
        else:
            disks_info = info
        
        health_summary = {}
        
        for disk_name, disk_info in disks_info.items():
            health = {
                'disk': disk_name,
                'status': 'UNKNOWN',
                'issues': [],
                'warnings': [],
                'recommendations': []
            }
            
            # Check SMART health status
            if disk_info.get('health_status'):
                if 'PASSED' in disk_info['health_status']:
                    health['status'] = 'GOOD'
                elif 'FAILED' in disk_info['health_status']:
                    health['status'] = 'FAILED'
                    health['issues'].append('SMART health test failed')
                    health['recommendations'].append('Backup data immediately and replace disk')
            
            # Check bad sectors
            bad_sectors = disk_info.get('bad_sectors', 0)
            pending = disk_info.get('pending_sectors', 0)
            uncorrectable = disk_info.get('uncorrectable_sectors', 0)
            
            if bad_sectors > 0:
                health['warnings'].append(f'Reallocated sectors: {bad_sectors}')
                if bad_sectors > 100:
                    health['issues'].append(f'High number of reallocated sectors: {bad_sectors}')
                    health['recommendations'].append('Consider replacing disk soon')
            
            if pending > 0:
                health['warnings'].append(f'Pending sectors: {pending}')
                health['recommendations'].append('Run full disk surface test')
            
            if uncorrectable > 0:
                health['issues'].append(f'Uncorrectable sectors: {uncorrectable}')
                health['recommendations'].append('Disk may have physical damage, replace immediately')
            
            # Check temperature
            temp = disk_info.get('temperature')
            if temp:
                if temp > 70:
                    health['warnings'].append(f'High temperature: {temp}°C')
                    health['recommendations'].append('Improve cooling')
                elif temp > 60:
                    health['warnings'].append(f'Elevated temperature: {temp}°C')
            
            # Check power on hours (for HDDs)
            if disk_info.get('type') == 'HDD' and disk_info.get('power_on_hours'):
                hours = disk_info['power_on_hours']
                years = hours / 8760  # 365 * 24
                if years > 5:
                    health['warnings'].append(f'Disk age: {years:.1f} years of power-on time')
                    health['recommendations'].append('Consider replacing due to age')
            
            if not health['issues'] and not health['warnings']:
                health['status'] = 'EXCELLENT'
            elif not health['issues']:
                health['status'] = 'WARNING'
            
            health_summary[disk_name] = health
        
        return health_summary
    
    @staticmethod
    def _check_command_exists(command: str) -> bool:
        """Check if a command exists in the system."""
        try:
            subprocess.run(
                ['which', command],
                capture_output=True,
                text=True,
                timeout=5
            )
            return True
        except Exception:
            return False
    
    @staticmethod
    def _bytes_to_human(bytes_val: Optional[int]) -> str:
        """Convert bytes to human-readable format."""
        if bytes_val is None:
            return "Unknown"
        
        for unit in ['B', 'KB', 'MB', 'GB', 'TB', 'PB']:
            if bytes_val < 1024.0:
                return f"{bytes_val:.2f} {unit}"
            bytes_val /= 1024.0
        return f"{bytes_val:.2f} EB"
    
    def refresh(self) -> 'DiskInfo':
        """Clear cache and refresh disk information."""
        self._info_cache = {}
        return self
    
    def __repr__(self) -> str:
        if self.disk_name:
            return f"DiskInfo(disk={self.disk_name})"
        else:
            disks = self.get_all_disks()
            return f"DiskInfo(disks={disks})"


# ------------------------------------------------------------------
# Example usage and test
# ------------------------------------------------------------------
if __name__ == '__main__':
    disk = DiskInfo()
    
    print("=== All Disks ===")
    all_disks = disk.get_all_disks()
    print(f"Found disks: {', '.join(all_disks)}")
    print()
    
    # Get summary for all disks
    print(disk.get_disk_summary())
    print()
    
    # Get detailed information for each disk
    for disk_name in all_disks:
        print(f"\n{'=' * 80}")
        print(f"Detailed information for {disk_name}:")
        print(f"{'=' * 80}")
        
        info = disk.get_disk_info(disk_name)
        if info:
            print(f"\nJSON-like structure (partial):")
            for key, value in info.items():
                if key != 'partitions' and value is not None:
                    print(f"  {key}: {value}")
            
            if info.get('partitions'):
                print(f"  partitions: {len(info['partitions'])} partitions found")
    
    # Check health for all disks
    print(f"\n{'=' * 80}")
    print("Disk Health Check")
    print(f"{'=' * 80}")
    health = disk.check_disk_health()
    for disk_name, status in health.items():
        print(f"\n{disk_name}:")
        print(f"  Status: {status['status']}")
        if status['warnings']:
            print(f"  Warnings: {', '.join(status['warnings'])}")
        if status['issues']:
            print(f"  Issues: {', '.join(status['issues'])}")
        if status['recommendations']:
            print(f"  Recommendations: {', '.join(status['recommendations'])}")
    
    # Example of setting a specific disk
    if all_disks:
        print(f"\n{'=' * 80}")
        print(f"Working with specific disk: {all_disks[0]}")
        print(f"{'=' * 80}")
        disk.set_disk(all_disks[0])
        print(f"Disk model: {disk.get_disk_info().get('model', 'Unknown')}")
        print(f"Disk size: {disk.get_disk_info().get('size_human', 'Unknown')}")