"""Local/remote pure protocol tests; no Brian, Torch, GPU or network execution.

Run with the already-provisioned Python and ``-B test_protocol.py``.  The
sidecar process is mocked with in-memory stdin and an anonymous stdout pipe.
These tests validate transport and declared arithmetic/input contracts, not
the scientific trajectory or independent correctness of the Torch component.
"""
from __future__ import annotations

import hashlib
import io
import json
import math
import os
from pathlib import Path
import struct
import sys
import threading
import types
import unittest
from unittest import mock

import numpy as np


HERE = Path(__file__).resolve().parent


def load_no_cache(path, name):
    module = types.ModuleType(name)
    module.__file__ = str(path)
    exec(compile(path.read_bytes(), str(path), "exec"), module.__dict__)
    return module


ORACLE = load_no_cache(HERE / "oracle_apl.py", "phase4ag_protocol_oracle")


def handshake(**updates):
    result = {"ok": True, "protocol": "phase4ag-apl-v1", "n_neurons": 166700,
        "n_apl": 2, "n_kc": 4064, "dt_ms": .1,
        "files_sha256": {name: "frozen-test-hash" for name in ORACLE.BUNDLE_FILES},
        "decay_f32_bits": f"{int(np.float32(math.exp(-.1 / 20.)).view(np.uint32)):08x}",
        "gain_f32_bits": f"{int(np.float32(.01).view(np.uint32)):08x}"}
    result.update(updates)
    return (json.dumps(result) + "\n").encode("utf-8")


def frame(tick=0, values=None, byte_count=None):
    values = np.arange(4068, dtype="<f4") if values is None else np.asarray(values, dtype="<f4")
    count = values.nbytes if byte_count is None else byte_count
    return struct.pack("<qQ", tick, count) + values.tobytes()


class MemoryInput(io.BytesIO):
    def close(self):
        if not self.closed:
            self.closed_data = self.getvalue()
        super().close()


class FakeProcess:
    """The only use of OS descriptors is an anonymous read-only response pipe."""
    def __init__(self, response, status=0, chunk_size=97):
        read_fd, write_fd = os.pipe()
        self.stdin = MemoryInput()
        self.stdout = os.fdopen(read_fd, "rb", buffering=0)
        self.status = status
        self.terminated = False

        def write_response():
            try:
                for begin in range(0, len(response), chunk_size):
                    view = memoryview(response[begin:begin + chunk_size])
                    while view:
                        view = view[os.write(write_fd, view):]
            except (BrokenPipeError, OSError):
                # Invalid-response tests deliberately close the reader early.
                pass
            finally:
                os.close(write_fd)

        self.thread = threading.Thread(target=write_response, daemon=True)
        self.thread.start()

    def poll(self):
        return self.status

    def wait(self, timeout=None):
        return self.status

    def terminate(self):
        self.terminated = True

    def kill(self):
        self.terminated = True

    def close(self):
        if not self.stdin.closed:
            self.stdin.close()
        if not self.stdout.closed:
            self.stdout.close()
        self.thread.join(timeout=2)


