#!/usr/bin/env python3
# group: rw auto quick
#
# Test NBD transmission commands against a qemu NBD export
#
# Copyright (C) 2026 Virtuozzo International GmbH
#
# SPDX-License-Identifier: GPL-2.0-or-later

import os
from types import ModuleType

import iotests
from iotests import qemu_img_create, qemu_img_map, qemu_io


base = os.path.join(iotests.test_dir, 'base')
top = os.path.join(iotests.test_dir, 'top')
# Larger than the maximum payload size an export can advertise
size = 64 * 1024 * 1024
pattern = 0xa5
nbd_sock = os.path.join(iotests.sock_dir, 'nbd_sock')
nbd_uri = 'nbd+unix:///exp?socket=' + nbd_sock
nbd: ModuleType

DEPTH_LOCAL = 1
DEPTH_BACKING = 2
CLUSTER = 65536


class TestNbdCommands(iotests.QMPTestCase):
    def setUp(self):
        qemu_img_create('-f', iotests.imgfmt, base, str(size))
        qemu_io('-c', f'write -P {pattern} 0 {size}', base)
        qemu_img_create('-f', iotests.imgfmt, '-b', base,
                        '-F', iotests.imgfmt, top, str(size))

        self.vm = iotests.VM()
        self.vm.launch()
        self.vm.cmd('blockdev-add', {
            'driver': iotests.imgfmt,
            'node-name': 'n',
            'file': {'driver': 'file', 'filename': top},
            'backing': {
                'driver': iotests.imgfmt,
                'node-name': 'base',
                'file': {'driver': 'file', 'filename': base},
            },
        })
        self.vm.cmd('nbd-server-start', {
            'addr': {'type': 'unix', 'data': {'path': nbd_sock}}
        })
        self.vm.cmd('block-export-add', {
            'type': 'nbd',
            'id': 'exp',
            'node-name': 'n',
            'name': 'exp',
            'writable': True,
            'allocation-depth': True,
        })

        self.h = None
        self.connect()

    def tearDown(self):
        self.disconnect()
        self.vm.shutdown()
        for f in (top, base, nbd_sock):
            try:
                os.remove(f)
            except OSError:
                pass

    def connect(self, structured=True, extended=True):
        self.disconnect()
        h = nbd.NBD()
        h.set_request_structured_replies(structured)
        h.set_request_extended_headers(extended)
        h.add_meta_context('base:allocation')
        h.add_meta_context('qemu:allocation-depth')
        # Let the server, not libnbd, reject the out of range requests below
        h.set_strict_mode(h.get_strict_mode() &
                          ~(nbd.STRICT_BOUNDS | nbd.STRICT_PAYLOAD))
        h.connect_uri(nbd_uri)
        self.assertEqual(h.get_structured_replies_negotiated(), structured)
        self.assertEqual(h.get_extended_headers_negotiated(), extended)
        self.h = h

    def disconnect(self):
        if self.h is not None:
            self.h.shutdown()
            self.h = None

    def block_status(self, count=size, wanted=None, flags=0):
        """Map each meta context in the reply to its list of extents."""
        reply = {}

        def cb(meta, _offset, entries, _err):
            reply.setdefault(meta, []).extend(zip(entries[0::2],
                                                  entries[1::2]))

        if wanted is None:
            self.h.block_status(count, 0, cb, flags)
        else:
            self.h.block_status_filter(count, 0, wanted, cb, flags)
        return reply

    def top_extents(self):
        """Which parts of the top image are local, once qemu has let go."""
        self.disconnect()
        self.vm.shutdown()
        return [(e['start'], e['length'], e['depth'])
                for e in qemu_img_map(top)]

    def test_cache_copies_on_read(self):
        maximum = self.h.get_block_size(nbd.SIZE_MAXIMUM)
        self.assertLess(maximum, size)
        self.assertEqual(self.block_status()['qemu:allocation-depth'],
                         [(size, DEPTH_BACKING)])

        self.h.cache(size, 0)

        self.assertEqual(self.top_extents(), [(0, size, 0)])
        qemu_io('-c', f'read -P {pattern} 0 {size}', top)

    def check_commands(self, structured, extended):
        self.connect(structured, extended)

        self.assertEqual(self.h.pread(4096, 4096), bytes([pattern]) * 4096)

        self.h.cache(size, 0)

        self.h.pwrite(b'x' * 4096, 4096)
        self.h.flush()
        self.assertEqual(self.h.pread(4096, 4096), b'x' * 4096)

        self.h.zero(4096, 8192)
        self.assertEqual(self.h.pread(4096, 8192), bytes(4096))

        self.h.trim(4096, 16384)

        if structured:
            self.assertEqual(self.block_status()['qemu:allocation-depth'],
                             [(size, DEPTH_LOCAL)])

    def test_commands_simple_replies(self):
        self.check_commands(structured=False, extended=False)

    def test_commands_structured_replies(self):
        self.check_commands(structured=True, extended=False)

    def test_commands_extended_headers(self):
        self.check_commands(structured=True, extended=True)

    def test_block_status_payload_filter(self):
        self.assertEqual(sorted(self.block_status()),
                         ['base:allocation', 'qemu:allocation-depth'])

        for wanted in (['base:allocation'], ['qemu:allocation-depth']):
            self.assertEqual(sorted(self.block_status(wanted=wanted)), wanted)

    def read_chunks(self, count, offset, flags=0):
        chunks = []

        def cb(subbuf, off, status, _err):
            chunks.append((off, len(subbuf), status))

        self.h.pread_structured(count, offset, cb, flags)
        return chunks

    def make_sparse(self):
        """Lay out data, a hole and data again, one cluster each."""
        self.h.pwrite(b'z' * CLUSTER, 0)
        self.h.zero(CLUSTER, CLUSTER)
        self.h.pwrite(b'z' * CLUSTER, 2 * CLUSTER)

    def test_read_sparse_chunks(self):
        self.make_sparse()

        first = (0, CLUSTER, nbd.READ_DATA)
        hole = (CLUSTER, CLUSTER, nbd.READ_HOLE)
        second = (2 * CLUSTER, CLUSTER, nbd.READ_DATA)

        # A hole is an offset and a length, not a cluster of zeroes,
        # wherever it falls in the reply
        self.assertEqual(self.read_chunks(2 * CLUSTER, 0), [first, hole])
        self.assertEqual(self.read_chunks(2 * CLUSTER, CLUSTER),
                         [hole, second])
        self.assertEqual(self.read_chunks(3 * CLUSTER, 0),
                         [first, hole, second])

        self.assertEqual(self.h.pread(CLUSTER, CLUSTER), bytes(CLUSTER))
        self.assertEqual(self.h.pread(CLUSTER, 0), b'z' * CLUSTER)

    def test_read_dont_fragment(self):
        self.make_sparse()

        self.assertEqual(self.read_chunks(3 * CLUSTER, 0, nbd.CMD_FLAG_DF),
                         [(0, 3 * CLUSTER, nbd.READ_DATA)])

    def test_block_status_req_one(self):
        hole = nbd.STATE_HOLE | nbd.STATE_ZERO
        self.make_sparse()

        alloc = self.block_status(3 * CLUSTER)['base:allocation']
        self.assertEqual(alloc, [(CLUSTER, 0), (CLUSTER, hole), (CLUSTER, 0)])

        # One extent, so the reply covers less than was asked for and the
        # client has to come back for the rest
        alloc = self.block_status(3 * CLUSTER, flags=nbd.CMD_FLAG_REQ_ONE)
        self.assertEqual(alloc['base:allocation'], [(CLUSTER, 0)])

    def test_write_and_trim_fua(self):
        self.h.pwrite(b'y' * 4096, 4096, nbd.CMD_FLAG_FUA)
        self.assertEqual(self.h.pread(4096, 4096), b'y' * 4096)
        self.h.trim(4096, 4096, nbd.CMD_FLAG_FUA)

    def test_fast_zero(self):
        self.h.zero(CLUSTER, CLUSTER, nbd.CMD_FLAG_FAST_ZERO)
        self.assertEqual(self.h.pread(CLUSTER, CLUSTER), bytes(CLUSTER))

        # Zeroing part of a cluster needs the fallback the flag forbids
        with self.assertRaises(nbd.Error) as caught:
            self.h.zero(4096, 4096, nbd.CMD_FLAG_FAST_ZERO)
        self.assertEqual(caught.exception.errno, 'ENOTSUP')

    def test_cache_past_end_of_export(self):
        self.assertRaises(nbd.Error, self.h.cache, size + 1, 0)

    def test_read_bound_by_max_payload(self):
        maximum = self.h.get_block_size(nbd.SIZE_MAXIMUM)
        self.assertRaises(nbd.Error, self.h.pread, maximum + 65536, 0)


if __name__ == '__main__':
    try:
        # Easier to use libnbd than to try and set up parallel
        # 'qemu-nbd --list' or 'qemu-io' processes, but not all systems
        # have libnbd installed.
        import nbd  # type: ignore

        iotests.main(supported_fmts=['qcow2'])
    except ImportError:
        iotests.notrun('Python bindings to libnbd are not installed')
