#!/usr/bin/env python3
#
# Script to compare machine type compatible properties (include/hw/core/boards.h).
# compat_props are applied to the driver during initialization to change
# default values, for instance, to maintain compatibility.
# This script constructs table with machines and values of their compat_props
# to compare and to find places for improvements or places with bugs. If
# during the comparison, some machine type doesn't have a property (it is in
# the comparison table because another machine type has it), then the
# appropriate method will be used to obtain the default value of this driver
# property via qmp command (e.g. query-cpu-model-expansion for x86_64-cpu).
# These methods are defined below in qemu_property_methods.
#
# Copyright (c) Yandex Technologies LLC, 2023
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, see <http://www.gnu.org/licenses/>.

import sys
from argparse import ArgumentParser, RawTextHelpFormatter, Namespace
import pandas as pd
from contextlib import ExitStack
from typing import Optional, List, Dict, Generator, Tuple, Union, Any, Set

try:
    from qemu.machine import QEMUMachine
except ModuleNotFoundError as exc:
    print(f"Module '{exc.name}' not found.", file=sys.stderr)
    print(f"Try $builddir/run {' '.join(sys.argv)}", file=sys.stderr)
    sys.exit(1)


default_qemu_args = '-enable-kvm -machine none'
default_qemu_binary = 'build/qemu-system-x86_64'


# Methods for gettig the right values of drivers properties
#
# Use these methods as a 'whitelist' and add entries only if necessary. It's
# important to be stable and predictable in analysis and tests.
# Be careful:
# * Class must be inherited from 'QEMUObject' and used in new_driver()
# * Class has to implement get_prop method in order to get values
# * Specialization always wins (with the given classes for 'device' and
#   'x86_64-cpu', method of 'x86_64-cpu' will be used for '486-x86_64-cpu')

class Driver():
    def __init__(self, vm: QEMUMachine, name: str, abstract: bool) -> None:
        self.vm = vm
        self.name = name
        self.abstract = abstract
        self.parent: Optional[Driver] = None
        self.property_getter: Optional[Driver] = None

    def get_prop(self, driver: str, prop: str) -> str:
        if self.property_getter:
            return self.property_getter.get_prop(driver, prop)
        else:
            return 'Unavailable method'

    def is_child_of(self, parent: 'Driver') -> bool:
        """Checks whether self is (recursive) child of @parent"""
        cur_parent = self.parent
        while cur_parent:
            if cur_parent is parent:
                return True
            cur_parent = cur_parent.parent

        return False

    def set_implementations(self, implementations: List['Driver']) -> None:
        self.implementations = implementations


class QEMUObject(Driver):
    def __init__(self, vm: QEMUMachine, name: str) -> None:
        super().__init__(vm, name, True)

    def set_implementations(self, implementations: List[Driver]) -> None:
        self.implementations = implementations

        # each implementation of the abstract driver has to use property getter
        # of this abstract driver unless it has specialization. (e.g. having
        # 'device' and 'x86_64-cpu', property getter of 'x86_64-cpu' will be
        # used for '486-x86_64-cpu')
        for impl in implementations:
            if not impl.property_getter or\
                    self.is_child_of(impl.property_getter):
                impl.property_getter = self


class QEMUDevice(QEMUObject):
    def __init__(self, vm: QEMUMachine) -> None:
        super().__init__(vm, 'device')
        self.cached: Dict[str, List[Dict[str, Any]]] = {}

    def get_prop(self, driver: str, prop_name: str) -> str:
        if driver not in self.cached:
            self.cached[driver] = self.vm.cmd('device-list-properties',
                                              typename=driver)
        for prop in self.cached[driver]:
            if prop['name'] == prop_name:
                return str(prop.get('default-value', 'No default value'))

        return 'Unknown property'


class QEMUx86CPU(QEMUObject):
    def __init__(self, vm: QEMUMachine) -> None:
        super().__init__(vm, 'x86_64-cpu')
        self.cached: Dict[str, Dict[str, Any]] = {}

    def get_prop(self, driver: str, prop_name: str) -> str:
        if not driver.endswith('-x86_64-cpu'):
            return 'Wrong x86_64-cpu name'

        # crop last 11 chars '-x86_64-cpu'
        name = driver[:-11]
        if name not in self.cached:
            self.cached[name] = self.vm.cmd(
                'query-cpu-model-expansion', type='full',
                model={'name': name})['model']['props']
        return str(self.cached[name].get(prop_name, 'Unknown property'))


