import struct
import tempfile
import time
import unittest
import zlib
from pathlib import Path
from unittest.mock import patch

import nand_put as x


class FakePort:
    def __init__(self, callback):
        self.callback = callback
        self.rx = bytearray()
        self.requests = []

    def write(self, data):
        self.requests.append(data)
        reply = self.callback(data, len(self.requests))
        if reply:
            self.rx.extend(reply)
        return len(data)

    def read(self, size):
        data = self.rx[:size]
        del self.rx[:size]
        return bytes(data)


def reply(request, status=0, payload=b""):
    magic, _, cmd, _, seq, size = x.HEADER.unpack(request[:16])
    assert magic == x.MAGIC and len(request) == 20 + size
    assert zlib.crc32(request[:-4]) == struct.unpack("<I", request[-4:])[0]
    return x.encode(cmd | 0x80, seq, bytes([status]) + payload)


class ProtocolTests(unittest.TestCase):
    def test_handshake_nonce_and_rejection(self):
        port = FakePort(lambda data, _: reply(data, payload=data[16:20]))
        port.rx.extend(b"NAND-XFER READY\n")
        with patch.object(x.secrets, "token_bytes", return_value=b"test"):
            x.Client(port).connect()
        self.assertEqual(port.requests[0][5], x.HELLO)
        broken = FakePort(lambda data, _: reply(data, payload=b"bad!"))
        broken.rx.extend(b"NAND-XFER READY\n")
        with self.assertRaisesRegex(ValueError, "handshake"):
            x.Client(broken).connect()

    def test_connect_hello_without_ready_banner(self):
        port = FakePort(lambda data, _: reply(data, payload=data[16:20]))
        client = x.Client(port)
        client.wait_ready = lambda timeout=0.4: False
        with patch.object(x.secrets, "token_bytes", return_value=b"test"):
            client.connect()
        self.assertEqual(port.requests[0][5], x.HELLO)

    def test_crc_and_large_frame(self):
        packet = x.encode(x.BLOCK, 0x12345678, bytes(range(256)) * 16)
        self.assertEqual(len(packet), 16 + 4096 + 4)
        self.assertEqual(zlib.crc32(packet[:-4]), struct.unpack("<I", packet[-4:])[0])
        with self.assertRaises(ValueError):
            x.encode(x.BLOCK, 1, bytes(4353))

    def test_utf8_paths_and_escape(self):
        self.assertEqual(x.path_bytes("/extflash/eaf/文件 1.eaf"),
                         b"/extflash/eaf/" + "文件 1.eaf".encode())
        for path in ("/extflash2/a", "/extflash/../a", "/extflash//a",
                     "/extflash/a.", "/extflash/a?b", "/extflash/a\x00b"):
            with self.subTest(path=path), self.assertRaises(ValueError):
                x.path_bytes(path)

    def test_lost_ack_reuses_sequence_and_packet(self):
        port = FakePort(lambda data, n: None if n == 1 else reply(data, payload=b"done"))
        result = x.Client(port).call(x.BLOCK, b"test", timeout=0.001)
        self.assertEqual(result, b"done")
        self.assertEqual(port.requests, [port.requests[0], port.requests[0]])

    def test_bad_crc_then_retry(self):
        port = FakePort(lambda data, n: reply(data)[:-1] + b"!" if n == 1 else reply(data))
        self.assertEqual(x.Client(port).call(x.INFO, timeout=0.01), b"")
        self.assertEqual(len(port.requests), 2)

    def test_device_error_is_not_retried(self):
        port = FakePort(lambda data, n: reply(data, status=4))
        with self.assertRaisesRegex(RuntimeError, "I/O error"):
            x.Client(port).call(x.BLOCK, b"x")
        self.assertEqual(len(port.requests), 1)

    def test_paged_directory(self):
        first = struct.pack("<HB", 1, 0) + struct.pack("<BBQ", 1, 3, 0) + b"eaf"
        second = struct.pack("<HB", 2, 1) + struct.pack("<BBQ", 0, 5, 12) + b"a.eaf"
        port = FakePort(lambda data, n: reply(data, payload=first if n == 1 else second))
        self.assertEqual(list(x.Client(port).listdir("/extflash")),
                         [(1, "eaf", 0), (0, "a.eaf", 12)])
        self.assertEqual(port.requests[0][16:18], b"\x00\x00")
        self.assertEqual(port.requests[1][16:18], b"\x01\x00")

    def test_format_cancel_never_contacts_device(self):
        class DenyClient:
            def call(self, command, *args, **kwargs):
                self.fail("device contacted")
        with self.assertRaisesRegex(ValueError, "cancelled"):
            x.format_with_confirmation(DenyClient(), input_fn=lambda _: "no")

    def test_format_challenge_and_single_confirm(self):
        calls = []
        class Client:
            seq = 123
            def call(self, cmd, payload=b"", **kwargs):
                calls.append((cmd, payload, kwargs))
                return b"\x12\x34\x56\x78" if cmd == x.PREPARE else b""
        x.format_with_confirmation(Client(), input_fn=lambda _: "FORMAT /extflash")
        self.assertEqual([(c, data) for c, data, _ in calls],
                         [(x.PREPARE, b""), (x.CONFIRM, b"\x12\x34\x56\x78" + struct.pack("<I", 123))])
        self.assertEqual(calls[1][2]["retries"], 1)

    def test_directory_rejects_link_and_keeps_relative_names(self):
        with tempfile.TemporaryDirectory() as folder:
            root = Path(folder)
            (root / "nested").mkdir()
            (root / "nested" / "a.eaf").write_bytes(b"abc")
            entries = list(x.directory_entries(root))
            self.assertEqual([(p, is_dir) for _, p, is_dir in entries],
                             [(('nested',), True), (('nested', 'a.eaf'), False)])
            with patch.object(Path, "is_junction", lambda self: self.name == "nested"):
                with self.assertRaisesRegex(ValueError, "junction"):
                    list(x.directory_entries(root))

    def test_upload_empty_file_readback(self):
        commands = []
        def callback(packet, _):
            _, _, cmd, _, _, _ = x.HEADER.unpack(packet[:16])
            commands.append(cmd)
            return reply(packet, payload=struct.pack("<QI", 0, 0) if cmd in (x.END, x.VERIFY) else b"")
        with tempfile.TemporaryDirectory() as folder:
            source = Path(folder) / "empty"
            source.touch()
            x.Client(FakePort(callback)).upload(source, "/extflash/empty")
        self.assertEqual(commands, [x.BEGIN, x.END, x.VERIFY])

    def test_upload_error_stops_quickly(self):
        def callback(packet, _):
            if packet[5] == x.BEGIN:
                return reply(packet, status=5)
            return None
        port = FakePort(callback)
        port.timeout = 0.05
        with tempfile.TemporaryDirectory() as folder:
            source = Path(folder) / "a.txt"
            source.write_bytes(b"hi")
            started = time.monotonic()
            with self.assertRaisesRegex(RuntimeError, "bad state"):
                x.Client(port).upload(source, "/extflash/a.txt")
        self.assertLess(time.monotonic() - started, 3)
        self.assertEqual(port.requests[0][5], x.BEGIN)
        self.assertEqual(port.requests[1][5], x.CANCEL)


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