"""WizardKit: ddrescue TUI - Block Pairs"""
# vim: sts=2 sw=2 ts=2

import logging
import math
import os
import pathlib
import plistlib
import re
import subprocess

from wk import cfg, exe, std
from wk.clone import menus
from wk.hw import disk as hw_disk
from wk.ui import cli


# STATIC VARIABLES
LOG = logging.getLogger(__name__)
DDRESCUE_LOG_REGEX = re.compile(
  r'^\s*(?P<key>\S+):\s+'
  r'(?P<size>\d+)\s+'
  r'(?P<unit>[PTGMKB]i?B?)'
  r'.*\(\s*(?P<percent>\d+\.?\d*)%\)$',
  re.IGNORECASE,
  )


# Classes
class BlockPair():
  """Object for tracking source to dest recovery data."""
  def __init__(
      self,
      source_dev: hw_disk.Disk,
      destination: pathlib.Path,
      working_dir: pathlib.Path,
      ):
    self.sector_size: int = source_dev.phy_sec
    self.source: pathlib.Path = pathlib.Path(source_dev.path)
    self.destination: pathlib.Path = destination
    self.map_data: dict[str, bool | int] = {}
    self.map_path: pathlib.Path = pathlib.Path()
    self.size: int = source_dev.size
    self.status: dict[str, float | int | str] = {
      'read-skip':  'Pending',
      'read-full':  'Pending',
      'trim':       'Pending',
      'scrape':     'Pending',
      }
    self.test_map: pathlib.Path | None = None
    self.view_map: bool = 'DISPLAY' in os.environ or 'WAYLAND_DISPLAY' in os.environ
    self.view_proc: subprocess.Popen | None = None

    # Set map path
    # e.g. '(Clone|Image)_Model[_p#]_Size[_Label].map'
    map_name = source_dev.model
    if source_dev.bus == 'Image':
      map_name = 'Image'
    if source_dev.parent:
      part_num = re.sub(r"^.*?(\d+)$", r"\1", self.source.name)
      map_name += f'_p{part_num}'
    size_str = std.bytes_to_string(
      size=self.size,
      use_binary=False,
      )
    map_name += f'_{size_str.replace(" ", "")}'
    if source_dev.raw_details.get('label', ''):
      map_name += f'_{source_dev.raw_details["label"]}'
    map_name = map_name.replace(' ', '_')
    map_name = map_name.replace('/', '_')
    map_name = map_name.replace('\\', '_')
    if destination.is_dir():
      # Imaging
      self.map_path = pathlib.Path(f'{destination}/Image_{map_name}.map')
      self.destination = self.map_path.with_suffix('.dd')
      self.destination.touch()
    else:
      # Cloning
      self.map_path = pathlib.Path(f'{working_dir}/Clone_{map_name}.map')

    # Create map file if needed
    # NOTE: We need to set the domain size for --complete-only to work
    if not self.map_path.exists():
      self.map_path.write_text(
        data=cfg.ddrescue.DDRESCUE_MAP_TEMPLATE.format(
          name=cfg.main.KIT_NAME_FULL,
          size=self.size,
          ),
        encoding='utf-8',
      )

    # Set initial status
    self.set_initial_status()

  def __getstate__(self):
    """Override to allow pickling ddrescue.State() objects."""
    bp_state = self.__dict__.copy()
    del bp_state['view_proc']
    return bp_state

  def get_error_size(self) -> int:
    """Get error size in bytes, returns int."""
    return self.size - self.get_rescued_size()

  def get_percent_recovered(self) -> float:
    """Get percent rescued from map_data, returns float."""
    return 100 * self.map_data.get('rescued', 0) / self.size

  def get_rescued_size(self) -> int:
    """Get rescued size using map data.

    NOTE: Returns 0 if no map data is available.
    """
    self.load_map_data()
    return self.map_data.get('rescued', 0)

  def load_map_data(self) -> None:
    """Load map data from file.

    NOTE: If the file is missing it is assumed that recovery hasn't
          started yet so default values will be returned instead.
    """
    data: dict[str, bool | int] = {'full recovery': False, 'pass completed': False}

    # Get output from ddrescuelog
    cmd = [
      'ddrescuelog',
      '--binary-prefixes',
      '--show-status',
      f'--size={self.size}',
      self.map_path,
      ]
    proc = exe.run_program(cmd, check=False)

    # Parse output
    for line in proc.stdout.splitlines():
      _r = DDRESCUE_LOG_REGEX.search(line)
      if _r:
        if _r.group('key') == 'rescued' and _r.group('percent') == '100':
          # Fix rounding errors from ddrescuelog output
          data['rescued'] = self.size
        else:
          data[_r.group('key')] = std.string_to_bytes(
            f'{_r.group("size")} {_r.group("unit")}',
            )
      data['pass completed'] = 'current status: finished' in line.lower()

    # Check if 100% done (only if map is present and non-zero size
    # NOTE: ddrescuelog returns 0 (i.e. 100% done) for empty files
    if self.map_path.exists() and self.map_path.stat().st_size != 0:
      cmd = [
        'ddrescuelog',
        '--done-status',
        f'--size={self.size}',
        self.map_path,
        ]
      proc = exe.run_program(cmd, check=False)
      data['full recovery'] = proc.returncode == 0

    # Done
    self.map_data.update(data)

  def pass_complete(self, pass_name) -> bool:
    """Check if pass_name is complete based on map data, returns bool."""
    pending_size = self.map_data['non-tried']

    # Full recovery
    if self.map_data.get('full recovery', False):
      return True

    # New recovery
    if 'non-tried' not in self.map_data:
      return False

    # Initial read skip pass
    if pass_name == 'read-skip':
      pass_threshold = cfg.ddrescue.AUTO_PASS_THRESHOLDS[pass_name]
      if self.get_percent_recovered() >= pass_threshold:
        return True

    # Recovery in progress
    if pass_name in ('trim', 'scrape'):
      pending_size += self.map_data['non-trimmed']
    if pass_name == 'scrape':
      pending_size += self.map_data['non-scraped']
    if pending_size == 0:
      # This is true when the previous and current passes are complete
      return True

    # This should never be reached
    return False

  def safety_check(self) -> None:
    """Run safety check and abort if necessary."""
    # TODO: Expand section to support non-Linux systems
    dest_size = -1
    if self.destination.is_block_device():
      cmd = [
        'lsblk', '--bytes', '--json',
        '--nodeps', '--noheadings', '--output=size',
        self.destination,
        ]
      json_data = exe.get_json_from_command(cmd)
      dest_size = json_data['blockdevices'][0]['size']
      del json_data

    # Check destination size if cloning
    if not self.destination.is_file() and dest_size < self.size:
      cli.print_error(f'Invalid destination: {self.destination}')
      raise std.GenericAbort()

  def set_initial_status(self) -> None:
    """Read map data and set initial statuses."""
    self.load_map_data()
    percent = self.get_percent_recovered()
    for name in self.status:
      if self.pass_complete(name):
        self.status[name] = percent
      else:
        # Stop checking
        if percent > 0:
          self.status[name] = percent
        break

  def skip_pass(self, pass_name) -> None:
    """Mark pass as skipped if applicable."""
    if self.status[pass_name] == 'Pending':
      self.status[pass_name] = 'Skipped'

  def update_progress(self, pass_name) -> None:
    """Update progress via map data."""
    self.load_map_data()

    # Update status
    percent = self.get_percent_recovered()
    if percent > 0:
      self.status[pass_name] = percent

    # Mark future passes as skipped if applicable
    if percent == 100:
      status_keys = list(self.status.keys())
      for pass_n in status_keys[status_keys.index(pass_name)+1:]:
        self.status[pass_n] = 'Skipped'


