"""Unit tests for the zpool-list and lsblk parsers in ssh_client.

These cover the structural cases that drive the pool-membership lock:
mirror/raidz/draid container vdevs, single-disk vdevs at depth 1, the
flattened-indentation behaviour of `zpool list -vHP` on TrueNAS, partition
suffix stripping for NVMe and SCSI, and the cache/log/spare/special
section markers (including plural variants).

Run with:  python -m unittest discover tests/ -v
"""

import unittest

from app.ssh_client import (
    _parse_zpool_list_output,
    _parse_lsblk_zfs_output,
    _parse_smart_health_batch,
)


class TestParseZpoolList(unittest.TestCase):

    def test_empty_output_returns_empty(self):
        self.assertEqual(_parse_zpool_list_output(""), {})

    def test_single_pool_with_mirror(self):
        # TrueNAS-flattened output: pool at depth 0, vdev type and devices
        # all at depth 1.
        out = _parse_zpool_list_output(
            "boot-pool\t232G\t8.4G\t224G\t-\t-\t17%\t3%\t1.00x\tONLINE\t-\n"
            "\tmirror-0\t232G\t8.4G\t224G\t-\t-\t17%\t3.6%\t-\tONLINE\n"
            "\t/dev/nvme0n1p3\t232G\t-\t-\t-\t-\t-\t-\t-\tONLINE\n"
            "\t/dev/sdd3\t232G\t-\t-\t-\t-\t-\t-\t-\tONLINE\n"
        )
        self.assertEqual(out, {
            "nvme0n1": {"pool": "boot-pool", "role": "data"},
            "sdd":     {"pool": "boot-pool", "role": "data"},
        })

    def test_raidz2_pool(self):
        out = _parse_zpool_list_output(
            "tank\t127T\t4.5T\t122T\t-\t-\t0%\t3%\t1.00x\tONLINE\t-\n"
            "\traidz2-0\t127T\t4.5T\t122T\t-\t-\t0%\t3%\t-\tONLINE\n"
            "\t/dev/sdc\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\n"
            "\t/dev/sde\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\n"
            "\t/dev/sdf\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\n"
        )
        self.assertEqual(set(out.keys()), {"sdc", "sde", "sdf"})
        for v in out.values():
            self.assertEqual(v, {"pool": "tank", "role": "data"})

    def test_draid_pool(self):
        out = _parse_zpool_list_output(
            "warm\t100T\t-\t-\t-\t-\t-\t-\t-\tONLINE\t-\n"
            "\tdraid2:8d:10c:1s-0\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\n"
            "\t/dev/sdg\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\n"
            "\t/dev/sdh\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\n"
        )
        self.assertEqual(out["sdg"], {"pool": "warm", "role": "data"})
        self.assertEqual(out["sdh"], {"pool": "warm", "role": "data"})

    def test_single_disk_vdev_at_depth_1(self):
        # No mirror/raidz wrapper — a `/dev/...` line itself sits at depth 1.
        out = _parse_zpool_list_output(
            "scratch\t1T\t-\t-\t-\t-\t-\t-\t-\tONLINE\t-\n"
            "\t/dev/sdi\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\n"
        )
        self.assertEqual(out, {"sdi": {"pool": "scratch", "role": "data"}})

    def test_section_markers_switch_role(self):
        # cache / log / spare / special / dedup all at depth 1; subsequent
        # /dev/... lines (also at depth 1) inherit that role.
        out = _parse_zpool_list_output(
            "tank\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\t-\n"
            "\tmirror-0\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\n"
            "\t/dev/sda\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\n"
            "\t/dev/sdb\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\n"
            "\tcache\n"
            "\t/dev/nvme1n1\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\n"
            "\tlog\n"
            "\t/dev/nvme2n1\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\n"
            "\tspare\n"
            "\t/dev/sdz\t-\t-\t-\t-\t-\t-\t-\t-\tAVAIL\n"
        )
        self.assertEqual(out["sda"],     {"pool": "tank", "role": "data"})
        self.assertEqual(out["sdb"],     {"pool": "tank", "role": "data"})
        self.assertEqual(out["nvme1n1"], {"pool": "tank", "role": "cache"})
        self.assertEqual(out["nvme2n1"], {"pool": "tank", "role": "log"})
        self.assertEqual(out["sdz"],     {"pool": "tank", "role": "spare"})

    def test_section_markers_plurals_normalize(self):
        # ZFS sometimes emits 'logs'/'spares' instead of 'log'/'spare'.
        out = _parse_zpool_list_output(
            "tank\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\t-\n"
            "\tlogs\n"
            "\t/dev/nvme0n1\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\n"
            "\tspares\n"
            "\t/dev/sdz\t-\t-\t-\t-\t-\t-\t-\t-\tAVAIL\n"
        )
        self.assertEqual(out["nvme0n1"]["role"], "log")
        self.assertEqual(out["sdz"]["role"],     "spare")

    def test_special_and_dedup_section(self):
        out = _parse_zpool_list_output(
            "tank\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\t-\n"
            "\tspecial\n"
            "\t/dev/sda\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\n"
            "\tdedup\n"
            "\t/dev/sdb\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\n"
        )
        self.assertEqual(out["sda"]["role"], "special")
        self.assertEqual(out["sdb"]["role"], "dedup")

    def test_partition_suffix_stripped(self):
        out = _parse_zpool_list_output(
            "tank\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\t-\n"
            "\tmirror-0\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\n"
            "\t/dev/sda3\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\n"
            "\t/dev/nvme0n1p3\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\n"
        )
        self.assertIn("sda", out)
        self.assertNotIn("sda3", out)
        self.assertIn("nvme0n1", out)
        self.assertNotIn("nvme0n1p3", out)

    def test_long_scsi_devname(self):
        # Past sdz: sdaa, sdab, ...
        out = _parse_zpool_list_output(
            "big\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\t-\n"
            "\traidz3-0\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\n"
            "\t/dev/sdaa\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\n"
            "\t/dev/sdab1\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\n"
        )
        self.assertEqual(out["sdaa"]["pool"], "big")
        self.assertEqual(out["sdab"]["pool"], "big")  # partition stripped

    def test_pool_name_with_dashes_dots_underscores(self):
        out = _parse_zpool_list_output(
            "my-cool_pool.v2\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\t-\n"
            "\t/dev/sda\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\n"
        )
        self.assertEqual(out["sda"]["pool"], "my-cool_pool.v2")

    def test_multiple_pools(self):
        out = _parse_zpool_list_output(
            "boot-pool\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\t-\n"
            "\tmirror-0\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\n"
            "\t/dev/nvme0n1p3\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\n"
            "\t/dev/sdd3\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\n"
            "tank\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\t-\n"
            "\traidz2-0\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\n"
            "\t/dev/sda\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\n"
            "\t/dev/sdb\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\n"
        )
        self.assertEqual(out["nvme0n1"]["pool"], "boot-pool")
        self.assertEqual(out["sdd"]["pool"],     "boot-pool")
        self.assertEqual(out["sda"]["pool"],     "tank")
        self.assertEqual(out["sdb"]["pool"],     "tank")

    def test_pool_role_resets_between_pools(self):
        # Section marker in pool A must not carry into pool B.
        out = _parse_zpool_list_output(
            "a\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\t-\n"
            "\tcache\n"
            "\t/dev/sda\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\n"
            "b\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\t-\n"
            "\t/dev/sdb\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\n"
        )
        self.assertEqual(out["sda"]["role"], "cache")
        self.assertEqual(out["sdb"]["role"], "data")

    def test_blank_lines_skipped(self):
        out = _parse_zpool_list_output(
            "\n"
            "tank\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\t-\n"
            "\n"
            "\t/dev/sda\t-\t-\t-\t-\t-\t-\t-\t-\tONLINE\n"
        )
        self.assertEqual(out, {"sda": {"pool": "tank", "role": "data"}})


