#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Oct 23 17:34:02 2017

@author: triley
"""


import errno
import os
import re
# import sys
# from subprocess import check_call, run
import subprocess
from pprint import pprint
from collections import OrderedDict  # not really needed for CPython >= 3.6 ?
from typing import Dict, List


class DiskPart(object):
    """  isolate hardware from everything else, abstract partitions
    """

    def __init__(
            self,
            dev='',
            uuid='',
            ptype='',
            puuid='',
            mount_point=''):

        """

        :rtype: DiskPart
        """

        self.dev = dev
        self.disk = dev[5:8]
        self.part_num = dev[8]
        self.uuid = uuid
        self.ptype = ptype
        self.puuid = puuid
        self.mount_point = mount_point

    def __str__(self) -> str:
        s = ['  H']
        s += [self.dev if self.dev else ' '*9]
        s += [self.ptype[:4] if self.ptype else ' '*4]
        s += [self.uuid[:8] if self.uuid else ' '*8]
        s += [self.mount_point if self.mount_point else 'None']
        rstr = str(' '.join(s))
        return rstr

#    def __repr__(self):        return_str = str('\n'.join(s))
#        self.__str__()  # cheat


class PartitionTable(object):
    """ Abstract Partition Table
    """

    def __init__(self) -> None:
        """ Extract info from the OS:

        """

        # one to one mapping {/dev/sd??:  partition}
        self.partitions: OrderedDict[str, DiskPart] = {}  # by /dev/sd??
        self.update_blkids()
        for p in self.partitions:  # sanity check
            assert '/dev/' in p, f'partition table init on partition {p}'

        # potentially many to one: 1 partition can be mounted multiple times
        self.mounts: OrderedDict[str, DiskPart] = {}  # by mount_point
        self.update_mounts()
        assert '/' in self.mounts
        assert '/boot/efi' in self.mounts

        self.sources: Dict[str, DiskPart] = {}  # by mount_point
        self.dests: OrderedDict[str, DiskPart] = {}  # by mount_point

    def __str__(self):
        s = ['\nPartition Table:']
        for name, part in self.partitions.items():
            s += [part.__str__()]
        return str('\n'.join(s))

#    def __repr__(self):
#        self.__str__()

    def add_mock_partition(self):
        pass

    def update_blkids(self) -> None:
        """ Read partitions from os using blkid.

            This seems to work fine for non-root user, provided that you poke
            it first with a 'sudo blkid' at the command line once. Hmm ...
        """
        re_uuid = re.compile(b'UUID="(.*?)"')
        re_type = re.compile(b'TYPE="(.*?)"')
        re_puuid = re.compile(b'PARTUUID="(.*?)"')

        result = subprocess.run('blkid', stdout=subprocess.PIPE)
        result.check_returncode()
        lines = result.stdout.split(b'\n')
        for line in lines:
            if b'/dev/' in line:  # ignore empty lines
                p = line.split(b':')
                dev = p[0].strip().decode("utf-8")  # python3 string, no blanks
                uuid_m = re_uuid.search(p[1])  # re.match objects
                ptype_m = re_type.search(p[1])
                puuid_m = re_puuid.search(p[1])
                self.partitions[dev] = DiskPart(
                    dev=dev,  # keep dev info in the DiskPart instance
                    uuid=uuid_m.group(1).decode("utf-8") if uuid_m else None,
                    ptype=ptype_m.group(1).decode("utf-8") if ptype_m else None,
                    puuid=puuid_m.group(1).decode("utf-8") if puuid_m else None)

    def update_mounts(self) -> None:
        """ filter mounted for disks """
        with open('/proc/mounts') as f:
            lines = f.readlines()
        for line in lines:  # things we might want to add
            if '/dev/' in line[:5]:  # then it's a disk partition
                line = line.split()
                if line[0] in self.partitions:  # ignore stale usb devices
                    self.mounts[line[1]] = self.partitions[line[0]]
                    self.partitions[line[0]].mount_point = line[1]
        removes = []  # don't delete while iterating
        for m in self.mounts:  # things we might want to be remove
            if m not in [line.split()[1] for line in lines]:
                removes += [m]
        for m in removes:
            del self.mounts[m]

    def check_mount_point(self, mount_point: str) -> None:
        # print('cmp: mount_point =', mount_point)
        if not os.path.isdir(mount_point):
            try:
                os.makedirs(mount_point)
            except OSError as e:
                if e.errno == errno.EEXIST:  # seems impossible
                    pass
                elif e.errno == errno.EACCES:
                    raise Exception('Must be run as root user')
                elif e.errno == 2:  # TODO look up non-numeric
                    # because we got crap for mount_point ?
                    raise Exception(
                            f'{mount_point} not found making mount_point')
                else:
                    pprint(e)
                    raise Exception(f'wtf checking mount_point {mount_point}')
            # print(f'created mount point {mount_point}')
            return
        else:
            pass

    def _mount_partition(self, part: str, mount_point: str) -> None:
        """ /bin/mount requires root priveledge
        """

        self.check_mount_point(mount_point)

        doit = None   # mounted allready ?
        cmd = ['/bin/mount', part, mount_point]
        msg = f'wtf mounting partition {part} at {mount_point}'
        if mount_point in self.mounts:
            if self.mounts[mount_point] is not self.partitions[part]:
                doit = cmd   # remount
        elif part not in [p.dev for p in self.mounts.values()]:
            doit = cmd
        else:
            raise Exception(msg)
        if doit:
            error_ret = subprocess.check_call(doit)
            if error_ret:
                raise Exception(msg)
            self.mounts[mount_point] = self.partitions[part]
            self.partitions[part].mount_point = mount_point

    def mock_mount(self, part: str, mount_point: str) -> None:

        raise Exception(f'not implemented yet {mount_point}')

    def mount_partition(self, part: str, mount_point: str) -> None:
        # print('mp:', part, mount_point)
        if os.getuid() == 0:
            self._mount_partition(part, mount_point)
        else:
            self.mock_mount(part, mount_point)

    def find_source_disk(self) -> None:
        """ Find out where to copy from.

            Includes the non-universal assumption that both partitions are on
            the same disk.
        """

        for mount in ['/', '/boot/efi']:
            self.sources[mount] = self.mounts[mount]

        assert len(self.sources) == 2, 'wtf: Failed to find sources.'

    def find_dest_disk(self) -> None:

        all_parts: List[DiskPart] = [p for i, p in self.partitions.items()]

        # Find available vfat partitions for EFI.
        efi_candidates: List[DiskPart] = []
        for partition in all_parts:
            # print('find_dest_disk:', partition.dev, partition.mount_point)
            if partition.ptype == 'vfat' and (
                    partition.mount_point != '/boot/efi'):
                efi_candidates += [partition]

        # The new '/' is to be put on the same disk as the new EFI.
        install_candidates: List[Dict[str, DiskPart]] = []  # str is by mount
        for efi_part in efi_candidates:
            for p in all_parts:  # should be same disk as efi partition
                if p.ptype != 'vfat' and efi_part.disk in p.dev:
                    efi_mount = f'/mnt/{p.disk}{p.part_num}/boot/efi'
                    root_mount = f'/mnt/{p.disk}{p.part_num}'
                    install_candidate = OrderedDict({
                        root_mount: p,  # mount root first, then EFI
                        efi_mount: efi_part})
                    install_candidates += [install_candidate]
        # print('\ninstall_candidates\n', install_candidates)
        if len(install_candidates) == 1:
            self.dests = install_candidates[0]
        elif len(install_candidates) == 0:
            raise Exception('No install candidates found')
        elif len(install_candidates) > 1:
            raise Exception('TODO fix later, choose from root_candidates')
        else:
            raise Exception('wtf in find_dest_disk')

    def mount_dest_disk(self):
        for m, p in self.dests.items():
            # print('mdd: mount', p.dev, m)
            if m not in self.mounts:
                self.mount_partition(self.dests[m].dev, m)


if __name__ == '__main__':
    # quick test for the most obvious case (for me right now)

    partition_table = PartitionTable()
    print('\n', partition_table)
    partition_table.find_source_disk()
    print('\nsources')
    pprint(partition_table.sources)

    partition_table.find_dest_disk()
    print('\ndests')
    pprint(partition_table.dests)

    partition_table.mount_dest_disk()
    print(partition_table)
