#!/usr/bin/python3
# group: quick
#
# Test cpr-exec migration mode (qemu-update) with various devices
#
# Copyright (c) 2026 Virtuozzo International GmbH.
#
# 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 logging
import os
from tempfile import NamedTemporaryFile

import iotests

disk_path = os.path.join(iotests.test_dir, 'disk')


class TestCprExec(iotests.QMPTestCase):
    def setUp(self) -> None:
        iotests.qemu_img_create('-f', iotests.imgfmt, disk_path, '64M')
        self.vm = iotests.VM()
        self.vm.add_args('-cpr-ram-rdonly')
        self.vm.add_drive(disk_path, opts='node-name=drive0-node',
                          interface='none')
        self.cprfile = None

    def tearDown(self) -> None:
        self.vm.shutdown()
        if os.path.isfile(disk_path):
            os.remove(disk_path)
        if os.path.isfile(self.cprfile):
            os.remove(self.cprfile)

    def _recreate_cprfile(self) -> None:
        if self.cprfile and os.path.isfile(self.cprfile):
            os.remove(self.cprfile)
        self.cprfile = NamedTemporaryFile(dir=iotests.test_dir,
                                          delete=False).name

    def _wait_migration_complete(self) -> None:
        with iotests.Timeout(5, 'Migration with mode=cpr-exec does not '
                             'complete in time'):
            self.vm.event_wait('RESUME')

    def _do_cpr(self) -> None:
        self._recreate_cprfile()

        cpr_exec_cmd = list(self.vm._qemu_full_args) + \
                       ['-incoming', f'file:{self.cprfile}']

        self.vm.qmp_check_log('migrate-set-parameters',
                              mode='cpr-exec', cpr_exec_command=cpr_exec_cmd)
        self.vm.qmp_check_log('migrate-set-capabilities',
                              capabilities=[{'capability': 'events',
                                             'state': True}])
        self.vm.qmp_check_log('migrate', uri=f'file:{self.cprfile}')
        self._wait_migration_complete()

    def _assert_cpu_count(self, expected: int) -> None:
        cpus = self.vm.qmp('query-cpus-fast')['return']
        self.assertEqual(len(cpus), expected)

    def _test_cpr(self, pre_cpr_hook=None, post_cpr_hook=None) -> None:
        self.vm.launch()

        if pre_cpr_hook is not None:
            pre_cpr_hook()

        self._do_cpr()
        self._do_cpr()

        if post_cpr_hook is not None:
            post_cpr_hook()

        self.vm.shutdown()

    def test_virtio_blk_hd(self) -> None:
        self.vm.add_args('-device', 'virtio-blk-pci,drive=drive0')
        self._test_cpr()

    def test_virtio_scsi_hd(self) -> None:
        self.vm.add_args('-device', 'virtio-scsi-pci')
        self.vm.add_args('-device', 'scsi-hd,drive=drive0')
        self._test_cpr()

    def test_usb_hd(self) -> None:
        self.vm.add_args('-device', 'nec-usb-xhci,id=usb')
        self.vm.add_args('-device', 'usb-storage,drive=drive0,bus=usb.0')
        self._test_cpr()

    def test_virtio_net(self) -> None:
        self.vm.add_args('-device', 'virtio-net-pci')
        self._test_cpr()

    def test_two_network_ifs(self) -> None:
        self.vm.add_args('-device', 'virtio-net-pci')
        self.vm.add_args('-device', 'e1000')
        self._test_cpr()

    def _cpr_with_tap_nic(self, model: str) -> None:
        """Attach virtual network card (non-virtio) to a real tap netdev.
        See #VSTOR-130135"""
        self.vm.add_args('-netdev', 'tap,script=no,downscript=no,id=hostnet0')
        self.vm.add_args('-device', f'{model},netdev=hostnet0')
        self._test_cpr()

    def test_tap_e1000(self) -> None:
        self._cpr_with_tap_nic('e1000')

    def test_tap_rtl8139(self) -> None:
        self._cpr_with_tap_nic('rtl8139')

    def test_tap_nontap(self) -> None:
        self.vm.add_args('-device', 'virtio-net-pci')
        self._cpr_with_tap_nic('e1000e')

    def test_tap_vhost(self) -> None:
        """A vhost=on tap netdev keeps its /dev/vhost-net fd across cpr-exec,
        but the kernel device stays owned by the mm which execve() replaces,
        so ownership has to be released before the exec.  Without that the new
        QEMU fails with 'vhost_set_owner failed: Device or resource busy'."""
        self.vm.add_args('-netdev',
                         'tap,script=no,downscript=no,id=hostnet0,vhost=on')
        self.vm.add_args('-device', 'virtio-net-pci,netdev=hostnet0')
        self._test_cpr()

    def test_tap_vhost_multiqueue(self) -> None:
        """Same, with several queues: every preserved vhost fd has to be
        handed over, not just the first one."""
        self.vm.add_args(
            '-netdev', 'tap,script=no,downscript=no,id=hostnet0,'
                       'vhost=on,queues=2',
            '-device', 'virtio-net-pci,netdev=hostnet0,mq=on,vectors=6')
        self._test_cpr()

    def test_tap_prefix_colliding_ids(self) -> None:
        """Attach 2 netdevs which share a name prefix.  Regression test.
        See #VSTOR-132222"""
        for nic_id in ('hostnet1', 'hostnet11'):
            self.vm.add_args(
                '-netdev', f'tap,script=no,downscript=no,id={nic_id}',
                '-device', f'e1000,netdev={nic_id}',
            )
        self._test_cpr()

    def test_usb_tablet(self) -> None:
        self.vm.add_args('-device', 'nec-usb-xhci,id=usb',
                         '-device', 'usb-tablet,bus=usb.0')
        self._test_cpr()

    def test_ide_cd(self) -> None:
        self.vm.add_args('-device', 'ide-cd,drive=drive0')
        self._test_cpr()

    def test_scsi_cd(self) -> None:
        self.vm.add_args('-device', 'virtio-scsi-pci')
        self.vm.add_args('-device', 'scsi-cd,drive=drive0')
        self._test_cpr()

    def test_floppy(self) -> None:
        self.vm.add_args('-device', 'floppy,drive=drive0')
        self._test_cpr()

    def test_null_block(self) -> None:
        self.vm.add_args('-blockdev', 'driver=null-co,read-zeroes=on,node-name=null')
        self.vm.add_args('-device', 'virtio-blk-pci,drive=null')
        self._test_cpr()

    def test_chardev_serial(self) -> None:
        self.vm.add_args('-chardev', 'null,id=char0',
                         '-device', 'isa-serial,chardev=char0')
        self._test_cpr()

    def test_chardev_srv_unix(self) -> None:
        chr_sock = iotests.file_path('chr_sock')
        self.vm.add_args('-chardev',
                         f'socket,id=char0,path={chr_sock},server=on,wait=off')
        self._test_cpr()

    def _test_chardev_change(self, from_type: str, to_type: str) -> None:
        """Hot-swap a serial chardev backend via chardev-change, then migrate.
        A pty target re-saves its master fd under the (label, 0) cpr key while
        the predecessor still owns it; regression test.  See #VSTOR-135364"""
        def change_backend():
            # chardev-change to a pty replies with the new pty path, so check
            # for success rather than asserting an empty {'return': {}} reply.
            res = self.vm.qmp('chardev-change', id='char0',
                              backend={'type': to_type, 'data': {}})
            assert 'return' in res, res
            # The new QEMU on cpr-exec must launch with the post-change backend
            args = list(self.vm._qemu_full_args)
            args[args.index(f'{from_type},id=char0')] = f'{to_type},id=char0'
            self.vm._qemu_full_args = tuple(args)

        self.vm.add_args('-chardev', f'{from_type},id=char0',
                         '-device', 'isa-serial,chardev=char0')
        self._test_cpr(pre_cpr_hook=change_backend)

    def test_chardev_change_pty_to_pty(self) -> None:
        self._test_chardev_change('pty', 'pty')

    def test_chardev_change_null_to_pty(self) -> None:
        self._test_chardev_change('null', 'pty')

    def test_chardev_change_pty_to_null(self) -> None:
        self._test_chardev_change('pty', 'null')

    def test_chardev_change_pty_symlink_rollback(self) -> None:
        """A chardev-change whose char_pty_open() fails at symlink() must roll
        back without dropping the surviving chardev's cpr fd.  The cpr-exec
        migration afterwards only succeeds if that fd was restored.  See
        #VSTOR-135364"""
        link = os.path.join(iotests.test_dir, 'pty-rollback.link')
        # Pre-create the target path so the new pty's symlink() fails (EEXIST).
        open(link, 'w').close()

        def failed_change():
            res = self.vm.qmp('chardev-change', id='char0',
                              backend={'type': 'pty', 'data': {'path': link}})
            assert 'error' in res, res
            # The change rolled back; the original char0 (no path) stays
            # active, so _qemu_full_args needs no adjustment.

        try:
            self.vm.add_args('-chardev', 'pty,id=char0',
                             '-device', 'isa-serial,chardev=char0')
            self._test_cpr(pre_cpr_hook=failed_change)
        finally:
            os.remove(link)

    def test_chardev_change_file_to_file(self) -> None:
        """chardev-change between file backends shares the (label, 0) cpr key
        like pty does.  Regression test.  See #VSTOR-135364"""
        out1 = os.path.join(iotests.test_dir, 'chr-file1.log')
        out2 = os.path.join(iotests.test_dir, 'chr-file2.log')

        def change_file():
            res = self.vm.qmp('chardev-change', id='char0',
                              backend={'type': 'file', 'data': {'out': out2}})
            assert 'return' in res, res
            args = list(self.vm._qemu_full_args)
            args[args.index(f'file,id=char0,path={out1}')] = \
                f'file,id=char0,path={out2}'
            self.vm._qemu_full_args = tuple(args)

        try:
            self.vm.add_args('-chardev', f'file,id=char0,path={out1}',
                             '-device', 'isa-serial,chardev=char0')
            self._test_cpr(pre_cpr_hook=change_file)
        finally:
            os.remove(out1)
            os.remove(out2)

    def test_chardev_change_socket_to_socket(self) -> None:
        """chardev-change between (server) socket backends re-saves the
        listen-socket fd while the predecessor still owns the cpr key.
        Regression test.  See #VSTOR-135364"""
        sock1 = iotests.file_path('chr-sock1')
        sock2 = iotests.file_path('chr-sock2')

        def change_socket():
            res = self.vm.qmp('chardev-change', id='char0',
                              backend={'type': 'socket', 'data': {
                                  'addr': {'type': 'unix',
                                           'data': {'path': sock2}},
                                  'server': True, 'wait': False}})
            assert 'return' in res, res
            args = list(self.vm._qemu_full_args)
            args[args.index(
                f'socket,id=char0,path={sock1},server=on,wait=off')] = \
                f'socket,id=char0,path={sock2},server=on,wait=off'
            self.vm._qemu_full_args = tuple(args)

        self.vm.add_args(
            '-chardev', f'socket,id=char0,path={sock1},server=on,wait=off',
            '-device', 'isa-serial,chardev=char0')
        self._test_cpr(pre_cpr_hook=change_socket)

    def test_nbd_attached_hd(self) -> None:
        nbd_img = os.path.join(iotests.test_dir, 'nbd_disk')
        iotests.qemu_img_create('-f', 'qcow2', nbd_img, '64M')

        nbd_sock = iotests.file_path('nbd.sock', base_dir=iotests.sock_dir)
        with iotests.change_log_level('qemu.iotests.diff_io', level=logging.WARNING), \
             iotests.qemu_nbd_popen('--persistent', '--socket', nbd_sock,
                                    '--export-name=default', '-f', 'qcow2', nbd_img):
            self.vm.add_args('-blockdev', 'driver=nbd,server.type=unix,' +
                             f'server.path={nbd_sock},export=default,node-name=nbd0')
            self.vm.add_args('-device', 'virtio-blk-pci,drive=nbd0')
            self._test_cpr()

        if os.path.isfile(nbd_img):
            os.remove(nbd_img)
        if os.path.exists(str(nbd_sock)):
            os.remove(str(nbd_sock))

    def test_live_external_snapshot(self) -> None:
        snapshot_path = os.path.join(iotests.test_dir, 'snapshot.qcow2')
        def create_snapshot():
            self.vm.qmp_check_log('blockdev-snapshot-sync',
                                  node_name='drive0-node',
                                  snapshot_file=snapshot_path, format='qcow2',
                                  snapshot_node_name='drive0-snap-node')
        self._test_cpr(pre_cpr_hook=create_snapshot)
        if os.path.isfile(snapshot_path):
            os.remove(snapshot_path)

    def test_vnc_server(self) -> None:
        self.vm.add_args('-vnc', ':1')
        self._test_cpr()

    def test_vcpu_hotplug(self) -> None:
        """Hotplug a vCPU and migrate with the command line left alone: the
        new QEMU starts with a single boot CPU, so it has to re-create the
        hotplugged one from the CPR state before the migration stream is
        loaded.  See #VSTOR-127496"""
        def hotplug_vcpu() -> None:
            cpus = self.vm.qmp('query-hotpluggable-cpus')['return']
            assert len(cpus) == 2
            cpu = next(c for c in cpus if 'qom-path' not in c)
            self.vm.qmp_check_log('device_add', driver=cpu['type'],
                            id='cpu1', **cpu['props'])
            self._assert_cpu_count(2)

        self.vm.add_args('-smp', '1,maxcpus=2')
        self._test_cpr(pre_cpr_hook=hotplug_vcpu,
                       post_cpr_hook=lambda: self._assert_cpu_count(2))

    def test_vcpu_hotplug_multiple(self) -> None:
        """Same as test_vcpu_hotplug, but every hotpluggable slot is filled,
        so each saved CPU has to be replayed on the destination.
        See #VSTOR-127496"""
        def hotplug_vcpus() -> None:
            cpus = self.vm.qmp('query-hotpluggable-cpus')['return']
            free = [c for c in cpus if 'qom-path' not in c]
            assert len(free) == 3
            for i, cpu in enumerate(free, start=1):
                self.vm.qmp_check_log('device_add', driver=cpu['type'],
                                      id=f'cpu{i}', **cpu['props'])
            self._assert_cpu_count(4)

        self.vm.add_args('-smp', '1,maxcpus=4')
        self._test_cpr(pre_cpr_hook=hotplug_vcpus,
                       post_cpr_hook=lambda: self._assert_cpu_count(4))


    def test_memory_hotplug(self) -> None:
        def hotplug_dimm() -> None:
            self.vm.qmp_check_log('object-add', qom_type='memory-backend-ram',
                                  id='mem1', size=1024 * 1024 * 1024)
            self.vm.qmp_check_log('device_add', driver='pc-dimm',
                                  id='dimm1', memdev='mem1')

            # Add the DIMM to _qemu_full_args so the new QEMU on cpr-exec
            # has the device present when loading migration state
            args = list(self.vm._qemu_full_args)
            args += ['-object', 'memory-backend-ram,id=mem1,size=1G',
                     '-device', 'pc-dimm,id=dimm1,memdev=mem1']
            self.vm._qemu_full_args = tuple(args)

        self.vm.add_args('-m', '1G,slots=2,maxmem=2G',
                         '-object', 'memory-backend-ram,id=ram0,size=1G',
                         '-numa', 'node,nodeid=0,memdev=ram0')
        self._test_cpr(pre_cpr_hook=hotplug_dimm)


if __name__ == '__main__':
    iotests.activate_logging()
    iotests.main(supported_fmts=["qcow2"],
                 supported_protocols=["file"])