class TestParseLsblkZfs(unittest.TestCase):

    def test_empty_returns_empty_set(self):
        self.assertEqual(_parse_lsblk_zfs_output(""), set())

    def test_partition_zfs_member(self):
        # Typical TrueNAS layout: zpool members are partitions.
        out = _parse_lsblk_zfs_output(
            "sda      \n"
            "sda1     \n"
            "sda3     zfs_member\n"
            "sdb      \n"
            "sdb3     zfs_member\n"
        )
        self.assertEqual(out, {"sda", "sdb"})

    def test_whole_disk_zfs_member(self):
        # Some configurations put zfs_member on the whole disk.
        out = _parse_lsblk_zfs_output(
            "sdc      zfs_member\n"
        )
        self.assertEqual(out, {"sdc"})

    def test_nvme_partitioned_and_whole(self):
        out = _parse_lsblk_zfs_output(
            "nvme0n1     \n"
            "nvme0n1p3   zfs_member\n"
            "nvme1n1     zfs_member\n"
        )
        self.assertEqual(out, {"nvme0n1", "nvme1n1"})

    def test_non_zfs_fstypes_ignored(self):
        out = _parse_lsblk_zfs_output(
            "sda1     ext4\n"
            "sda2     swap\n"
            "sdb1     btrfs\n"
        )
        self.assertEqual(out, set())

    def test_long_scsi_devnames(self):
        out = _parse_lsblk_zfs_output(
            "sdaa     zfs_member\n"
            "sdab1    zfs_member\n"
        )
        self.assertEqual(out, {"sdaa", "sdab"})

    def test_short_lines_skipped(self):
        out = _parse_lsblk_zfs_output(
            "sda\n"
            "\n"
            "sdb1     zfs_member\n"
        )
        self.assertEqual(out, {"sdb"})


class TestParseSmartHealthBatch(unittest.TestCase):

    def test_passed_drive(self):
        out = _parse_smart_health_batch(
            "@@sda@@\n"
            "smartctl 7.4 2023-08-01 r5530 [x86_64-linux-6.6]\n"
            "SMART overall-health self-assessment test result: PASSED\n"
            "@@END@@\n"
        )
        self.assertEqual(out, {"sda": "PASSED"})

    def test_failed_drive(self):
        out = _parse_smart_health_batch(
            "@@sdb@@\n"
            "SMART overall-health self-assessment test result: FAILED!\n"
            "@@END@@\n"
        )
        self.assertEqual(out, {"sdb": "FAILED"})

    def test_unknown_when_no_marker(self):
        out = _parse_smart_health_batch(
            "@@sdc@@\n"
            "/dev/sdc: Unknown USB bridge\n"
            "@@END@@\n"
        )
        self.assertEqual(out, {"sdc": "UNKNOWN"})

    def test_multiple_drives_mixed_states(self):
        out = _parse_smart_health_batch(
            "@@sda@@\n"
            "SMART overall-health self-assessment test result: PASSED\n"
            "@@END@@\n"
            "@@sdb@@\n"
            "SMART overall-health self-assessment test result: FAILED!\n"
            "@@END@@\n"
            "@@nvme0n1@@\n"
            "SMART overall-health self-assessment test result: PASSED\n"
            "@@END@@\n"
        )
        self.assertEqual(out, {"sda": "PASSED", "sdb": "FAILED", "nvme0n1": "PASSED"})

    def test_empty_returns_empty(self):
        self.assertEqual(_parse_smart_health_batch(""), {})


if __name__ == "__main__":
    unittest.main()