"""Tests for the functional GMMU walker (sim/gmmu_walk.py, docs/48).

The golden vector is a captured channel PDB (numeric entries); the synthetic
trees are built with the module's own encoders and must round-trip.  No
hardware, no modules, no guest, no network.

Run::

    python3 -m unittest tests.test_gmmu_walk -v
"""
from __future__ import annotations

import contextlib
import io
import json
import os
import struct
import sys
import unittest

HERE = os.path.dirname(os.path.abspath(__file__))
TOOL_ROOT = os.path.abspath(os.path.join(HERE, ".."))
if TOOL_ROOT not in sys.path:
    sys.path.insert(0, TOOL_ROOT)

from sim.gmmu_walk import (  # noqa: E402
    FAULT_INVALID_APERTURE,
    FAULT_PTE_INVALID,
    FAULT_TABLE_MISSING,
    FAULT_UNSUPPORTED,
    FAULT_OUT_OF_RANGE,
    WALK_FORMAT,
    MemoryView,
    WalkError,
    channel_pdb_fixture,
    encode_dual_slot,
    encode_pde_pointer,
    encode_pte,
    main,
    walk,
)


def _tree(root_pa: int, va: int, target_pa: int | None, *, base: int = 0x50000000,
          read_only: bool = False) -> MemoryView:
    i0 = (va >> 47) & 0x3
    i1 = (va >> 38) & 0x1FF
    i2 = (va >> 29) & 0x1FF
    i3 = (va >> 21) & 0xFF
    i4 = (va >> 12) & 0x1FF
    pd2, pd1, pd0, pt = base + 0x1000, base + 0x2000, base + 0x3000, base + 0x4000
    mem = MemoryView()
    page = bytearray(0x1000)
    struct.pack_into("<Q", page, i0 * 8, encode_pde_pointer(pd2))
    mem.add(root_pa, bytes(page))
    page = bytearray(0x1000)
    struct.pack_into("<Q", page, i1 * 8, encode_pde_pointer(pd1))
    mem.add(pd2, bytes(page))
    page = bytearray(0x1000)
    struct.pack_into("<Q", page, i2 * 8, encode_pde_pointer(pd0))
    mem.add(pd1, bytes(page))
    page = bytearray(0x1000)
    struct.pack_into("<16s", page, i3 * 16, encode_dual_slot(None, pt))
    mem.add(pd0, bytes(page))
    page = bytearray(0x1000)
    if target_pa is not None:
        struct.pack_into("<Q", page, i4 * 8, encode_pte(target_pa, read_only=read_only))
    mem.add(pt, bytes(page))
    return mem


def _run_main(argv):
    out, err = io.StringIO(), io.StringIO()
    with contextlib.redirect_stdout(out), contextlib.redirect_stderr(err):
        rc = main(argv)
    return rc, out.getvalue(), err.getvalue()


class GoldenVectorTests(unittest.TestCase):
    def test_captured_channel_pdb_walks_to_logged_dmas(self):
        root_pa, mem = channel_pdb_fixture()
        expected = {
            0x100000: 0x2045D9000,
            0x101000: 0x121BD9000,
            0x102000: 0x152115000,
            0x103000: 0x15A5F5000,
            0x104000: 0x10CF27000,
        }
        for va, pa in expected.items():
            report = walk(va, root_pa=root_pa, mem=mem)
            self.assertEqual(report["status"], "ok", report)
            self.assertEqual(report["format"], WALK_FORMAT)
            self.assertEqual(int(report["pa"], 16), pa)
            self.assertEqual(report["page_size"], "4k")
            self.assertTrue(report["writable"])
            self.assertEqual(report["evidence"], "validated")
            self.assertEqual(len(report["steps"]), 5)

    def test_walk_is_deterministic(self):
        root_pa, mem = channel_pdb_fixture()
        first = json.dumps(walk(0x100000, root_pa=root_pa, mem=mem), sort_keys=True)
        second = json.dumps(walk(0x100000, root_pa=root_pa, mem=mem), sort_keys=True)
        self.assertEqual(first, second)