class ProtocolTests(unittest.TestCase):
    def create_worker(self, response, status=0):
        process = FakeProcess(response, status=status)
        self.addCleanup(process.close)
        source_patch = mock.patch.object(ORACLE, "_digest", return_value="frozen-test-hash")
        process_patch = mock.patch.object(ORACLE.subprocess, "Popen", return_value=process)
        source_patch.start()
        process_mock = process_patch.start()
        self.addCleanup(source_patch.stop)
        self.addCleanup(process_patch.stop)
        worker = ORACLE._APLWorker(Path("science-python"), Path("legacy_apl_worker.py"), Path("bundle"), None, timeout=2)
        self.addCleanup(worker.abort)
        return worker, process, process_mock

    def test_exact_handshake_frame_and_clean_eof(self):
        worker, process, process_mock = self.create_worker(handshake() + frame())
        globals_ = np.array([1, 40, 166699], dtype=np.int32)
        drive, state, inhibition = worker.step(0, globals_)
        values = np.arange(4068, dtype=np.float32)
        np.testing.assert_array_equal(drive, values[:2])
        np.testing.assert_array_equal(state, values[2:4])
        np.testing.assert_array_equal(inhibition, values[4:])
        request = process.stdin.getvalue()
        self.assertEqual(struct.unpack("<qQ", request[:16]), (0, 3))
        np.testing.assert_array_equal(np.frombuffer(request[16:], dtype="<i4"), globals_)
        self.assertEqual(worker.finish(1), 0)
        command = process_mock.call_args.args[0]
        self.assertEqual(command, ["science-python", "-B", "legacy_apl_worker.py", "--bundle-dir", "bundle", "--dt-ms", "0.1"])
        self.assertFalse(process_mock.call_args.kwargs["shell"])

    def test_empty_spike_frame_still_advances_one_tick(self):
        worker, process, _ = self.create_worker(handshake() + frame(values=np.zeros(4068, np.float32)))
        arrays = worker.step(0, np.empty(0, np.int32))
        self.assertTrue(all(np.all(value == 0) for value in arrays))
        self.assertEqual(process.stdin.getvalue(), struct.pack("<qQ", 0, 0))
        self.assertEqual(worker.finish(1), 0)

    def test_two_consecutive_ticks(self):
        worker, _, _ = self.create_worker(handshake() + frame(0) + frame(1))
        worker.step(0, np.array([4], np.int32))
        worker.step(1, np.array([8, 10], np.int32))
        self.assertEqual(worker.finish(2), 0)

    def test_rejects_wrong_protocol_or_dimensions(self):
        for update in ({"protocol": "other"}, {"ok": False}, {"n_neurons": 42},
                       {"n_kc": 3}, {"dt_ms": .2}, {"gain_f32_bits": "00000000"},
                       {"files_sha256": {}}):
            with self.subTest(update=update):
                with self.assertRaises(ValueError):
                    self.create_worker(handshake(**update))

    def test_rejects_wrong_tick_or_payload_length(self):
        for response in (frame(tick=1), frame(byte_count=8)):
            with self.subTest(response_header=response[:16]):
                worker, _, _ = self.create_worker(handshake() + response)
                with self.assertRaises(ValueError):
                    worker.step(0, np.empty(0, np.int32))

    def test_rejects_truncated_response(self):
        worker, _, _ = self.create_worker(handshake() + frame()[:-1])
        with self.assertRaises(EOFError):
            worker.step(0, np.empty(0, np.int32))

    def test_rejects_nonfinite_or_negative_apl_payload(self):
        for bad_value in (np.nan, np.inf, -1):
            with self.subTest(value=bad_value):
                values = np.zeros(4068, np.float32)
                values[-1] = bad_value
                worker, _, _ = self.create_worker(handshake() + frame(values=values))
                with self.assertRaises(ValueError):
                    worker.step(0, np.empty(0, np.int32))

    def test_rejects_invalid_spike_dtype_order_and_range(self):
        for values in (np.array([1], np.int64), np.array([2, 1], np.int32),
                       np.array([1, 1], np.int32), np.array([-1], np.int32),
                       np.array([166700], np.int32)):
            with self.subTest(values=values):
                worker, process, _ = self.create_worker(handshake() + frame())
                with self.assertRaises(ValueError):
                    worker.step(0, values)
                self.assertEqual(process.stdin.getvalue(), b"")

    def test_rejects_skipped_request_tick(self):
        worker, process, _ = self.create_worker(handshake() + frame())
        with self.assertRaises(ValueError):
            worker.step(1, np.empty(0, np.int32))
        self.assertEqual(process.stdin.getvalue(), b"")

    def test_rejects_failed_exit_and_trailing_stdout(self):
        for suffix, status in ((b"", 7), (b"unexpected", 0)):
            with self.subTest(status=status, suffix=suffix):
                worker, _, _ = self.create_worker(handshake() + frame() + suffix, status=status)
                worker.step(0, np.empty(0, np.int32))
                with self.assertRaises(ValueError):
                    worker.finish(1)

    def test_pipe_reader_timeout(self):
        read_fd, write_fd = os.pipe()
        stream = os.fdopen(read_fd, "rb", buffering=0)
        reader = ORACLE._PipeReader(stream, .05)
        try:
            with self.assertRaises(TimeoutError):
                reader.read(1)
        finally:
            os.close(write_fd)
            reader.thread.join(timeout=2)
            stream.close()
        self.assertFalse(reader.thread.is_alive())


