"""WizardKit: Disk object and functions"""
# vim: sts=2 sw=2 ts=2

import copy
import logging
import pathlib
import platform
import plistlib
import re

from dataclasses import dataclass, field
from typing import Any, Union

from wk.cfg.main import KIT_NAME_SHORT
from wk.cfg.python import DATACLASS_DECORATOR_KWARGS
from wk.exe import get_json_from_command, run_program
from wk.hw.test import Test
from wk.hw.smart import (
  generate_attribute_report,
  get_known_disk_attributes,
  )
from wk.std import PLATFORM
from wk.ui import ansi


# STATIC VARIABLES
LOG = logging.getLogger(__name__)
WK_LABEL_REGEX = re.compile(
  fr'{KIT_NAME_SHORT}_(LINUX|UFD)',
  re.IGNORECASE,
  )


# Classes
@dataclass(**DATACLASS_DECORATOR_KWARGS)
class Disk:
  """Object for tracking disk specific data."""
  attributes: dict[Any, dict] =         field(init=False, default_factory=dict)
  bus: str =                            field(init=False)
  children: list[dict] =                field(init=False, default_factory=list)
  description: str =                    field(init=False)
  filesystem: str =                     field(init=False)
  initial_attributes: dict[Any, dict] = field(init=False)
  known_attributes: dict[Any, dict] =   field(init=False, default_factory=dict)
  log_sec: int =                        field(init=False)
  model: str =                          field(init=False)
  name: str =                           field(init=False)
  notes: list[str] =                    field(init=False, default_factory=list)
  path: Union[pathlib.Path, str]
  parent: str  =                        field(init=False)
  phy_sec: int =                        field(init=False)
  raw_details: dict[str, Any] =         field(init=False)
  raw_smartctl: dict[str, Any] =        field(init=False)
  serial: str =                         field(init=False)
  size: int =                           field(init=False)
  ssd: bool =                           field(init=False)
  tests: list[Test] =                   field(init=False, default_factory=list)
  use_sat: bool =                       field(init=False, default=False)

  def __post_init__(self) -> None:
    self.path = pathlib.Path(self.path).resolve()
    self.update_details()
    self.set_description()
    self.known_attributes = get_known_disk_attributes(self.model)
    if not self.attributes and self.bus == 'USB':
      # Try using SAT
      LOG.warning('Using SAT for smartctl for %s', self.path)
      self.notes = []
      self.use_sat = True
    self.initial_attributes = copy.deepcopy(self.attributes)
    if not self.is_4k_aligned():
      self.add_note('One or more partitions are not 4K aligned', 'YELLOW')

  def add_note(self, note, color=None) -> None:
    """Add note that will be included in the disk report."""
    if color:
      note = ansi.color_string(note, color)
    if note not in self.notes:
      self.notes.append(note)
      self.notes.sort()

  def contains_note(self, note_str) -> bool:
    """Check if note is already present."""
    present = False
    for note in self.notes:
      if note_str == ansi.strip_colors(note):
        present = True
    return present

  def disable_disk_tests(self) -> None:
    """Disable all tests."""
    LOG.warning('Disabling all tests for: %s', self.path)
    for test in self.tests:
      if test.status in ('Pending', 'Working'):
        test.set_status('Denied')
      test.disabled = True

  def generate_report(self, header=True) -> list[str]:
    """Generate Disk report, returns list."""
    report = []
    if header:
      report.append(ansi.color_string(f'Device ({self.path.name})', 'BLUE'))
      report.append(f'  {self.description}')

    # Attributes
    if self.attributes:
      if header:
        report.append(ansi.color_string('Attributes', 'BLUE'))
      report.extend(generate_attribute_report(self))

    # Notes
    if self.notes:
      report.append(ansi.color_string('Notes', 'BLUE'))
    for note in self.notes:
      report.append(f'  {note}')

    # Tests
    for test in self.tests:
      report.extend(test.report)

    return report

  def get_labels(self) -> list[str]:
    """Build list of labels for this disk, returns list."""
    labels = []

    # Add all labels from raw_details
    for details in [self.raw_details, *self.raw_details.get('children', [])]:
      labels.append(details.get('label', ''))
      labels.append(details.get('partlabel', ''))

    # Remove empty labels
    labels = [str(label) for label in labels if label]

    # Done
    return labels

  def is_4k_aligned(self) -> bool:
    """Check that all disk partitions are aligned, returns bool."""
    aligned = True
    if PLATFORM == 'Darwin':
      aligned = is_4k_aligned_macos(self.raw_details)
    elif PLATFORM == 'Linux':
      aligned = is_4k_aligned_linux(self.path, self.phy_sec)

    return aligned

  @property
  def present(self) -> bool:
    """Verify this device is still present, returns bool."""
    if not self.path.exists():
      self.add_note('Device disconnected', 'RED')
      return False
    return True

  def set_description(self) -> None:
    """Set disk description from details."""
    decimals = 1
    suffix = ' '

    # Set size_str (try binary scale first)
    for scale in (1024, 1000):
      size = float(self.size)
      units = list('KMGTPEZY')
      while units:
        if abs(size) < scale:
          break
        size /= scale
        suffix = units.pop(0)
      size = ((size * 10) // 1) / 10
      if size % 1 == 0:
        # Found an exact whole number, drop the decimal
        decimals = 0
        break
      if size % 1 == 0.5:
        break

    # Done
    self.description = (
      f'{size:0.{decimals}f} {suffix}B ({self.bus}) {self.model} {self.serial}'
      )

  def update_details(self, skip_children=True) -> None:
    """Update disk details using OS specific methods.

    Required details default to generic descriptions
    and are converted to the correct type.
    """
    if not self.present:
      return

    if PLATFORM == 'Darwin':
      self.raw_details = get_disk_details_macos(
        self.path, skip_children=skip_children,
        )
    elif PLATFORM == 'Linux':
      self.raw_details = get_disk_details_linux(
        self.path, skip_children=skip_children,
        )

    # Set necessary details
    self.bus =    str(self.raw_details.get('bus', '???')).upper()
    self.bus =        self.bus.replace('IMAGE', 'Image')
    self.bus =        self.bus.replace('NVME', 'NVMe')
    self.children =   self.raw_details.get('children', [])
    self.filesystem = self.raw_details.get('fstype', 'Unknown')
    self.log_sec =    self.raw_details.get('log-sec', 512)
    self.model =      self.raw_details.get('model', 'Unknown Model')
    self.name =       self.raw_details.get('name', self.path)
    self.parent =     self.raw_details.get('parent', None)
    self.phy_sec =    self.raw_details.get('phy-sec', 512)
    self.serial =     self.raw_details.get('serial', 'Unknown Serial')
    self.size =       self.raw_details.get('size', -1)
    self.ssd =        self.raw_details.get('ssd', False)

    # Ensure certain attributes types
    ## NOTE: This is ugly, deal.
    for attr in ['bus', 'model', 'name', 'serial']:
      setattr(self, attr, str(getattr(self, attr)))
    for attr in ['log_sec', 'phy_sec', 'size']:
      try:
        setattr(self, attr, int(getattr(self, attr)))
      except (TypeError, ValueError):
        LOG.error('Invalid disk %s: %s', attr, getattr(self, attr))
        if attr == 'size':
          setattr(self, attr, -1)


# Functions
def get_disk_details_linux(disk_path, skip_children=True) -> dict[Any, Any]:
  """Get disk details using lsblk, returns dict."""
  def _flatten_json(dev) -> list:
    """Convert lsblk JSON tree to a flat list of items, returns list."""
    devs = [dev]
    for child in dev.pop('children', []):
      devs.extend(_flatten_json(child))
    return devs

	# Get lsblk info
  cmd = ['lsblk', '--bytes', '--json', '--output-all', '--paths', disk_path]
  if skip_children:
    cmd.append('--nodeps')
  json_data = get_json_from_command(cmd, check=False)
  dev_list = _flatten_json(json_data.get('blockdevices', [{}])[0])

  # Fix details
  for dev in dev_list:
    dev['bus'] = dev.pop('tran', '???')
    dev['parent'] = dev.pop('pkname', None)
    dev['ssd'] = not dev.pop('rota', True)
    if 'loop' in str(disk_path) and dev['bus'] is None:
      dev['bus'] = 'Image'
      dev['model'] = ''
      dev['serial'] = ''

  # Convert to dict
  details = dev_list.pop(0)
  details['children'] = dev_list

  # Done
  return details


def get_disk_details_macos(disk_path, skip_children=True) -> dict:
  """Get disk details using diskutil, returns dict."""
  details = {}
  disk_path = pathlib.Path(disk_path)

  # Get "list" details
  cmd = ['diskutil', 'list', '-plist', disk_path]
  proc = run_program(cmd, check=False, encoding=None, errors=None)
  try:
    plist_data = plistlib.loads(proc.stdout)
  except (TypeError, ValueError):
    # Invalid / corrupt plist data? return empty dict to avoid crash
    LOG.error('Failed to get diskutil list for %s', disk_path)
    return details

  # Parse "list" details
  details = plist_data.get('AllDisksAndPartitions', [{}])[0]
  details['path'] = disk_path
  if skip_children:
    details['children'] = []
  else:
    details['children'] = details.pop('Partitions', [])
    details['children'].extend(details.pop('APFSVolumes', []))
    for child in details['children']:
      child['path'] = disk_path.with_name(child['DeviceIdentifier'])

  # Get "info" details
  for dev in [details, *details['children']]:
    cmd = ['diskutil', 'info', '-plist', dev['path']]
    proc = run_program(cmd, check=False, encoding=None, errors=None)
    try:
      plist_data = plistlib.loads(proc.stdout)
    except (TypeError, ValueError):
      LOG.error('Failed to get diskutil info for %s', dev['path'])
      continue

    # Parse "info" details
    dev.update(plist_data)
    dev['bus'] = dev.pop('BusProtocol', '???')
    dev['fstype'] = dev.pop('FilesystemType', '')
    dev['label'] = dev.pop('VolumeName', '')
    dev['model'] = dev.pop('MediaName', 'Unknown')
    dev['mountpoint'] = dev.pop('MountPoint', '')
    dev['name'] = dev.get('name', str(dev['path']))
    dev['phy-sec'] = dev.pop('DeviceBlockSize', 512)
    dev['serial'] = get_disk_serial_macos(dev['path'])
    dev['size'] = dev.pop('Size', -1)
    dev['ssd'] = dev.pop('SolidState', False)
    dev['vendor'] = ''
    if dev.get('WholeDisk', True):
      dev['parent'] = None
    else:
      dev['parent'] = dev.pop('ParentWholeDisk', None)

  # Fix details if main dev is a child
  for child in details['children']:
    if disk_path == child['path']:
      for key in ('fstype', 'label', 'name', 'size'):
        details[key] = child[key]
      break

  # Done
  return details


def get_disk_serial_macos(path) -> str:
  """Get disk serial using system_profiler, returns str."""
  cmd = ['sudo', 'smartctl', '--info', '--json', path]
  smart_info = get_json_from_command(cmd)
  return smart_info.get('serial_number', 'Unknown Serial')


def get_disks(skip_kits=False) -> list[Disk]:
  """Get disks using OS-specific methods, returns list."""
  disks = []
  if PLATFORM == 'Darwin':
    disks = get_disks_macos()
  elif PLATFORM == 'Linux':
    disks = get_disks_linux()

  # Skip WK disks
  if skip_kits:
    for disk in disks:
      disk.update_details(skip_children=False)
    disks = [
      disk_obj for disk_obj in disks
      if not any(
        WK_LABEL_REGEX.search(label) for label in disk_obj.get_labels()
        )
      ]

  # Done
  return disks


def get_disks_linux() -> list[Disk]:
  """Get disks via lsblk, returns list."""
  cmd = ['lsblk', '--json', '--nodeps', '--paths']
  disks = []

  # Add valid disks
  json_data = get_json_from_command(cmd)
  for disk in json_data.get('blockdevices', []):
    disk_obj = Disk(disk['name'])

    # Skip loopback devices, optical devices, etc
    if disk_obj.raw_details.get('type', '???') != 'disk':
      continue

    # Skip empty devices (usually card readers)
    if disk_obj.size <= 0:
      continue

    # Add disk
    disks.append(disk_obj)

  # Done
  return disks


def get_disks_macos() -> list[Disk]:
  """Get disks via diskutil, returns list."""
  cmd = ['diskutil', 'list', '-plist', 'physical']
  disks = []

  # El Capitan workaround
  if platform.mac_ver()[0].startswith('10.11'):
    cmd.pop()

  # Get info from diskutil
  proc = run_program(cmd, encoding=None, errors=None, check=False)

  # Parse plist data
  try:
    plist_data = plistlib.loads(proc.stdout)
  except (TypeError, ValueError):
    # Invalid / corrupt plist data? return empty list to avoid crash
    LOG.error('Failed to get diskutil list')
    return disks

  # Add valid disks
  for disk in plist_data['AllDisksAndPartitions']:
    name = disk['DeviceIdentifier']
    if name not in plist_data['WholeDisks']:
      # Only check "WholeDisks"
      continue
    if not disk['Content']:
      # This lists GPT or MBR for device, blank should mean it's an image
      continue
    disks.append(Disk(f'/dev/{name}'))

  # Remove virtual disks
  disks = [
    d for d in disks if d.raw_details.get('VirtualOrPhysical') != 'Virtual'
    ]

  # Done
  return disks


def is_4k_aligned_macos(disk_details) -> bool:
  """Check partition alignment using diskutil info, returns bool."""
  aligned = True

  # Check partitions
  for part in disk_details.get('children', []):
    offset = part.get('PartitionMapPartitionOffset', 0)
    if not offset:
      # Assuming offset couldn't be found and it defaulted to 0
      # NOTE: Just logging the error, not bailing
      LOG.error('Failed to get partition offset for %s', part['path'])
    aligned = aligned and offset >= 0 and offset % 4096 == 0

  # Done
  return aligned


def is_4k_aligned_linux(dev_path, physical_sector_size) -> bool:
  """Check partition alignment using lsblk, returns bool."""
  aligned = True
  cmd = [
    'sudo',
    'sfdisk',
    '--json',
    dev_path,
    ]

  # Get partition details
  json_data = get_json_from_command(cmd)

  # Check partitions
  for part in json_data.get('partitiontable', {}).get('partitions', []):
    offset = physical_sector_size * part.get('start', -1)
    aligned = aligned and offset >= 0 and offset % 4096 == 0

  # Done
  return aligned


if __name__ == '__main__':
  print("This file is not meant to be called directly.")