## @file
# Unit tests for the Firmware Module Management Tool.
#
# Copyright (c) Microsoft Corporation.
#
# SPDX-License-Identifier: BSD-2-Clause-Patent
##

import logging
import os
import sys
import tempfile
import unittest
import uuid
from pathlib import Path
from types import SimpleNamespace
from unittest import mock

import TestTools


_TESTS_DIR = Path(__file__).resolve().parent
_PYTHON_SOURCE = _TESTS_DIR.parent / "Source" / "Python"
_FMMT_SOURCE = _PYTHON_SOURCE / "FMMT"
sys.path.insert(0, str(_FMMT_SOURCE))
sys.path.insert(0, str(_PYTHON_SOURCE))
os.environ.setdefault("FmmtConfPath", "")

# FMMT opens a build log when imported. Unit tests do not need a persistent log.
with mock.patch("logging.FileHandler", return_value=logging.NullHandler()):
    import core.FMMTOperation as FMMTOperation
    from core.BiosTree import BIOSTREE
    from core.BiosTree import FFS_FREE_SPACE
    from core.BiosTree import FFS_TREE
    from core.BiosTree import SEC_FV_TREE
    from core.BiosTreeNode import FreeSpaceNode
    from core.FvHandler import FvHandler
    from FirmwareStorageFormat.Common import ModifyGuidFormat
    from FirmwareStorageFormat.FfsFileHeader import EFI_FFS_FILE_HEADER


class _FirmwareVolumeData:
    def __init__(self, size, free_space=0, block_size=0x40):
        self.Data = b""
        self.Free_Space = free_space
        self.Header = SimpleNamespace(
            Attributes=0,
            BlockMap=[SimpleNamespace(Length=block_size)],
            ExtHeaderOffset=0,
            FileSystemGuid=ModifyGuidFormat(
                "00000000-0000-0000-0000-000000000000"
            ),
            FvLength=size,
        )
        self.Size = size

    def ModCheckSum(self):
        pass

    def ModExtHeaderData(self):
        pass

    def ModFvExt(self):
        pass

    def ModFvSize(self):
        pass


class _TestFvHandler(FvHandler):
    def CompressData(self, target_tree):
        self.Status = True

    def ModifyTest(self, parent_tree, needed_space):
        self.Status = True


class _TestRecursiveFvHandler(FvHandler):
    def CompressData(self, target_tree):
        self.Status = True


def _create_ffs(name, size, data=b""):
    node = BIOSTREE(name)
    node.type = FFS_TREE
    header = EFI_FFS_FILE_HEADER()
    header.Name = ModifyGuidFormat(str(name))
    header.Size[0] = size & 0xFF
    header.Size[1] = (size >> 8) & 0xFF
    header.Size[2] = (size >> 16) & 0xFF
    node.Data = SimpleNamespace(
        Data=data,
        Header=header,
        Name=name,
        PadData=b"",
        Size=size,
    )
    return node


