"""Tests for the launch memory-contract audit (harness/launch_audit.py, docs/48).

Offline only: synthetic page-table trees plus the captured channel-PDB vector.

Run::

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

import contextlib
import io
import json
import os
import struct
import sys
import tempfile
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 harness.launch_audit import (  # noqa: E402
    AUDIT_FORMAT,
    CHECK_FORMAT,
    SPEC_FORMAT,
    VERDICT_ACCESS_DENIED,
    VERDICT_OK,
    VERDICT_PA_MISMATCH,
    VERDICT_UNMAPPED,
    AuditError,
    _tree_for_va,
    audit,
    check,
    main,
    run_demo,
)
from sim.gmmu_walk import (  # noqa: E402
    MemoryView,
    encode_dual_slot,
    encode_pde_pointer,
    encode_pte,
)


def _spec(root_pa: int, mem: MemoryView, contract: list, **extra) -> dict:
    spec = {
        "format": SPEC_FORMAT,
        "name": "test",
        "gmmu": "ver2",
        "page_mode": "4k",
        "pdb_pa": f"{root_pa:#x}",
        "memory": [{"pa": f"{pa:#x}", "hex": b.hex()} for pa, b in _blobs(mem)],
        "contract": contract,
    }
    spec.update(extra)
    return spec


def _blobs(mem: MemoryView):
    return [(int(item["pa"], 16), mem.read(int(item["pa"], 16), item["size"]))
            for item in mem.spans()]


def _drop(mem: MemoryView, drop_pas) -> MemoryView:
    """Copy a memory view without the given table pages (a partial dump)."""
    out = MemoryView()
    for pa, blob in _blobs(mem):
        if pa not in drop_pas:
            out.add(pa, blob)
    return out


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 AuditTests(unittest.TestCase):
    def test_shadow_maps_to_expected_pa_is_ok(self):
        va, shadow = 0x204608000, 0x15E8FA000
        root = 0x60000000
        report = audit(_spec(root, _tree_for_va(root, va, shadow), [
            {"label": "qmd-dest", "va": f"{va:#x}", "access": "write",
             "expect_pa": f"{shadow:#x}"},
        ]))
        self.assertEqual(report["format"], AUDIT_FORMAT)
        row = report["contract"][0]
        self.assertEqual(row["verdict"], VERDICT_OK)
        self.assertTrue(report["ok"])

    def test_host_checks_guest_page_is_pa_mismatch(self):
        va, shadow = 0x204608000, 0x15E8FA000
        root = 0x60000000
        report = audit(_spec(root, _tree_for_va(root, va, shadow), [
            {"label": "qmd-dest", "va": f"{va:#x}", "access": "write",
             "expect_pa": f"{va:#x}"},
        ]))
        row = report["contract"][0]
        self.assertEqual(row["verdict"], VERDICT_PA_MISMATCH)
        self.assertFalse(report["ok"])
        self.assertIn("0x15e8fa000", row["detail"])

    def test_unmapped_leaf_is_unmapped(self):
        va = 0x204608000
        root = 0x60000000
        report = audit(_spec(root, _tree_for_va(root, va, None), [
            {"label": "qmd-dest", "va": f"{va:#x}", "access": "write"},
        ]))
        self.assertEqual(report["contract"][0]["verdict"], VERDICT_UNMAPPED)
        self.assertFalse(report["ok"])

    def test_read_only_leaf_for_write_is_access_denied(self):
        va, target = 0x204608000, 0x15E8FA000
        # build the same tree as _tree_for_va but with a read-only leaf
        i0, i1 = (va >> 47) & 0x3, (va >> 38) & 0x1FF
        i2, i3 = (va >> 29) & 0x1FF, (va >> 21) & 0xFF
        i4 = (va >> 12) & 0x1FF
        base = 0x50000000
        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(base, 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)
        struct.pack_into("<Q", page, i4 * 8, encode_pte(target, read_only=True))
        mem.add(pt, bytes(page))
        root = 0x60000000  # separate root alias is not needed; spec uses 'base'
        report = audit(_spec(base, mem, [
            {"label": "release-slot", "va": f"{va:#x}", "access": "write"},
        ]))
        self.assertEqual(report["contract"][0]["verdict"], VERDICT_ACCESS_DENIED)
        self.assertFalse(report["ok"])
        self.assertNotEqual(root, base)

    def test_spec_errors(self):
        with self.assertRaises(AuditError):
            audit({"format": "nope"})
        with self.assertRaises(AuditError):
            audit({"format": SPEC_FORMAT})
        with self.assertRaises(AuditError):
            audit(_spec(0x60000000, MemoryView(), [], gmmu="ver9"))
        with self.assertRaises(AuditError):
            audit(_spec(0x60000000, MemoryView(), [{"label": "x"}]))

    def test_multiple_contract_legs_counted(self):
        va, shadow = 0x204608000, 0x15E8FA000
        root = 0x60000000
        report = audit(_spec(root, _tree_for_va(root, va, shadow), [
            {"label": "qmd", "va": f"{va:#x}", "access": "write",
             "expect_pa": f"{shadow:#x}"},
            {"label": "missing", "va": f"{va + 0x1000:#x}", "access": "read"},
        ]))
        self.assertEqual(report["summary"][VERDICT_OK], 1)
        self.assertEqual(report["summary"][VERDICT_UNMAPPED], 1)
        self.assertFalse(report["ok"])


class CheckTests(unittest.TestCase):
    """The `check` mode: which table pages does a partial dump still need?"""

    def test_check_empty_memory_reports_root_per_leg(self):
        spec = {
            "format": SPEC_FORMAT,
            "name": "empty-memory",
            "gmmu": "ver2",
            "page_mode": "4k",
            "pdb_pa": "0x12f16e000",
            "contract": [
                {"label": "qmd", "va": "0x204608000", "access": "write"},
                {"label": "ring", "va": "0x100000", "access": "read"},
            ],
        }
        report = check(spec)
        self.assertEqual(report["format"], CHECK_FORMAT)
        self.assertEqual(
            set(report),
            {"format", "name", "firmware", "gmmu", "page_mode", "pdb_pa",
             "memory_pages", "contract", "needed_pages", "complete", "limits"})
        self.assertEqual(report["memory_pages"], 0)
        self.assertFalse(report["complete"])
        for row in report["contract"]:
            self.assertEqual(row["status"], "missing")
            self.assertEqual(row["first_missing"]["pa"], "0x12f16e000")
            self.assertEqual(row["first_missing"]["table_pa"], "0x12f16e000")
            self.assertEqual(row["first_missing"]["level"], "root")
            self.assertEqual(row["first_missing"]["depth"], 0)
            self.assertFalse(row["first_missing"]["present"])
            self.assertEqual(row["path"][-1]["pa"], "0x12f16e000")
            self.assertFalse(row["path"][-1]["present"])
        needs = report["needed_pages"]
        self.assertEqual([page["pa"] for page in needs], ["0x12f16e000"])
        self.assertEqual(needs[0]["levels"], ["root"])
        self.assertEqual(needs[0]["legs"], ["qmd", "ring"])

    def test_check_partial_memory_reports_first_missing_with_path(self):
        va = 0x204608000
        root = 0x60000000
        pde1 = 0x50000000 + 0x2000
        full = _tree_for_va(root, va, 0x12345000)
        partial = _drop(full, {pde1, 0x50000000 + 0x4000})
        report = check(_spec(root, partial, [
            {"label": "qmd", "va": f"{va:#x}", "access": "write"},
        ]))
        row = report["contract"][0]
        self.assertEqual(row["status"], "missing")
        missing = row["first_missing"]
        index = (va >> 29) & 0x1FF
        self.assertEqual(missing["pa"], f"{pde1:#x}")
        self.assertEqual(missing["table_pa"], f"{pde1:#x}")
        self.assertEqual(missing["level"], "pde1")
        self.assertEqual(missing["depth"], 2)
        self.assertEqual(missing["index"], index)
        self.assertEqual(missing["entry_pa"], f"{pde1 + index * 8:#x}")
        self.assertEqual([hop["level"] for hop in row["path"]],
                         ["root", "pde2", "pde1"])
        self.assertEqual(row["path"][0]["table_pa"], f"{root:#x}")
        self.assertTrue(row["path"][0]["present"])
        self.assertFalse(row["path"][-1]["present"])
        self.assertEqual(report["needed_pages"][0]["pa"], f"{pde1:#x}")
        self.assertEqual(report["needed_pages"][0]["legs"], ["qmd"])

    def test_check_partial_memory_marks_complete_and_missing_legs(self):
        root = 0x60000000
        base = 0x51000000
        pd2, pd1 = base + 0x1000, base + 0x2000
        pd0_done, pt_done = base + 0x3000, base + 0x4000
        pd0_missing = base + 0x5000
        target = 0x12345000
        va_done, va_miss = 0x100000, 0x204608000

        def idx(va):
            return ((va >> 47) & 3, (va >> 38) & 0x1FF, (va >> 29) & 0x1FF,
                    (va >> 21) & 0xFF, (va >> 12) & 0x1FF)

        i0d, i1d, i2d, i3d, i4d = idx(va_done)
        i0m, i1m, i2m, _i3m, _i4m = idx(va_miss)
        self.assertEqual((i0d, i1d), (i0m, i1m))  # share root[i0] and pd2[i1]
        self.assertNotEqual(i2d, i2m)

        mem = MemoryView()
        page = bytearray(0x1000)
        struct.pack_into("<Q", page, i0d * 8, encode_pde_pointer(pd2))
        mem.add(root, bytes(page))
        page = bytearray(0x1000)
        struct.pack_into("<Q", page, i1d * 8, encode_pde_pointer(pd1))
        mem.add(pd2, bytes(page))
        page = bytearray(0x1000)
        struct.pack_into("<Q", page, i2d * 8, encode_pde_pointer(pd0_done))
        struct.pack_into("<Q", page, i2m * 8, encode_pde_pointer(pd0_missing))
        mem.add(pd1, bytes(page))
        page = bytearray(0x1000)
        struct.pack_into("<16s", page, i3d * 16, encode_dual_slot(None, pt_done))
        mem.add(pd0_done, bytes(page))
        page = bytearray(0x1000)
        struct.pack_into("<Q", page, i4d * 8, encode_pte(target))
        mem.add(pt_done, bytes(page))

        report = check(_spec(root, mem, [
            {"label": "done", "va": f"{va_done:#x}", "access": "read"},
            {"label": "miss", "va": f"{va_miss:#x}", "access": "write"},
        ]))
        rows = {row["label"]: row for row in report["contract"]}
        self.assertEqual(rows["done"]["status"], "complete")
        self.assertEqual(rows["done"]["resolved_pa"], f"{target:#x}")
        self.assertEqual(rows["miss"]["status"], "missing")
        self.assertEqual(rows["miss"]["first_missing"]["pa"], f"{pd0_missing:#x}")
        self.assertEqual(rows["miss"]["first_missing"]["level"], "pde0")
        self.assertEqual(rows["miss"]["first_missing"]["index"], idx(va_miss)[3])
        self.assertEqual([page["pa"] for page in report["needed_pages"]],
                         [f"{pd0_missing:#x}"])
        self.assertFalse(report["complete"])

    def test_check_needed_pages_dedup_and_ordering(self):
        root = 0x60000000
        pd2_a, pd2_b = 0x70001000, 0x70000000
        va_a1, va_a2 = 0x204608000, 0x100000
        va_b = 0x800000000000
        self.assertEqual((va_a1 >> 47) & 3, 0)
        self.assertEqual((va_a2 >> 47) & 3, 0)
        self.assertEqual((va_b >> 47) & 3, 1)
        mem = MemoryView()
        page = bytearray(0x1000)
        struct.pack_into("<Q", page, 0, encode_pde_pointer(pd2_a))
        struct.pack_into("<Q", page, 8, encode_pde_pointer(pd2_b))
        mem.add(root, bytes(page))
        report = check(_spec(root, mem, [
            {"label": "b-leg", "va": f"{va_b:#x}", "access": "read"},
            {"label": "a1-leg", "va": f"{va_a1:#x}", "access": "write"},
            {"label": "a2-leg", "va": f"{va_a2:#x}", "access": "read"},
        ]))
        pages = report["needed_pages"]
        self.assertEqual([page["pa"] for page in pages],
                         [f"{pd2_b:#x}", f"{pd2_a:#x}"])
        self.assertEqual(pages[0]["levels"], ["pde2"])
        self.assertEqual(pages[0]["legs"], ["b-leg"])
        self.assertEqual(pages[1]["levels"], ["pde2"])
        self.assertEqual(pages[1]["legs"], ["a1-leg", "a2-leg"])
        self.assertFalse(report["complete"])

    def test_check_fault_leg_needs_no_pages(self):
        va = 0x204608000
        root = 0x60000000
        mem = _tree_for_va(root, va, None)  # zero leaf entry -> pte_invalid
        report = check(_spec(root, mem, [
            {"label": "qmd", "va": f"{va:#x}", "access": "write"},
        ]))
        row = report["contract"][0]
        self.assertEqual(row["status"], "fault")
        self.assertEqual(row["fault"], "pte_invalid")
        self.assertIsNone(row["first_missing"])
        self.assertEqual(report["needed_pages"], [])
        self.assertTrue(report["complete"])

    def test_check_64k_dual_missing_big_table(self):
        va = 0x204600000
        root, pd2, pd1, pd0 = 0x60000000, 0x50001000, 0x50002000, 0x50003000
        big_pt, small_pt = 0x80000000, 0x80001000
        i0, i1, i2 = (va >> 47) & 3, (va >> 38) & 0x1FF, (va >> 29) & 0x1FF
        i3 = (va >> 21) & 0xFF
        mem = MemoryView()
        page = bytearray(0x1000)
        struct.pack_into("<Q", page, i0 * 8, encode_pde_pointer(pd2))
        mem.add(root, 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(big_pt, small_pt))
        mem.add(pd0, bytes(page))
        report = check(_spec(root, mem, [
            {"label": "prog", "va": f"{va:#x}", "access": "read"},
        ], page_mode="64k"))
        row = report["contract"][0]
        self.assertEqual(row["status"], "missing")
        missing = row["first_missing"]
        self.assertEqual(missing["pa"], f"{big_pt:#x}")
        self.assertEqual(missing["level"], "pte")
        self.assertEqual(missing["depth"], 4)
        self.assertEqual(missing["index"], (va >> 16) & 0x1FF)

    def test_check_is_deterministic(self):
        va = 0x204608000
        root = 0x60000000
        full = _tree_for_va(root, va, 0x12345000)
        partial = _drop(full, {0x50002000})
        spec = _spec(root, partial, [
            {"label": "qmd", "va": f"{va:#x}", "access": "write"},
        ])
        self.assertEqual(json.dumps(check(spec), sort_keys=True),
                         json.dumps(check(spec), sort_keys=True))

    def test_check_spec_errors(self):
        with self.assertRaises(AuditError):
            check({"format": SPEC_FORMAT})
        with self.assertRaises(AuditError):
            check({"format": "nope", "pdb_pa": "0x1000", "contract": []})
        with self.assertRaises(AuditError):
            check(_spec(0x60000000, MemoryView(), [], gmmu="ver9"))
        with self.assertRaises(AuditError):
            check(_spec(0x60000000, MemoryView(), [{"label": "x"}]))

    def test_cli_check_exit_codes_and_artifact(self):
        va = 0x204608000
        root = 0x60000000
        full = _tree_for_va(root, va, 0x12345000)
        partial = _drop(full, {0x50002000})
        spec = _spec(root, partial, [
            {"label": "qmd", "va": f"{va:#x}", "access": "write"},
        ])
        with tempfile.TemporaryDirectory() as tmp:
            spec_path = os.path.join(tmp, "spec.json")
            with open(spec_path, "w", encoding="utf-8") as fh:
                json.dump(spec, fh)
            rc, out, _ = _run_main(["check", "--spec", spec_path, "--out", tmp])
            self.assertEqual(rc, 3)
            self.assertIn("need 0x50002000", out)
            self.assertIn("dump next:", out)
            report_path = os.path.join(tmp, "launch-audit-check.json")
            self.assertTrue(os.path.isfile(report_path))
            with open(report_path, encoding="utf-8") as fh:
                report = json.load(fh)
            self.assertEqual(report["format"], CHECK_FORMAT)
            self.assertEqual(report["needed_pages"][0]["pa"], "0x50002000")
            # a complete view -> exit 0
            spec["memory"] = [{"pa": f"{pa:#x}", "hex": b.hex()}
                              for pa, b in _blobs(full)]
            with open(spec_path, "w", encoding="utf-8") as fh:
                json.dump(spec, fh)
            rc, out, _ = _run_main(["check", "--spec", spec_path])
            self.assertEqual(rc, 0)
            self.assertIn("complete", out)


class DemoAndCliTests(unittest.TestCase):
    def test_demo_is_green_and_deterministic(self):
        first = run_demo()
        second = run_demo()
        self.assertEqual(first["format"], "greeneye-launch-audit-demo/1")
        self.assertTrue(first["ok"])
        self.assertEqual(json.dumps(first, sort_keys=True),
                         json.dumps(second, sort_keys=True))

    def test_audit_is_deterministic(self):
        va, shadow = 0x204608000, 0x15E8FA000
        root = 0x60000000
        spec = _spec(root, _tree_for_va(root, va, shadow), [
            {"label": "qmd", "va": f"{va:#x}", "access": "write"},
        ])
        self.assertEqual(json.dumps(audit(spec), sort_keys=True),
                         json.dumps(audit(spec), sort_keys=True))

    def test_cli_demo_and_hints(self):
        rc, out, _ = _run_main(["demo"])
        self.assertEqual(rc, 0)
        self.assertIn("captured channel contract: ok=True", out)
        self.assertIn("engine-write-shadow-host-checks-guest", out)
        rc, out, _ = _run_main(["capture-hints"])
        self.assertEqual(rc, 0)
        self.assertIn("CAPTURE HINTS", out)

    def test_engine_write_shadow_demo_verdicts(self):
        demo = run_demo()
        scenarios = {s["name"]: s for s in demo["engine_write_shadow"]}
        mismatch = scenarios["engine-write-shadow-host-checks-guest"]
        aligned = scenarios["engine-write-shadow-bookkeeping-aligned"]
        self.assertEqual(mismatch["expected_verdict"], VERDICT_PA_MISMATCH)
        self.assertEqual(
            mismatch["report"]["contract"][0]["verdict"], VERDICT_PA_MISMATCH)
        self.assertEqual(aligned["expected_verdict"], VERDICT_OK)
        self.assertEqual(
            aligned["report"]["contract"][0]["verdict"], VERDICT_OK)
        self.assertTrue(demo["ok"])
        # the original demo sections are still present and unchanged in shape
        self.assertTrue(demo["captured_channel"]["ok"])
        self.assertEqual(len(demo["qmd_hypotheses"]), 3)

    def test_cli_run_writes_report_and_exit_codes(self):
        va, shadow = 0x204608000, 0x15E8FA000
        root = 0x60000000
        spec = _spec(root, _tree_for_va(root, va, shadow), [
            {"label": "qmd", "va": f"{va:#x}", "access": "write"},
        ])
        with tempfile.TemporaryDirectory() as tmp:
            spec_path = os.path.join(tmp, "spec.json")
            with open(spec_path, "w", encoding="utf-8") as fh:
                json.dump(spec, fh)
            rc, out, _ = _run_main(["run", "--spec", spec_path, "--out", tmp])
            self.assertEqual(rc, 0)
            self.assertTrue(os.path.isfile(os.path.join(tmp, "launch-audit.json")))
            # mismatch -> exit 3
            spec["contract"][0]["expect_pa"] = f"{va:#x}"
            with open(spec_path, "w", encoding="utf-8") as fh:
                json.dump(spec, fh)
            rc, _, _ = _run_main(["run", "--spec", spec_path])
            self.assertEqual(rc, 3)


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