class SyntheticRoundTripTests(unittest.TestCase):
    def test_encoder_round_trip_4k(self):
        va, pa = 0x204608000, 0x15E8FA000
        mem = _tree(0x60000000, va, pa)
        report = walk(va, root_pa=0x60000000, mem=mem)
        self.assertEqual(report["status"], "ok")
        self.assertEqual(int(report["pa"], 16), pa)
        self.assertEqual(report["page_size"], "4k")
        self.assertTrue(report["writable"])

    def test_read_only_leaf_is_reported(self):
        va, pa = 0x204608000, 0x15E8FA000
        mem = _tree(0x60000000, va, pa, read_only=True)
        report = walk(va, root_pa=0x60000000, mem=mem)
        self.assertEqual(report["status"], "ok")
        self.assertTrue(report["read_only"])
        self.assertFalse(report["writable"])

    def test_2m_leaf(self):
        va, base_pa = 0x40000000, 0x20000000
        mem = MemoryView()
        i0 = (va >> 47) & 0x3
        i1 = (va >> 38) & 0x1FF
        i2 = (va >> 29) & 0x1FF
        i3 = (va >> 21) & 0xFF
        pd2, pd1, pd0 = 0x51001000, 0x51002000, 0x51003000
        page = bytearray(0x1000)
        struct.pack_into("<Q", page, i0 * 8, encode_pde_pointer(pd2))
        mem.add(0x51000000, bytes(page))
        page = bytearray(0x1000)
        struct.pack_into("<Q", page, i1 * 8, encode_pde_pointer(pd1))
        mem.add(pd2, bytes(page))
        page = bytearray(0x1000)
        struct.pack_into("<Q", page, i2 * 8, encode_pde_pointer(pd0))
        mem.add(pd1, bytes(page))
        page = bytearray(0x1000)
        struct.pack_into("<Q", page, i3 * 16, encode_pte(base_pa, page_size="2m"))
        mem.add(pd0, bytes(page))
        report = walk(va + 0x12345, root_pa=0x51000000, mem=mem)
        self.assertEqual(report["status"], "ok")
        self.assertEqual(report["page_size"], "2m")
        self.assertEqual(int(report["pa"], 16), base_pa + 0x12345)

    def test_auto_page_mode_prefers_prefers_small_when_big_invalid(self):
        va = 0x100000  # 64K-aligned; the dual BIG half is invalid (0)
        root_pa, mem = channel_pdb_fixture()
        report = walk(va, root_pa=root_pa, mem=mem, page_mode="auto")
        self.assertEqual(report["status"], "ok")
        self.assertEqual(report["page_mode_used"], "4k")

    def test_auto_page_mode_uses_big_when_valid(self):
        va, pa = 0x100000, 0x30000000
        # build a tree whose dual BIG half points at a 64K-granularity table
        i0 = (va >> 47) & 0x3
        i1 = (va >> 38) & 0x1FF
        i2 = (va >> 29) & 0x1FF
        i3 = (va >> 21) & 0xFF
        pd2, pd1, pd0, pt = 0x52001000, 0x52002000, 0x52003000, 0x52004000
        mem = MemoryView()
        page = bytearray(0x1000)
        struct.pack_into("<Q", page, i0 * 8, encode_pde_pointer(pd2))
        mem.add(0x52000000, bytes(page))
        page = bytearray(0x1000)
        struct.pack_into("<Q", page, i1 * 8, encode_pde_pointer(pd1))
        mem.add(pd2, bytes(page))
        page = bytearray(0x1000)
        struct.pack_into("<Q", page, i2 * 8, encode_pde_pointer(pd0))
        mem.add(pd1, bytes(page))
        page = bytearray(0x1000)
        struct.pack_into("<16s", page, i3 * 16, encode_dual_slot(pt, pt))
        mem.add(pd0, bytes(page))
        page = bytearray(0x1000)
        # BIG path indexes the 64K table by VA bits 20:16
        struct.pack_into("<Q", page, ((va >> 16) & 0x1FF) * 8,
                         encode_pte(pa, page_size="64k"))
        mem.add(pt, bytes(page))
        report = walk(va, root_pa=0x52000000, mem=mem, page_mode="auto")
        self.assertEqual(report["status"], "ok")
        self.assertEqual(report["page_mode_used"], "64k")
        self.assertEqual(report["evidence"], "modeled")


class FaultTests(unittest.TestCase):
    def test_missing_table_page(self):
        root_pa, mem = channel_pdb_fixture()
        report = walk(0x100000, root_pa=root_pa, mem=MemoryView())
        self.assertEqual(report["status"], "fault")
        self.assertEqual(report["fault"], FAULT_TABLE_MISSING)

    def test_invalid_leaf(self):
        mem = _tree(0x60000000, 0x204608000, None)  # leaf entry stays zero
        report = walk(0x204608000, root_pa=0x60000000, mem=mem)
        self.assertEqual(report["status"], "fault")
        self.assertEqual(report["fault"], FAULT_PTE_INVALID)

    def test_invalid_pde_aperture(self):
        root_pa = 0x60000000
        mem = MemoryView()
        page = bytearray(0x1000)
        struct.pack_into("<Q", page, 0, 0)  # aperture 0 = invalid
        mem.add(root_pa, bytes(page))
        report = walk(0x100000, root_pa=root_pa, mem=mem)
        self.assertEqual(report["status"], "fault")
        self.assertEqual(report["fault"], FAULT_INVALID_APERTURE)

    def test_unsupported_kind_fails_closed(self):
        root_pa, mem = channel_pdb_fixture()
        report = walk(0x100000, root_pa=root_pa, mem=mem, kind="ver9")
        self.assertEqual(report["status"], "fault")
        self.assertEqual(report["fault"], FAULT_UNSUPPORTED)

    def test_va_out_of_range(self):
        root_pa, mem = channel_pdb_fixture()
        report = walk(1 << 60, root_pa=root_pa, mem=mem)
        self.assertEqual(report["status"], "fault")
        self.assertEqual(report["fault"], FAULT_OUT_OF_RANGE)

    def test_bad_page_mode_is_clean_error(self):
        root_pa, mem = channel_pdb_fixture()
        with self.assertRaises(WalkError):
            walk(0x100000, root_pa=root_pa, mem=mem, page_mode="8k")


class CliTests(unittest.TestCase):
    def test_demo_returns_zero(self):
        rc, out, _ = _run_main(["demo"])
        self.assertEqual(rc, 0)
        self.assertIn('"ok": true', out)
        self.assertIn("PASS", out)

    def test_run_returns_three_on_fault(self):
        rc, out, _ = _run_main(["run", "--pdb-pa", "0x60000000", "--va", "0x1000"])
        self.assertEqual(rc, 3)
        self.assertIn("FAULT", out)


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