# Copyright (C) 2020-2024 Soleta Networks <opengnsys@soleta.eu>
#
# This program is free software: you can redistribute it and/or modify it under
# the terms of the GNU Affero General Public License as published by the
# Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.

import argparse
import re

from cli.utils import print_json


class OgDisk():

    @staticmethod
    def list_disks(rest, args):
        parser = argparse.ArgumentParser(prog='ogcli list disks')
        parser.add_argument('--client-ip',
                            nargs='?',
                            required=True,
                            help='Client IP to query')
        parsed_args = parser.parse_args(args)
        payload = {'client': [parsed_args.client_ip]}

        r = rest.get('/client/setup', payload=payload)
        print_json(r.text)

    @staticmethod
    def setup_disk(rest, args):
        def parse_size(size):
            size = size.upper()
            units = {"M": 1024, "G": 1024**2, "T": 1024**3}  # Mapped to K
            # size = re.sub(r'(\d+)([MGT])', r'\1 \2', size)
            match = re.match(r'(\d+)([MGT])', size)
            if match:
                if len(match.groups()) == 2:
                    number, unit = match.groups()
                    return str(int(float(number)*units[unit]))
            raise ValueError(f'error parsing size {size}. Examples of valid sizes: 200M, 2G, 1T')

        disk_type_map = {'dos': 'MSDOS', 'gpt': 'GPT'}
        part_types = ['LINUX', 'EFI', 'NTFS', 'CACHE']
        fs_types = ['EXT4', 'FAT32', 'NTFS', 'CACHE']

        parser = argparse.ArgumentParser(prog='ogcli setup disk')
        parser.add_argument('--type',
                            nargs='?',
                            required=True,
                            choices=['dos', 'gpt'],
                            help='Disk partition scheme')
        parser.add_argument('--num',
                            nargs='?',
                            default=1,
                            help='Disk number (defaults to 1)')
        parser.add_argument('--part',
                            nargs='+',
                            action='append',
                            type=lambda x: x.split(','),
                            required=True,
                            help='Partition definition (syntax: "num,part_scheme,fs,size")')
        group = parser.add_argument_group('clients', 'Client selection args')
        group.add_argument('--center-id',
                           type=int,
                           action='append',
                           default=[],
                           required=False,
                           help='Clients from given center id')
        group.add_argument('--room-id',
                           type=int,
                           action='append',
                           default=[],
                           required=False,
                           help='Clients from given room id')
        group.add_argument('--client-ip',
                           action='append',
                           default=[],
                           required=False,
                           help='Specific client IP')

        parsed_args = parser.parse_args(args)
        r = rest.get('/scopes')
        scopes = r.json()
        ips = set()

        for center in parsed_args.center_id:
            center_scope = scope_lookup(center, 'center', scopes)
            ips.update(ips_in_scope(center_scope))
        for room in parsed_args.room_id:
            room_scope = scope_lookup(room, 'room', scopes)
            ips.update(ips_in_scope(room_scope))
        for l in parsed_args.client_ip:
            ips.add(l)
        if not ips:
            print("Missing --client-ip, or --room-id/--center-id. No clients provided.")
            return None

        if not parsed_args.num.isdigit():
            print(f'Invalid disk number: must be an integer value')
            return

        payload = {'clients': parsed_args.client_ip, 'type': disk_type_map[parsed_args.type], 'disk': str(parsed_args.num),
                   'cache': '0', 'cache_size': '0', 'partition_setup': []}

        defined_part_indices = set()
        for i, p in enumerate(parsed_args.part, start=1):

            p = p[0]

            if len(p) != 4:
                print(f'Invalid partition: requires "num,part_scheme,fs,size", "{",".join(p)}" provided')
                return
            part_num, code, fs, size = p[0], p[1].upper(), p[2].upper(), p[3]

            if not part_num.isdigit():
                print(f'Invalid partition: the first parameter must be a number, "{part_num}" provided')
                return

            if part_num in defined_part_indices:
                print(f'Invalid partition: partition number "{part_num}" has multiple definitions')
                return

            defined_part_indices.add(part_num)

            if code not in part_types:
                print(f'Invalid partition {i}: specified partition type {code} is not supported. The supported formats are {part_types}')
                return

            if fs not in fs_types:
                print(f'Invalid partition {i}: specified filesystem {fs} is not supported. The supported formats are {fs_types}')
                return
            try:
                size = parse_size(size)
            except ValueError as error:
                print(f'Invalid partition {i}: {str(error)}')
                return

            for j in range(i, int(part_num)):
                part = {'partition': str(j), 'code': 'EMPTY',
                        'filesystem': 'EMPTY', 'size': '0',
                        'format': '0'}
                payload['partition_setup'].append(part)

            if fs == 'CACHE':
                # Assuming flag specifying if there's cache in the setup
                payload['cache'] = '1'
                payload['cache_size'] = size
            part = {'partition': str(p[0]), 'code': code.upper(),
                    'filesystem': fs.upper(), 'size': size,
                    'format': '0'}
            payload['partition_setup'].append(part)

        last_partnum = int(parsed_args.part[-1][0][0])
        # Pad with empty partitions if no 4th part was defined
        for i in range(last_partnum + 1, 5):
            part = {'partition': str(i), 'code': 'EMPTY',
                    'filesystem': 'EMPTY', 'size': '0',
                    'format': '0'}
            payload['partition_setup'].append(part)

        rest.post('/setup', payload=payload)