# Now it's stub, because all memory_backend types don't have default values
# but this behaviour can be changed
class QEMUMemoryBackend(QEMUObject):
    def __init__(self, vm: QEMUMachine) -> None:
        super().__init__(vm, 'memory-backend')
        self.cached: Dict[str, List[Dict[str, Any]]] = {}

    def get_prop(self, driver: str, prop_name: str) -> str:
        if driver not in self.cached:
            self.cached[driver] = self.vm.cmd('qom-list-properties',
                                              typename=driver)
        for prop in self.cached[driver]:
            if prop['name'] == prop_name:
                return str(prop.get('default-value', 'No default value'))

        return 'Unknown property'


def new_driver(vm: QEMUMachine, name: str, is_abstr: bool) -> Driver:
    if name == 'object':
        return QEMUObject(vm, 'object')
    elif name == 'device':
        return QEMUDevice(vm)
    elif name == 'x86_64-cpu':
        return QEMUx86CPU(vm)
    elif name == 'memory-backend':
        return QEMUMemoryBackend(vm)
    else:
        return Driver(vm, name, is_abstr)
# End of methods definition


class VMPropertyGetter:
    """It implements the relationship between drivers and how to get their
    properties"""
    def __init__(self, vm: QEMUMachine) -> None:
        self.drivers: Dict[str, Driver] = {}

        qom_all_types = vm.cmd('qom-list-types', abstract=True)
        self.drivers = {t['name']: new_driver(vm, t['name'],
                                              t.get('abstract', False))
                        for t in qom_all_types}

        for t in qom_all_types:
            drv = self.drivers[t['name']]
            if 'parent' in t:
                drv.parent = self.drivers[t['parent']]

        for drv in self.drivers.values():
            imps = vm.cmd('qom-list-types', implements=drv.name)
            # only implementations inherit property getter
            drv.set_implementations([self.drivers[imp['name']]
                                     for imp in imps])

    def get_prop(self, driver: str, prop: str) -> str:
        # wrong driver name or disabled in config driver
        try:
            drv = self.drivers[driver]
        except KeyError:
            return 'Unavailable driver'

        assert not drv.abstract

        return drv.get_prop(driver, prop)

    def get_implementations(self, driver: str) -> List[str]:
        return [impl.name for impl in self.drivers[driver].implementations]


class Machine:
    """A short QEMU machine type description. It contains only processed
    compat_props (properties of abstract classes are applied to its
    implementations)
    """
    # raw_mt_dict - dict produced by `query-machines`
    def __init__(self, raw_mt_dict: Dict[str, Any],
                 qemu_drivers: VMPropertyGetter) -> None:
        self.name = raw_mt_dict['name']
        self.compat_props: Dict[str, Any] = {}
        # properties are applied sequentially and can rewrite values like in
        # QEMU. Also it has to resolve class relationships to apply appropriate
        # values from abstract class to all implementations
        for prop in raw_mt_dict['compat-props']:
            driver = prop['qom-type']
            try:
                # implementation adds only itself, abstract class adds
                #  lementation (abstract classes are uninterestiong)
                impls = qemu_drivers.get_implementations(driver)
                for impl in impls:
                    if impl not in self.compat_props:
                        self.compat_props[impl] = {}
                    self.compat_props[impl][prop['property']] = prop['value']
            except KeyError:
                # QEMU doesn't know this driver thus it has to be saved
                if driver not in self.compat_props:
                    self.compat_props[driver] = {}
                self.compat_props[driver][prop['property']] = prop['value']