# Functions
def add_clone_block_pairs(state) -> None:
  """Add device to device block pairs and set settings if necessary."""
  source_sep = get_partition_separator(state.source.path.name)
  dest_sep = get_partition_separator(state.destination.path.name)
  settings = {}

  # Clone settings
  settings = state.load_settings(discard_unused_settings=True)

  # Add pairs from previous run
  if settings['Partition Mapping']:
    # Resume previous run, load pairs from settings file
    for part_map in settings['Partition Mapping']:
      bp_source = hw_disk.Disk(
        f'{state.source.path}{source_sep}{part_map[0]}',
        )
      bp_dest = pathlib.Path(
        f'{state.destination.path}{dest_sep}{part_map[1]}',
        )
      state.add_block_pair(bp_source, bp_dest)
      return

  # Add pairs from selection
  source_parts = menus.select_disk_parts('Clone', state.source)
  if state.source.path.samefile(source_parts[0].path):
    # Whole disk (or single partition via args), skip settings
    bp_dest = state.destination.path
    state.add_block_pair(state.source, bp_dest)
    return

  # New run, use new settings file
  settings['Needs Format'] = True
  offset = 0
  user_choice = cli.choice(
    'Format clone using GPT, MBR, or match Source type?',
    ['G', 'M', 'S'],
    )
  if user_choice == 'G':
    settings['Table Type'] = 'GPT'
  elif user_choice == 'M':
    settings['Table Type'] = 'MBR'
  else:
    # Match source type
    settings['Table Type'] = get_table_type(state.source.path)
  if cli.ask('Create an empty Windows boot partition on the clone?'):
    settings['Create Boot Partition'] = True
    offset = 2 if settings['Table Type'] == 'GPT' else 1

  # Add pairs
  for dest_num, part in enumerate(source_parts):
    dest_num += offset + 1
    bp_dest = pathlib.Path(
      f'{state.destination.path}{dest_sep}{dest_num}',
      )
    state.add_block_pair(part, bp_dest)

    # Add to settings file
    source_num = re.sub(r'^.*?(\d+)$', r'\1', part.path.name)
    settings['Partition Mapping'].append([source_num, dest_num])

    # Save settings
    state.save_settings(settings)


