import hashlib
import os
import struct
import tempfile
import unittest
from pathlib import Path
from unittest.mock import patch

from lightnest.binary import (BinaryFormatError, BinaryReader, Limits, SLOT_MAGIC,
                              find_slots, index_elf, index_lnw1, inspect_binary, read_small_program)


def blob(shape=(2,), data=None):
    name = b'weight'
    header = b'LNW1' + struct.pack('<10I', 1, 8, 1, 2, 32, 32, 16, 16, 1, 1)
    prefix = header + struct.pack('<I', 1) + b'a' + struct.pack('<I', len(name)) + name
    prefix += struct.pack('<I', len(shape)) + struct.pack('<' + 'I' * len(shape), *shape)
    count = 1
    for d in shape:
        count *= d
    return prefix, data if data is not None else b'\0' * (4 * count)


class BinaryTests(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.path = Path(self.tmp.name) / 'input'

    def tearDown(self):
        self.tmp.cleanup()

    def write(self, data):
        self.path.write_bytes(data)

    def test_range_and_chunk_guards(self):
        self.write(bytes(range(100)))
        with BinaryReader(self.path, max_read=16) as r:
            for off, size in ((-1, 1), (100, 1), (0, 17), (1, -1)):
                with self.assertRaises(BinaryFormatError):
                    r.read_at(off, size)
            self.assertEqual(b''.join(r.iter_chunks(10, 20, chunk_size=3)), bytes(range(10, 30)))
            with self.assertRaises(ValueError):
                list(r.iter_chunks(chunk_size=0))
            self.assertEqual(r.sha256(), hashlib.sha256(bytes(range(100))).hexdigest())

    def test_magic_cross_boundary_and_overlaps(self):
        self.write(b'xxxxxxMAGICaaaaaMAGICaaa')
        with BinaryReader(self.path, max_read=8) as r:
            self.assertEqual(list(r.find_magic(b'MAGIC')), [6, 16])
            self.assertEqual(list(r.find_magic(b'aaa')), [11, 12, 13, 21])
            with self.assertRaises(BinaryFormatError):
                list(r.find_magic(b'aaa', max_matches=1))
            self.assertEqual(list(r.find_magic(b'MAGIC', 7, 13)), [])

    def test_sparse_gigabyte_tensor_never_read(self):
        # 4 GiB sparse payload: actual allocated storage stays tiny.
        prefix, _ = blob((1,), b'')
        prefix = prefix[:-4] + struct.pack('<I', 1 << 30)
        with self.path.open('wb') as f:
            f.write(prefix)
            f.truncate(len(prefix) + (1 << 32))
        real = os.pread
        observed = []
        def guarded(fd, n, offset):
            observed.append((offset, n))
            self.assertLessEqual(n, 32)
            self.assertLess(offset, len(prefix))
            return real(fd, n, offset)
        with BinaryReader(self.path, max_read=32) as r, patch('os.pread', guarded):
            info = index_lnw1(r)
        self.assertEqual(info['tensors'][0]['nbytes'], 1 << 32)
        self.assertEqual(info['end_offset'], self.path.stat().st_size)
        self.assertLess(sum(n for _, n in observed), 100)

    def test_lnw1_embedded_offsets(self):
        prefix, data = blob()
        self.write(b'prefix' + prefix + data + b'suffix')
        with BinaryReader(self.path) as r:
            info = index_lnw1(r, 6, end=6+len(prefix)+len(data))
        self.assertEqual(info['tensors'][0]['shape'], [2])
        self.assertEqual(info['tensors'][0]['data_offset'], 6+len(prefix))
        self.assertEqual(info['architecture']['heads'], 2)

    def test_truncated_every_metadata_boundary(self):
        prefix, data = blob()
        for length in range(len(prefix) + len(data)):
            self.write((prefix + data)[:length])
            with BinaryReader(self.path) as r, self.assertRaises(BinaryFormatError):
                index_lnw1(r)

    def test_lnw1_guards(self):
        prefix, data = blob()
        cases = []
        for offset, value in ((4, 2), (8, 0), (16, 3), (36, 100001), (40, 10001), (44, 4097)):
            b = bytearray(prefix + data)
            struct.pack_into('<I', b, offset, value)
            cases.append(b)
        cases.append(prefix[:-4] + struct.pack('<I', 0) + data)
        cases.append(prefix[:48] + b'\xff' + prefix[49:] + data)
        for value in cases:
            self.write(value)
            with BinaryReader(self.path) as r, self.assertRaises(BinaryFormatError):
                index_lnw1(r)
        self.write(prefix + data)
        with BinaryReader(self.path) as r, self.assertRaises(BinaryFormatError):
            index_lnw1(r, limits=Limits(max_metadata_bytes=20))

    def test_nonregular_and_concurrent_truncate(self):
        with self.assertRaises(BinaryFormatError):
            BinaryReader(self.tmp.name)
        fifo = Path(self.tmp.name) / 'fifo'
        os.mkfifo(fifo)
        with self.assertRaises(BinaryFormatError):
            BinaryReader(fifo)
        self.write(b'1234')
        with BinaryReader(self.path) as r:
            self.write(b'')
            with self.assertRaises(BinaryFormatError):
                r.read_at(0, 4)

    def test_slots_bounded_and_ambiguous(self):
        slot = SLOT_MAGIC + struct.pack('<II', 256, 4) + b'RV32\0'.ljust(16, b'\0') + bytes(256)
        self.write(bytes(65530) + slot + slot)
        with BinaryReader(self.path) as r:
            slots = find_slots(r)
        self.assertEqual([s['offset'] for s in slots], [65530, 65826])
        self.assertEqual(slots[0]['data_offset'], 65570)
        self.write(SLOT_MAGIC + struct.pack('<II', 255, 0) + bytes(272))
        with BinaryReader(self.path) as r:
            self.assertEqual(find_slots(r), [])

    def test_unknown_and_explicit_raw(self):
        self.write(b'\x73\x00\x10\x00')
        self.assertEqual(inspect_binary(self.path)['format'], 'unknown')
        self.assertEqual(inspect_binary(self.path, raw_arch='RV32')['architecture_source'], 'explicit hint')
        self.write(b'abc')
        with self.assertRaises(BinaryFormatError):
            inspect_binary(self.path, raw_arch='RV32')

    def test_streaming_sparse_hash_read_bounds(self):
        with self.path.open('wb') as f:
            f.truncate(2 * 1024 * 1024)
        real = os.pread
        sizes = []
        def guarded(fd, size, offset):
            sizes.append(size)
            self.assertLessEqual(size, 1024)
            return real(fd, size, offset)
        with BinaryReader(self.path, max_read=1024) as r, patch('os.pread', guarded):
            digest = r.sha256()
        expected = hashlib.sha256()
        for _ in range(2048):
            expected.update(bytes(1024))
        self.assertEqual(digest, expected.hexdigest())
        self.assertEqual(len(sizes), 2048)

    def test_elf32_big_endian_and_nobits(self):
        ident = b'\x7fELF' + bytes([1, 2, 1]) + bytes(9)
        header = ident + struct.pack('>HHIIIIIHHHHHH', 2, 243, 1, 0, 0, 52, 0, 52, 32, 0, 40, 1, 0)
        section = struct.pack('>IIIIIIIIII', 0, 8, 3, 4096, 92, 1 << 30, 0, 0, 4, 0)
        self.write(header + section)
        with BinaryReader(self.path) as r:
            info = index_elf(r)
        self.assertEqual(info['bits'], 32)
        self.assertEqual(info['byteorder'], 'big')
        self.assertEqual(info['architecture'], 'riscv')
        self.assertFalse(info['sections'][0]['file_backed'])
        self.assertEqual(info['sections'][0]['nbytes'], 1 << 30)

    def test_unknown_100gb_default_reads_only_magic(self):
        with self.path.open('wb') as f:
            f.truncate(100 * 1024**3)
        real = os.pread
        reads = []
        def guarded(fd, n, offset):
            reads.append((offset, n))
            self.assertLessEqual(offset+n, 4)
            return real(fd, n, offset)
        with patch('os.pread', guarded):
            info = inspect_binary(self.path)
        self.assertEqual(reads, [(0, 4)])
        self.assertEqual(info['weight_search']['scope'], 'not searched')
        self.assertFalse(info['slot_search']['whole_file'])

    def test_elf_huge_embedded_payload_prefix_budget(self):
        prefix, _ = blob((1,), b'')
        prefix = prefix[:-4] + struct.pack('<I', 1 << 30)
        payload_size = len(prefix) + (1 << 32)
        ident = b'\x7fELF' + bytes([2, 1, 1]) + bytes(9)
        header = ident + struct.pack('<HHIQQQIHHHHHH', 2, 62, 1, 0, 0, 64, 0, 64, 56, 0, 64, 3, 1)
        names = b'\0.shstrtab\0.rodata\0'
        def section(name, kind, flags, offset, size):
            return struct.pack('<IIQQQQIIQQ', name, kind, flags, 0, offset, size, 0, 0, 1, 0)
        with self.path.open('wb') as f:
            f.write(header + section(0, 0, 0, 0, 0) + section(1, 3, 0, 256, len(names))
                    + section(11, 1, 2, 4096, payload_size) + names)
            f.seek(4096)
            f.write(prefix)
            f.truncate(4096+payload_size)
        real = os.pread
        reads = []
        def guarded(fd, n, offset):
            reads.append((offset, n))
            self.assertLess(offset, 4096+65536)
            return real(fd, n, offset)
        with patch('os.pread', guarded):
            info = inspect_binary(self.path, scan_bytes=4096)
        self.assertEqual(len(info['embedded_weights']), 1)
        self.assertEqual(info['embedded_weights'][0]['tensors'][0]['nbytes'], 1 << 32)
        self.assertEqual(info['weight_search']['search_bytes_read'], 4096)
        self.assertFalse(info['weight_search']['whole_file'])
        self.assertFalse(info['weight_search']['complete_within_selected_ranges'])
        self.assertLess(sum(n for _, n in reads), 5000)
        no_scan = inspect_binary(self.path, scan_bytes=0)
        self.assertEqual(no_scan['embedded_weights'], [])
        self.assertFalse(no_scan['weight_search']['complete_within_selected_ranges'])

    def test_elf_direct_slot_lookup(self):
        ident = b'\x7fELF' + bytes([2, 1, 1]) + bytes(9)
        header = ident + struct.pack('<HHIQQQIHHHHHH', 2, 62, 1, 0, 0, 64, 0, 64, 56, 0, 64, 3, 1)
        names = b'\0.shstrtab\0.lnself\0'
        def section(name, kind, flags, offset, size):
            return struct.pack('<IIQQQQIIQQ', name, kind, flags, 0, offset, size, 0, 0, 1, 0)
        slot = SLOT_MAGIC + struct.pack('<II', 256, 4) + b'RV32\0'.ljust(16, b'\0') + bytes(256)
        self.write(header + section(0, 0, 0, 0, 0) + section(1, 3, 0, 256, len(names))
                   + section(11, 1, 3, 512, len(slot)) + names)
        with self.path.open('r+b') as f:
            f.seek(512)
            f.write(slot)
        info = inspect_binary(self.path, scan_bytes=0)
        self.assertEqual(info['slots'][0]['offset'], 512)
        self.assertEqual(info['slots'][0]['arch'], 'RV32')
        self.assertEqual(info['weight_search']['search_bytes_read'], 0)
        self.assertFalse(info['slot_search']['uniqueness_verified'])

    def test_deep_scan_is_explicit(self):
        prefix, payload = blob()
        slot = SLOT_MAGIC + struct.pack('<II', 256, 0) + bytes(272)
        self.write(b'unknown-prefix' + prefix + payload + slot)
        fast = inspect_binary(self.path)
        self.assertEqual(fast['slots'], [])
        self.assertEqual(fast['embedded_weights'], [])
        deep = inspect_binary(self.path, deep_scan=True)
        self.assertEqual(len(deep['slots']), 1)
        self.assertEqual(len(deep['embedded_weights']), 1)
        self.assertTrue(deep['slot_search']['whole_file'])
        self.assertTrue(deep['weight_search']['complete_within_selected_ranges'])

    def test_small_program_limit_before_read(self):
        self.write(b'\x73\x00\x10\x00')
        self.assertEqual(read_small_program(self.path), self.path.read_bytes())
        with self.path.open('wb') as f:
            f.truncate(1 << 32)
        with patch('os.pread', side_effect=AssertionError('oversized file payload read')):
            with self.assertRaisesRegex(ValueError, 'inspect_binary'):
                read_small_program(self.path)
        for data in (b'', b'abc', b'\x7fELF', b'LNW1'):
            self.write(data)
            with self.assertRaisesRegex(ValueError, 'inspect_binary'):
                read_small_program(self.path)

    def test_minimal_elf_and_guards(self):
        ident = b'\x7fELF' + bytes([2, 1, 1]) + bytes(9)
        def header(**kw):
            vals = dict(kind=2, machine=62, version=1, entry=0, phoff=0, shoff=0, flags=0,
                        ehsize=64, phsize=56, phnum=0, shsize=64, shnum=0, shstr=0)
            vals.update(kw)
            return ident + struct.pack('<HHIQQQIHHHHHH', *vals.values())
        self.write(header())
        with BinaryReader(self.path) as r:
            self.assertEqual(index_elf(r)['architecture'], 'x86_64')
        for bad in (header(phnum=1, phoff=64), header(phnum=65535), header(ehsize=1), header(shoff=64), header(phnum=8193)):
            self.write(bad)
            with BinaryReader(self.path) as r, self.assertRaises(BinaryFormatError):
                index_elf(r)
        segment = struct.pack('<IIQQQQQQ', 1, 5, 0, 0, 0, 120, 120, 4096)
        self.write(header(phoff=64, phnum=1) + segment)
        with BinaryReader(self.path) as r:
            self.assertEqual(index_elf(r)['segments'][0]['nbytes'], 120)


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