class Configuration():
    """Class contains all necessary components to generate table and is used
    to compare different binaries"""
    def __init__(self, vm: QEMUMachine,
                 req_mt: List[str], all_mt: bool) -> None:
        self._vm = vm
        self._binary = vm.binary
        self._qemu_args = args.qemu_args.split(' ')

        self._qemu_drivers = VMPropertyGetter(vm)
        self.req_mt = get_req_mt(self._qemu_drivers, vm, req_mt, all_mt)

    def get_implementations(self, driver_name: str) -> List[str]:
        return self._qemu_drivers.get_implementations(driver_name)

    def get_table(self, req_props: List[Tuple[str, str]]) -> pd.DataFrame:
        table: List[pd.DataFrame] = []
        for mt in self.req_mt:
            name = f'{self._binary}\n{mt.name}'
            column = []
            for driver, prop in req_props:
                try:
                    # values from QEMU machine type definitions
                    column.append(mt.compat_props[driver][prop])
                except KeyError:
                    # values from QEMU type definitions
                    column.append(self._qemu_drivers.get_prop(driver, prop))
            table.append(pd.DataFrame({name: column}))

        return pd.concat(table, axis=1)


script_desc = """Script to compare machine types (their compat_props).

Examples:
* save info about all machines:  ./scripts/compare-machine-types.py --all \
--format csv --raw > table.csv
* compare machines: ./scripts/compare-machine-types.py --mt pc-q35-2.12 \
pc-q35-3.0
* compare binaries and machines: ./scripts/compare-machine-types.py \
--mt pc-q35-6.2 pc-q35-7.0 --qemu-binary build/qemu-system-x86_64 \
build/qemu-exp
  ╒════════════╤══════════════════════════╤════════════════════════════\
╤════════════════════════════╤══════════════════╤══════════════════╕
  │   Driver   │         Property         │  build/qemu-system-x86_64  \
│  build/qemu-system-x86_64  │  build/qemu-exp  │  build/qemu-exp  │
  │            │                          │         pc-q35-6.2         \
│         pc-q35-7.0         │    pc-q35-6.2    │    pc-q35-7.0    │
  ╞════════════╪══════════════════════════╪════════════════════════════\
╪════════════════════════════╪══════════════════╪══════════════════╡
  │  PIIX4_PM  │ x-not-migrate-acpi-index │            True            \
│           False            │      False       │      False       │
  ├────────────┼──────────────────────────┼────────────────────────────\
┼────────────────────────────┼──────────────────┼──────────────────┤
  │ virtio-mem │  unplugged-inaccessible  │           False            \
│            auto            │      False       │       auto       │
  ╘════════════╧══════════════════════════╧════════════════════════════\
╧════════════════════════════╧══════════════════╧══════════════════╛

If a property from QEMU machine defintion applies to an abstract class (e.g. \
x86_64-cpu) this script will compare all implementations of this class.

"Unavailable method" - means that this script doesn't know how to get \
default values of the driver. To add method use the construction described \
at the top of the script.
"Unavailable driver" - means that this script doesn't know this driver. \
For instance, this can happen if you configure QEMU without this device or \
if machine type definition has error.
"No default value" - means that the appropriate method can't get the default \
value and most likely that this property doesn't have it.
"Unknown property" - means that the appropriate method can't find property \
with this name."""


def parse_args() -> Namespace:
    parser = ArgumentParser(formatter_class=RawTextHelpFormatter,
                            description=script_desc)
    parser.add_argument('--format', choices=['human-readable', 'json', 'csv'],
                        default='human-readable',
                        help='returns table in json format')
    parser.add_argument('--raw', action='store_true',
                        help='prints ALL defined properties without value '
                             'transformation. By default, only rows '
                             'with different values will be printed and '
                             'values will be transformed(e.g. "on" -> True)')
    parser.add_argument('--qemu-args', default=default_qemu_args,
                        help='command line to start qemu. '
                             f'Default: "{default_qemu_args}"')
    parser.add_argument('--qemu-binary', nargs="*", type=str,
                        default=[default_qemu_binary],
                        help='list of qemu binaries that will be compared. '
                             f'Deafult: {default_qemu_binary}')

    mt_args_group = parser.add_mutually_exclusive_group()
    mt_args_group.add_argument('--all', action='store_true',
                               help='prints all available machine types (list '
                                    'of machine types will be ignored)')
    mt_args_group.add_argument('--mt', nargs="*", type=str,
                               help='list of Machine Types '
                                    'that will be compared')

    return parser.parse_args()


def mt_comp(mt: Machine) -> Tuple[str, int, int, int]:
    """Function to compare and sort machine by names.
    It returns socket_name, major version, minor version, revision"""
    # none, microvm, x-remote and etc.
    if '-' not in mt.name or '.' not in mt.name:
        return mt.name, 0, 0, 0

    socket, ver = mt.name.rsplit('-', 1)
    ver_list = list(map(int, ver.split('.', 2)))
    ver_list += [0] * (3 - len(ver_list))
    return socket, ver_list[0], ver_list[1], ver_list[2]