class ArithmeticAndIdentityContractTests(unittest.TestCase):
    def test_si_conversion_must_follow_float64_promotion(self):
        h32 = np.array([.01, 1 / 7, 123.456], dtype=np.float32)
        correct = h32.astype(np.float64) * np.float64(.001)
        early_conversion = (h32 * np.float32(.001)).astype(np.float64)
        self.assertFalse(np.array_equal(correct.view(np.uint64), early_conversion.view(np.uint64)))
        for actual, original in zip(correct, h32):
            self.assertEqual(actual.hex(), (float(original) * float(.001)).hex())

    def test_declared_overlay_does_not_guard_refractory_targets(self):
        before = np.array([0, -.002, .001], dtype=np.float64)
        refractory = np.array([22, 21, 0], dtype=np.int32)
        h32 = np.array([1, 2, 3], dtype=np.float32)
        delta = h32.astype(np.float64) * np.float64(.001)
        after = before - delta
        self.assertTrue(np.all(after != before))
        np.testing.assert_array_equal(refractory, [22, 21, 0])
        self.assertEqual(after[0], -.001)

    def test_af_queue_hash_keeps_original_noncontiguous_ids_and_rejections(self):
        path = HERE.parent / "author-fullgraph-lab-v1" / "oracle_full.py"
        self.assertEqual(hashlib.sha256(path.read_bytes()).hexdigest(), ORACLE.AF_SHA256)
        af = load_no_cache(path, "phase4ag_protocol_af")
        edges = np.array([81, 2, 500], dtype=np.int64)
        accepted = np.array([1, 0, 1], dtype=np.uint8)
        original_edges, original_mask = edges.copy(), accepted.copy()
        result = af.queue_summary(18, edges, accepted, chunk_bytes=8)
        packet = b"phase4af-queue-v1\0" + struct.pack("<qq", 18, 3) + edges.astype("<i8").tobytes() + accepted.tobytes()
        self.assertEqual(result["joint_sha256"], hashlib.sha256(packet).hexdigest())
        self.assertEqual((result["count"], result["accepted_count"]), (3, 2))
        self.assertEqual(result, af.queue_summary(18, edges, accepted, chunk_bytes=8192))
        np.testing.assert_array_equal(edges, original_edges)
        np.testing.assert_array_equal(accepted, original_mask)

    def test_empty_queue_hash_includes_tick(self):
        af = load_no_cache(HERE.parent / "author-fullgraph-lab-v1" / "oracle_full.py", "phase4ag_empty_af")
        first = af.queue_summary(0, np.empty(0, np.int64), np.empty(0, np.uint8))
        second = af.queue_summary(1, np.empty(0, np.int64), np.empty(0, np.uint8))
        self.assertEqual(first["ordered_edge_sha256"], hashlib.sha256(b"").hexdigest())
        self.assertNotEqual(first["joint_sha256"], second["joint_sha256"])

    def test_oracle_import_has_no_brian_or_torch_dependency(self):
        before = set(sys.modules)
        load_no_cache(HERE / "oracle_apl.py", "phase4ag_second_import")
        added = set(sys.modules) - before
        self.assertFalse(any(name == "brian2" or name.startswith("brian2.") or
                             name == "torch" or name.startswith("torch.") for name in added))


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