def add_image_block_pairs(state) -> None:
  """Add device to image file block pairs."""
  source_parts = menus.select_disk_parts(state.mode, state.source)
  for part in source_parts:
    state.add_block_pair(part, state.destination)


def build_sfdisk_partition_line(table_type, dev_path, size, details) -> str:
  """Build sfdisk partition line using passed details, returns str."""
  line = f'{dev_path} : size={size}'
  dest_type = ''
  source_filesystem = str(details.get('fstype', '')).upper()
  source_table_type = ''
  source_type = details.get('parttype', '')

  # Set dest type
  if re.match(r'^0x\w+$', source_type):
    # Both source and dest are MBR
    source_table_type = 'MBR'
    if table_type == 'MBR':
      dest_type = source_type.replace('0x', '').lower()
  elif re.match(r'^\w{8}-\w{4}-\w{4}-\w{4}-\w{12}$', source_type):
    # Source is a GPT type
    source_table_type = 'GPT'
    if table_type == 'GPT':
      dest_type = source_type.upper()
  if not dest_type:
    # Assuming changing table types, set based on FS
    if source_filesystem in cfg.ddrescue.PARTITION_TYPES.get(table_type, {}):
      dest_type = cfg.ddrescue.PARTITION_TYPES[table_type][source_filesystem]
  line += f', type={dest_type}'

  # Safety Check
  if not dest_type:
    cli.print_error(f'Failed to determine partition type for: {dev_path}')
    raise std.GenericAbort()

  # Add extra details
  if details.get('partlabel', ''):
    line += f', name="{details["partlabel"]}"'
  if details.get('partuuid', '') and source_table_type == table_type:
    # Only add UUID if source/dest table types match
    line += f', uuid={details["partuuid"].upper()}'

  # Done
  return line


def get_partition_separator(name) -> str:
  """Get partition separator based on device name, returns str."""
  separator = ''
  if re.search(r'(loop|mmc|nvme)', name, re.IGNORECASE):
    separator = 'p'

  return separator


