"""Pure-Python BLAKE3 / zstd / XXH64 used by the independent Oak reader.

Vectors: tests/data/blake3_vectors.json holds BLAKE3 of the spec's test input

pattern (byte i = i % 251) at chunk/block boundaries, produced once with the

reference ``blake3`` package; tests/data/zstd_vectors.json holds frames made

by the ``zstd`` CLI (levels 1-19, multi-block, RLE, no-checksum). When the

``zstd`` CLI is installed, a live differential also runs.

"""

import hashlib

import json

import os

from pathlib import Path

import shutil

import subprocess

import sys

import unittest

ROOT = Path(__file__).resolve().parents[1]

sys.path.insert(0, str(ROOT / "scripts"))

from oakbench import pureblake3, purezstd  # noqa: E402

DATA = ROOT / "tests" / "data"

def pattern(n):

    return bytes(i % 251 for i in range(n))

class Blake3Tests(unittest.TestCase):

    def test_reference_vectors_across_chunk_and_tree_boundaries(self):

        vectors = json.loads((DATA / "blake3_vectors.json").read_text())

        self.assertGreaterEqual(len(vectors), 20)

        for length, expected in vectors.items():

            with self.subTest(length=length):

                self.assertEqual(pureblake3.hexdigest(pattern(int(length))), expected)

    def test_single_bit_flip_changes_digest(self):

        data = bytearray(pattern(2049))

        before = pureblake3.hexdigest(bytes(data))

        data[1500] ^= 1

        self.assertNotEqual(before, pureblake3.hexdigest(bytes(data)))

class Xxh64Tests(unittest.TestCase):

    def test_known_values(self):

        self.assertEqual(purezstd.xxh64(b""), 0xEF46DB3751D8E999)

        self.assertEqual(purezstd.xxh64(b"a"), 0xD24EC4F1A98C6E5B)

        self.assertEqual(purezstd.xxh64(b"abc"), 0x44BC2CF5AD770999)

class ZstdVectorTests(unittest.TestCase):

    def setUp(self):

        self.vectors = json.loads((DATA / "zstd_vectors.json").read_text())["vectors"]

    def test_cli_generated_frames_decode_exactly(self):

        for vector in self.vectors:

            with self.subTest(vector=vector["name"]):

                out = purezstd.decompress(bytes.fromhex(vector["frame_hex"]))

                self.assertEqual(len(out), vector["size"])

                self.assertEqual(hashlib.sha256(out).hexdigest(), vector["sha256"])

    def test_every_single_byte_corruption_is_rejected_or_detected(self):

        # Flip each byte of a checksummed compressed frame: the decoder must

        # raise or (never) silently return the original plaintext's digest

        # for different input bytes being accepted as the same content.

        vector = next(v for v in self.vectors if v["name"] == "small-text-l19")

        frame = bytearray(bytes.fromhex(vector["frame_hex"]))

        silent = 0

        for index in range(len(frame)):

            corrupt = bytearray(frame)

            corrupt[index] ^= 0x41

            try:

                out = purezstd.decompress(bytes(corrupt))

            except (purezstd.ZstdError, purezstd.ZstdBoundError):

                continue

            if hashlib.sha256(out).hexdigest() != vector["sha256"]:

                silent += 1

        # The XXH64 content checksum must catch every altered output.

        self.assertEqual(silent, 0)

    def test_truncation_and_trailing_garbage_fail(self):

        frame = bytes.fromhex(self.vectors[0]["frame_hex"])

        for cut in (1, 4, 5, len(frame) // 2, len(frame) - 1):

            with self.assertRaises(purezstd.ZstdError):

                purezstd.decompress(frame[:cut])

        with self.assertRaises(purezstd.ZstdError):

            purezstd.decompress(frame + b"\x00garbage")

    def test_declared_size_and_window_bounds_are_bound_errors(self):

        frame = b"\x28\xb5\x2f\xfd\xa0" + (1 << 30).to_bytes(4, "little") + b"\x01\x00\x00"

        with self.assertRaises(purezstd.ZstdBoundError):

            purezstd.decompress(frame, max_output=1 << 20, max_window=1 << 20)

    def test_raw_encoder_roundtrip_and_multiblock(self):

        for payload in (b"", b"x", os.urandom(300_000)):

            self.assertEqual(purezstd.decompress(purezstd.encode_raw(payload)), payload)

        self.assertEqual(purezstd.frame_content_size(purezstd.encode_raw(b"abc")), 3)

    # A hand-built compressed block: raw literals of size 0, one sequence with

    # RLE tables (LL=0, OF code 3 -> offset 5, ML=3) and a 3-bit bitstream.

    MATCH_BLOCK = bytes([0x3D, 0x00, 0x00, 0x00, 0x01, 0x54, 0x00, 0x03, 0x00, 0x08])

    FRAME_HEADER = b"\x28\xb5\x2f\xfd\x00\x00"  # no content size, 1 KiB window

    def test_match_within_frame_decodes(self):

        raw_block = bytes([(6 << 3), 0x00, 0x00]) + b"abcdef"

        self.assertEqual(purezstd.decompress(self.FRAME_HEADER + raw_block + self.MATCH_BLOCK), b"abcdefbcd")

    def test_match_into_previous_frame_is_rejected(self):

        with self.assertRaisesRegex(purezstd.ZstdError, "outside frame window"):

            purezstd.decompress(purezstd.encode_raw(b"abcdef") + self.FRAME_HEADER + self.MATCH_BLOCK)

    def test_match_beyond_declared_window_is_rejected(self):

        # Single-segment frame whose window (= content size 9) is fine, then a

        # frame declaring content size 4 but reaching back 5 bytes.

        frame = b"\x28\xb5\x2f\xfd\x20\x09" + bytes([(6 << 3), 0x00, 0x00]) + b"abcdef" + self.MATCH_BLOCK

        self.assertEqual(purezstd.decompress(frame), b"abcdefbcd")

        small = b"\x28\xb5\x2f\xfd\x20\x04" + bytes([(6 << 3), 0x00, 0x00]) + b"abcdef" + self.MATCH_BLOCK

        with self.assertRaises(purezstd.ZstdError):

            purezstd.decompress(small)

    @unittest.skipUnless(shutil.which("zstd"), "zstd CLI not installed")

    def test_live_differential_against_zstd_cli(self):

        corpus = [pattern(70_000), os.urandom(140_000), b"oak tree blob " * 20_000,

                  (ROOT / "scripts" / "bench.py").read_bytes()]

        for data in corpus:

            for level in ("-1", "-6", "-19", "--fast=3", "--long=20"):

                with self.subTest(size=len(data), level=level):

                    frame = subprocess.run(["zstd", "-q", "-c", level], input=data, capture_output=True,

                                           check=True).stdout

                    self.assertEqual(purezstd.decompress(frame), data)

if __name__ == "__main__":

    unittest.main()