class FmmtReplaceTests(unittest.TestCase):
    def test_replace_multiple_matches_uses_independent_ffs_trees(self):
        target_guid = uuid.UUID("11111111-2222-3333-4444-555555555555")
        targets = [SimpleNamespace(), SimpleNamespace()]
        replacement = SimpleNamespace(
            Data=SimpleNamespace(PadData=b"", Size=0x20),
            Parent=None,
        )

        class _Root:
            def __init__(self):
                self.Findlist = []

            def FindNode(self, name, matches):
                matches.extend(targets)

        class _Parser:
            def __init__(self, root):
                self.FinalData = b"output"
                self.WholeFvTree = root

            def Encapsulation(self, root, compress_status):
                pass

            def ParserFromRoot(self, root, data):
                pass

        input_parser = _Parser(_Root())
        replacement_parser = _Parser(SimpleNamespace(Child=[replacement]))
        replacements = []

        class _Handler:
            def __init__(self, new_ffs, target_ffs):
                replacements.append(new_ffs)
                new_ffs.Parent = target_ffs

            def ReplaceFfs(self):
                return True

        with tempfile.TemporaryDirectory() as directory:
            input_path = Path(directory) / "input.fd"
            new_ffs_path = Path(directory) / "new.ffs"
            output_path = Path(directory) / "output.fd"
            input_path.write_bytes(b"input")
            new_ffs_path.write_bytes(b"replacement")

            with mock.patch.object(
                FMMTOperation,
                "FMMTParser",
                side_effect=[input_parser, replacement_parser],
            ), mock.patch.object(FMMTOperation, "FvHandler", _Handler):
                FMMTOperation.ReplaceFfs(
                    str(input_path),
                    target_guid,
                    str(new_ffs_path),
                    str(output_path),
                )

        self.assertEqual(2, len(replacements))
        self.assertIsNot(replacements[0], replacements[1])
        self.assertIs(targets[0], replacements[0].Parent)
        self.assertIs(targets[1], replacements[1].Parent)

    def test_smaller_replacement_creates_free_space_node(self):
        prefix = _create_ffs(
            uuid.UUID("10000000-0000-0000-0000-000000000001"), 0x18
        )
        target = _create_ffs(
            uuid.UUID("20000000-0000-0000-0000-000000000002"), 0x40
        )
        replacement = _create_ffs(target.Data.Name, 0x18)
        volume = BIOSTREE("FV")
        volume.type = SEC_FV_TREE
        volume.Data = _FirmwareVolumeData(0x80)
        volume.insertChild(prefix)
        volume.insertChild(target)

        self.assertTrue(_TestFvHandler(replacement, target).ReplaceFfs())

        self.assertIs(replacement, volume.Child[1])
        free_space = volume.Child[-1]
        self.assertEqual(FFS_FREE_SPACE, free_space.type)
        self.assertIsInstance(free_space.Data, FreeSpaceNode)
        self.assertEqual(b"\xFF" * 0x28, free_space.Data.Data)

    def test_growing_nested_fv_preserves_final_ffs(self):
        prefix = _create_ffs(
            uuid.UUID("10000000-0000-0000-0000-000000000001"), 0x18, b"prefix"
        )
        target = _create_ffs(
            uuid.UUID("20000000-0000-0000-0000-000000000002"), 0x20, b"target"
        )
        trailing = _create_ffs(
            uuid.UUID("30000000-0000-0000-0000-000000000003"),
            0x18,
            b"trailing",
        )
        replacement = _create_ffs(target.Data.Name, 0x30, b"replacement")
        volume = BIOSTREE("FV")
        volume.type = SEC_FV_TREE
        volume.Data = _FirmwareVolumeData(0x80)
        volume.insertChild(prefix)
        volume.insertChild(target)
        volume.insertChild(trailing)

        self.assertTrue(_TestFvHandler(replacement, target).ReplaceFfs())

        self.assertEqual(b"trailing", trailing.Data.Data)
        self.assertIn(trailing, volume.Child)
        self.assertIs(replacement, volume.Child[1])
        free_space = volume.Child[-1]
        self.assertEqual(FFS_FREE_SPACE, free_space.type)
        self.assertIsInstance(free_space.Data, FreeSpaceNode)
        self.assertEqual(b"\xFF" * 0x30, free_space.Data.Data)

    def test_block_aligned_growth_of_only_ffs_does_not_crash(self):
        target = _create_ffs(
            uuid.UUID("20000000-0000-0000-0000-000000000002"), 0x20, b"target"
        )
        replacement = _create_ffs(target.Data.Name, 0x60, b"replacement")
        volume = BIOSTREE("FV")
        volume.type = SEC_FV_TREE
        volume.Data = _FirmwareVolumeData(0x40)
        volume.insertChild(target)

        self.assertTrue(_TestFvHandler(replacement, target).ReplaceFfs())

        self.assertEqual([replacement], volume.Child)
        self.assertEqual(0, volume.Data.Free_Space)

    def test_block_aligned_growth_preserves_ffs_order(self):
        prefix = _create_ffs(
            uuid.UUID("10000000-0000-0000-0000-000000000001"), 0x18, b"prefix"
        )
        target = _create_ffs(
            uuid.UUID("20000000-0000-0000-0000-000000000002"), 0x20, b"target"
        )
        trailing = _create_ffs(
            uuid.UUID("30000000-0000-0000-0000-000000000003"),
            0x18,
            b"trailing",
        )
        replacement = _create_ffs(target.Data.Name, 0x60, b"replacement")
        volume = BIOSTREE("FV")
        volume.type = SEC_FV_TREE
        volume.Data = _FirmwareVolumeData(0x80)
        volume.insertChild(prefix)
        volume.insertChild(target)
        volume.insertChild(trailing)

        self.assertTrue(_TestFvHandler(replacement, target).ReplaceFfs())

        self.assertEqual([prefix, replacement, trailing], volume.Child)
        self.assertEqual(0, volume.Data.Free_Space)

    def test_growing_ancestor_fv_preserves_final_ffs(self):
        prefix = _create_ffs(
            uuid.UUID("10000000-0000-0000-0000-000000000001"), 0x18, b"prefix"
        )
        trailing = _create_ffs(
            uuid.UUID("30000000-0000-0000-0000-000000000003"),
            0x18,
            b"trailing",
        )
        volume = BIOSTREE("FV")
        volume.type = SEC_FV_TREE
        volume.Data = _FirmwareVolumeData(0x80)
        volume.insertChild(prefix)
        volume.insertChild(trailing)

        parent = BIOSTREE("ROOT")
        parent.type = FMMTOperation.ROOT_TREE
        parent.insertChild(volume)

        _TestRecursiveFvHandler(None).ModifyTest(volume, 0x10)

        self.assertEqual(b"trailing", trailing.Data.Data)
        self.assertIn(trailing, volume.Child)
        free_space = volume.Child[-1]
        self.assertEqual(FFS_FREE_SPACE, free_space.type)
        self.assertIsInstance(free_space.Data, FreeSpaceNode)
        self.assertEqual(b"\xFF" * 0x30, free_space.Data.Data)


TheTestSuite = TestTools.MakeTheTestSuite(locals())


if __name__ == "__main__":
    unittest.TextTestRunner(verbosity=2).run(TheTestSuite())