def get_mt_definitions(qemu_drivers: VMPropertyGetter,
                       vm: QEMUMachine) -> List[Machine]:
    """Constructs list of machine definitions (primarily compat_props) via
    info from QEMU"""
    raw_mt_defs = vm.cmd('query-machines', compat_props=True)
    mt_defs = []
    for raw_mt in raw_mt_defs:
        mt_defs.append(Machine(raw_mt, qemu_drivers))

    mt_defs.sort(key=mt_comp)
    return mt_defs


def get_req_mt(qemu_drivers: VMPropertyGetter, vm: QEMUMachine,
               req_mt: Optional[List[str]], all_mt: bool) -> List[Machine]:
    """Returns list of requested by user machines"""
    mt_defs = get_mt_definitions(qemu_drivers, vm)
    if all_mt:
        return mt_defs

    if req_mt is None:
        print('Enter machine types for comparision')
        exit(0)

    matched_mt = []
    for mt in mt_defs:
        if mt.name in req_mt:
            matched_mt.append(mt)

    return matched_mt


def get_affected_props(configs: List[Configuration]) -> Generator[Tuple[str,
                                                                        str],
                                                                  None, None]:
    """Helps to go through all affected in machine definitions drivers
    and properties"""
    driver_props: Dict[str, Set[Any]] = {}
    for config in configs:
        for mt in config.req_mt:
            compat_props = mt.compat_props
            for driver, prop in compat_props.items():
                if driver not in driver_props:
                    driver_props[driver] = set()
                driver_props[driver].update(prop.keys())

    for driver, props in sorted(driver_props.items()):
        for prop in sorted(props):
            yield driver, prop


def transform_value(value: str) -> Union[str, bool]:
    true_list = ['true', 'on']
    false_list = ['false', 'off']

    out = value.lower()

    if out in true_list:
        return True

    if out in false_list:
        return False

    return value


def simplify_table(table: pd.DataFrame) -> pd.DataFrame:
    """transforms values to make it easier to compare it and drops rows
    with the same values for all columns"""

    table = table.map(transform_value)

    return table[~table.iloc[:, 3:].eq(table.iloc[:, 2], axis=0).all(axis=1)]


# constructs table in the format:
#
# Driver  | Property  | binary1  | binary1  | ...
#         |           | machine1 | machine2 | ...
# ------------------------------------------------------ ...
# driver1 | property1 |  value1  |  value2  | ...
# driver1 | property2 |  value3  |  value4  | ...
# driver2 | property3 |  value5  |  value6  | ...
#   ...   |    ...    |   ...    |   ...    | ...
#
def fill_prop_table(configs: List[Configuration],
                    is_raw: bool) -> pd.DataFrame:
    req_props = list(get_affected_props(configs))
    if not req_props:
        print('No drivers to compare. Check machine names')
        exit(0)

    driver_col, prop_col = tuple(zip(*req_props))
    table = [pd.DataFrame({'Driver': driver_col}),
             pd.DataFrame({'Property': prop_col})]

    table.extend([config.get_table(req_props) for config in configs])

    df_table = pd.concat(table, axis=1)

    if is_raw:
        return df_table

    return simplify_table(df_table)


def print_table(table: pd.DataFrame, table_format: str) -> None:
    if table_format == 'json':
        print(comp_table.to_json())
    elif table_format == 'csv':
        print(comp_table.to_csv())
    else:
        print(comp_table.to_markdown(index=False, stralign='center',
                                     colalign=('center',), headers='keys',
                                     tablefmt='fancy_grid',
                                     disable_numparse=True))


if __name__ == '__main__':
    args = parse_args()
    with ExitStack() as stack:
        vms = [stack.enter_context(QEMUMachine(binary=binary, qmp_timer=15,
               args=args.qemu_args.split(' '))) for binary in args.qemu_binary]

        configurations = []
        for vm in vms:
            vm.launch()
            configurations.append(Configuration(vm, args.mt, args.all))

        comp_table = fill_prop_table(configurations, args.raw)
        if not comp_table.empty:
            print_table(comp_table, args.format)