def get_table_type(disk_path) -> str:
  """Get disk partition table type, returns str.

  NOTE: If resulting table type is not GPT or MBR
        then an exception is raised.
  """
  disk_path = str(disk_path)
  table_type = None

  # Linux
  if std.PLATFORM == 'Linux':
    cmd = f'lsblk --json --output=pttype --nodeps {disk_path}'.split()
    json_data = exe.get_json_from_command(cmd)
    table_type = json_data['blockdevices'][0].get('pttype', '').upper()
    table_type = table_type.replace('DOS', 'MBR')

  # macOS
  if std.PLATFORM == 'Darwin':
    cmd = ['diskutil', 'list', '-plist', disk_path]
    proc = exe.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
      pass
    else:
      disk_details = plist_data.get('AllDisksAndPartitions', [{}])[0]
      table_type = disk_details['Content']
      table_type = table_type.replace('FDisk_partition_scheme', 'MBR')
      table_type = table_type.replace('GUID_partition_scheme', 'GPT')

  # Check type
  if table_type not in ('GPT', 'MBR'):
    cli.print_error(f'Unsupported partition table type: {table_type}')
    raise std.GenericAbort()

  # Done
  return table_type


def prep_destination(
    state,
    source_parts: list[hw_disk.Disk],
    dry_run: bool = True,
    ) -> None:
  """Prep destination as necessary."""
  # TODO: Split into Linux and macOS
  #       logical sector size is not easily found under macOS
  #       It might be easier to rewrite this section using macOS tools
  dest_prefix = str(state.destination.path)
  dest_prefix += get_partition_separator(state.destination.path.name)
  esp_type = 'C12A7328-F81F-11D2-BA4B-00A0C93EC93B'
  msr_type = 'E3C9E316-0B5C-4DB8-817D-F92DF00215AE'
  part_num = 0
  sfdisk_script = []
  settings = state.load_settings()

  # Bail early
  if not settings['Needs Format']:
    return

  # Add partition table settings
  if settings['Table Type'] == 'GPT':
    sfdisk_script.append('label: gpt')
  else:
    sfdisk_script.append('label: dos')
  sfdisk_script.append('unit: sectors')
  sfdisk_script.append('')

  # Add boot partition if requested
  if settings['Create Boot Partition']:
    if settings['Table Type'] == 'GPT':
      part_num += 1
      sfdisk_script.append(
        build_sfdisk_partition_line(
          table_type='GPT',
          dev_path=f'{dest_prefix}{part_num}',
          size='260MiB',
          details={'parttype': esp_type, 'partlabel': 'EFI System'},
          ),
        )
      part_num += 1
      sfdisk_script.append(
        build_sfdisk_partition_line(
          table_type=settings['Table Type'],
          dev_path=f'{dest_prefix}{part_num}',
          size='16MiB',
          details={'parttype': msr_type, 'partlabel': 'Microsoft Reserved'},
          ),
        )
    elif settings['Table Type'] == 'MBR':
      part_num += 1
      sfdisk_script.append(
        build_sfdisk_partition_line(
          table_type='MBR',
          dev_path=f'{dest_prefix}{part_num}',
          size='100MiB',
          details={'parttype': '0x7', 'partlabel': 'System Reserved'},
          ),
        )

  # Add selected partition(s)
  for part in source_parts:
    num_sectors = part.size / state.destination.log_sec
    num_sectors = math.ceil(num_sectors)
    part_num += 1
    sfdisk_script.append(
      build_sfdisk_partition_line(
        table_type=settings['Table Type'],
        dev_path=f'{dest_prefix}{part_num}',
        size=num_sectors,
        details=part.raw_details,
        ),
      )

  # Save sfdisk script
  script_path = (
    f'{state.working_dir}/'
    f'sfdisk_{state.destination.path.name}.script'
    )
  with open(script_path, 'w', encoding='utf-8') as _f:
    _f.write('\n'.join(sfdisk_script))

  # Skip real format for dry runs
  if dry_run:
    LOG.info('Dry run, refusing to format destination')
    return

  # Format disk
  LOG.warning('Formatting destination: %s', state.destination.path)
  with open(script_path, 'r', encoding='utf-8') as _f:
    proc = exe.run_program(
      cmd=['sudo', 'sfdisk', state.destination.path],
      stdin=_f,
      check=False,
      )
    if proc.returncode != 0:
      cli.print_error('Error(s) encoundtered while formatting destination')
      raise std.GenericAbort()

  # Update settings
  settings['Needs Format'] = False
  state.save_settings(settings)


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