"""Synthetic AST/array diagnostics only. No real brain/GPU or old-file writes."""
import ast
import copy
import json
from pathlib import Path
import tempfile
import types
import unittest
from unittest.mock import patch

import numpy as np

import runtime_light as r


AJ = Path(__file__).absolute().parent.parent / 'apl-timing-lab-v1/gpu_timing.py'
AM = Path(__file__).absolute().parent.parent / 'wholebrain-plastic-interface-v1/overlay.py'


class ScalarTensor:
    """Tiny NumPy scalar stub for diagnostic algebra, not a neuron simulator."""
    def __init__(self, value):
        self.value = np.asarray(value)

    def logical_and_(self, other):
        self.value = self.value & other.value
        return self

    def __lt__(self, other):
        return ScalarTensor(self.value < other)

    def __and__(self, other):
        return ScalarTensor(self.value & other.value)

    def __invert__(self):
        return ScalarTensor(~self.value)

    def all(self):
        return ScalarTensor(self.value.all())

    def item(self):
        return self.value.item()


class TorchScalars:
    bool, int64 = np.bool_, np.int64
    ones = staticmethod(lambda shape, device, dtype: ScalarTensor(np.ones(shape, dtype)))
    zeros = staticmethod(lambda shape, device, dtype: ScalarTensor(np.zeros(shape, dtype)))
    full = staticmethod(lambda shape, value, device, dtype: ScalarTensor(np.full(shape, value, dtype)))
    ones_like = staticmethod(lambda value: ScalarTensor(np.ones_like(value.value)))
    full_like = staticmethod(lambda value, fill: ScalarTensor(np.full_like(value.value, fill)))
    isfinite = staticmethod(lambda value: ScalarTensor(np.isfinite(value)))
    where = staticmethod(lambda condition, first, other: ScalarTensor(np.where(condition.value, first.value, other.value)))


class AstTests(unittest.TestCase):
    def source(self):
        if not AJ.is_file():
            self.skipTest('Original pinned AJ source is not installed')
        return AJ.read_bytes()

    def test_math_projection_equal_light_and_reference_and_guard_explicit(self):
        for mode in ('light', 'reference'):
            tree, audit = r.derive_ast(self.source(), mode)
            self.assertTrue(audit['mathematical_ast_identical'])
            self.assertEqual(audit['math_ast_sha256'], audit['derived_math_ast_sha256'])
            self.assertIn("spec['max_steps']", ast.unparse(tree.body[0]))
            compile(tree, '<synthetic-AST-check>', 'exec')

    def test_reference_entire_tick_loop_remains_identical(self):
        source = ast.parse(self.source())
        original = next(node for node in source.body if isinstance(node, ast.FunctionDef) and node.name == 'run_gpu')
        tree, _ = r.derive_ast(self.source(), 'reference')
        tick = lambda f: next(n for n in ast.walk(f) if isinstance(n, ast.For) and isinstance(n.target, ast.Name) and n.target.id == 'tick')
        self.assertEqual(r._dump(tick(original)), r._dump(tick(tree.body[-1])))

    def test_light_retains_events_requests_kernels_and_no_full_state_cpu_copy(self):
        tree, _ = r.derive_ast(self.source(), 'light')
        function = tree.body[-1]
        tick = next(n for n in ast.walk(function) if isinstance(n, ast.For) and isinstance(n.target, ast.Name) and n.target.id == 'tick')
        text = ast.unparse(tick)
        for expected in ('ext.state_step', 'ext.fill_edges', 'ext.apply_arrivals', 'apl.advance',
                         'event_steps.append', 'event_nodes.append', 'writer.writerows', '_runtime.after_tick',
                         'apl_ring[slot].clone()', 'apl_ring[slot].copy_(h)'):
            self.assertIn(expected, text)
        for forbidden in ('v.cpu()', 'g.cpu()', 'ref.cpu()', 'drive.cpu()', 'h.cpu()', 'queue_summary(',
                          'queue_output', 'hash_output', 'np.isfinite'):
            self.assertNotIn(forbidden, text)

    def test_altered_math_is_detected_and_source_pin_rejects_comments(self):
        tree, audit = r.derive_ast(self.source(), 'light')
        function = copy.deepcopy(tree.body[-1])
        node = next(n for n in ast.walk(function) if isinstance(n, ast.Assign) and
                    ast.unparse(n) == 'g[kc_local_gpu] = g[kc_local_gpu] - delta_si')
        node.value.op = ast.Add()
        self.assertNotEqual(r.math_projection(function), r.math_projection(tree.body[-1]))
        with self.assertRaises(ValueError):
            r.derive_ast(self.source() + b'\n# altered\n', 'light')

    def test_fresh_namespace_no_fake_file_no_old_globals_changed(self):
        self.source()
        module = types.ModuleType('synthetic_AJ')
        module.__file__ = str(AJ)
        module.marker = object()
        before = dict(module.__dict__)
        runtime = types.SimpleNamespace(write_json=lambda *_: None)
        function, audit = r.compile_runtime(module, runtime, 'light')
        self.assertEqual(before, module.__dict__)
        self.assertEqual(function.__globals__['__file__'], str(Path(r.__file__).absolute()))
        self.assertNotEqual(function.__globals__['__file__'], str(AJ))
        self.assertFalse(audit['old_globals_patched'])


class ContractTests(unittest.TestCase):
    def spec(self, steps=1200):
        return dict(max_steps=steps, steps=steps, dt_ms=.1, duration_ms=steps*.1,
                    apl_delay_steps=18, record_history=False)

    def test_directory_parent_traversal_cannot_bypass_input_or_source_boundaries(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory).resolve()
            case, models, graph, bundle = (root / name for name in ('case', 'models', 'graph', 'bundle'))
            for path in (case, models, graph, bundle, root / 'new'):
                path.mkdir()
            safe = root / 'new'
            for protected in (case, models, graph, bundle):
                alias = safe / '..' / protected.name / 'cache'
                for output, build in ((alias, safe / 'build'), (safe / 'out', alias)):
                    with self.subTest(protected=protected.name, output=output), self.assertRaisesRegex(ValueError, 'cannot write'):
                        r.execution_paths(case, output, build, (models, graph, bundle))
            actual = r.execution_paths(case, safe / '..' / 'safe-output', safe / '..' / 'safe-build', (models, graph, bundle))
            self.assertEqual(actual, (case, root / 'safe-output', root / 'safe-build'))
            self.assertFalse((root / 'safe-output').exists())
            self.assertFalse((root / 'safe-build').exists())

    def test_directory_link_alias_cannot_bypass_protected_boundary(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory).resolve()
            case, models = root / 'case', root / 'models'
            case.mkdir(); models.mkdir()
            alias = root / 'alias-to-models'
            try:
                alias.symlink_to(models, target_is_directory=True)
            except (OSError, NotImplementedError) as error:
                self.skipTest('Directory symlink unavailable on this host: ' + str(error))
            for output, build in ((alias / 'out', root / 'build'), (root / 'out', alias / 'build')):
                with self.subTest(output=output), self.assertRaisesRegex(ValueError, 'cannot write'):
                    r.execution_paths(case, output, build, (models,))
            # A protected root supplied through an alias is also canonicalized.
            with self.assertRaisesRegex(ValueError, 'cannot write'):
                r.execution_paths(case, root / 'out', models / 'build', (alias,))

    def test_explicit_bounds_1000_1200_retention_no_hidden_duration(self):
        with tempfile.TemporaryDirectory() as directory:
            for steps in (1000, 1200, 100000):
                self.assertEqual(r.validate_spec(self.spec(steps), directory)['steps'], steps)
            for update in ({'max_steps': 999}, {'max_steps': True}, {'steps': 1200.},
                           {'duration_ms': 100.}, {'dt_ms': .2}, {'apl_delay_steps': 0},
                           {'within_trial_weight_updates': True}, {'checkpoint': None}):
                with self.subTest(update=update), self.assertRaises(ValueError):
                    r.validate_spec(dict(self.spec(), **update), directory)
            Path(directory, 'arrays.npz').write_bytes(b'unsupported checkpoint')
            with self.assertRaises(ValueError):
                r.validate_spec(self.spec(), directory)

    def test_finite_observer_remembers_nan_that_later_disappears(self):
        observer = r.FiniteObserver(TorchScalars, 'CPU-test-stub')
        initial = np.array([1., -2.], np.float64)
        saved = initial.tobytes()
        observer.observe(0, initial)
        observer.observe(1, np.array([np.nan], np.float32))
        observer.observe(2, np.array([0.], np.float64))
        result = observer.result()
        self.assertFalse(result['all_finite'])
        self.assertEqual(result['first_failure_tick'], 1)
        self.assertEqual(result['finite_observed_ticks'], 3)
        self.assertEqual(saved, initial.tobytes())
        with self.assertRaises(ValueError):
            observer.observe(4, initial)

    def test_finite_observer_all_valid_reports_minus_one(self):
        observer = r.FiniteObserver(TorchScalars, 'CPU-test-stub')
        for tick in range(4):
            observer.observe(tick, np.array([1, 0], np.int32), np.array([-0., 1.], np.float64))
        self.assertTrue(observer.result()['all_finite'])
        self.assertEqual(observer.result()['first_failure_tick'], -1)

    def test_summary_has_no_fake_outer_learning_claim_and_light_audit_is_not_full(self):
        obj = object.__new__(r._Runtime)
        obj.mode, obj.spec = 'light', dict(max_steps=1200, steps=1200, outer_protocol='conditioning-owned-by-parent')
        obj.result = dict(first_failure_tick=-1, arriving_edges=11, accepted_arriving_edges=7)
        inherited = dict(learning=False, arriving_edges=0, accepted_arriving_edges=0, complete=True,
                         apl_timing={'generated_trace': 'absent.npy', 'applied_trace': 'absent.npy'})
        final = obj.summary(inherited)
        self.assertNotIn('learning', final)
        self.assertFalse(final['within_trial_weight_updates'])
        self.assertFalse(final['learning_pass_claimed'])
        self.assertFalse(final['full_sequence_bitwise_audit_performed'])
        self.assertEqual(final['arriving_edges'], 11)
        self.assertIsNone(final['apl_timing']['generated_trace'])
        self.assertIsNone(final['apl_timing']['applied_trace'])
        self.assertFalse(obj.timing(1200, np.zeros((18, 4064), np.float32))['full_sequence_bitwise_audit_performed'])
        self.assertEqual(inherited['arriving_edges'], 0)

    def test_terminal_queue_schema_chronological_pending_and_fresh_tick(self):
        class CPUArray:
            def __init__(self, array):
                self.value = array

            def cpu(self):
                return self

            def numpy(self):
                return self.value

        ring = np.zeros((19, 166700), np.uint8)
        ring[2, 3] = 1
        apl_ring = np.repeat(np.arange(18, dtype=np.float32)[:, None], 4064, axis=1)
        apl = types.SimpleNamespace(state=CPUArray(np.array([1., 2.], np.float32)))
        obj = object.__new__(r._Runtime)
        obj.mode, obj.installations, obj.finished = 'light', 1, False
        obj.observer = types.SimpleNamespace(result=lambda: dict(all_finite=True, first_failure_tick=-1,
            finite_observed_ticks=1200, arriving_edges=11, accepted_arriving_edges=7))
        with tempfile.TemporaryDirectory() as directory:
            obj.finish(CPUArray(ring), CPUArray(apl_ring), apl, 1200, Path(directory))
            with np.load(Path(directory)/'terminal-queues.npz', allow_pickle=False) as output:
                self.assertEqual(set(output.files), {'spike_ring', 'apl_state', 'apl_pending', 'next_tick'})
                self.assertEqual(output['spike_ring'].dtype, np.uint8)
                self.assertEqual(output['spike_ring'].tobytes(), ring.tobytes())
                self.assertEqual(output['next_tick'].dtype, np.int64)
                self.assertEqual(output['next_tick'].item(), 1200)
                expected = np.concatenate((apl_ring[1200 % 18:], apl_ring[:1200 % 18]))
                np.testing.assert_array_equal(output['apl_pending'], expected)
            self.assertTrue(obj.finished)
            with self.assertRaises(ValueError):
                obj.finish(CPUArray(ring), CPUArray(apl_ring), apl, 1200, Path(directory))

    def test_pinned_overlay_helpers_toy_units_and_missing_nonzero_fail_closed(self):
        if not AM.is_file():
            self.skipTest('Original frozen AM overlay helper not installed')
        helper = r._load('phase4an_test_overlay_helper', AM, r.AM_OVERLAY_SHA)
        ids, base32 = np.array([3, 9], np.int64), np.array([.275, 1.1], np.float32)
        delta = np.float32(-.25) * base32
        plan = helper.validate_arrays(ids, delta, ids, base32)
        base = np.array([.000275, .0011], np.float64)
        result, _ = helper.materialize_numpy(base, ids, np.zeros(2, bool), plan)
        self.assertTrue(helper.bit_equal(result, base + delta.astype(np.float64) * .001))
        with self.assertRaises(ValueError):
            helper.map_installation(base[:1], ids[:1], np.zeros(1, bool), plan)
        with self.assertRaises(ValueError):
            helper.validate_arrays(ids, np.array([.1, 0], np.float32), ids, base32)

    def test_manifest_cannot_self_authorize_replacement_mask(self):
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory)
            manifest = dict(schema='phase4an-overlay-v1', overlay_name='learned', overlay_sha256='0'*64,
                mask_path=str(path/'mask.npy'), mask_sha256='0'*64, base_candidate_path=str(path/'base.npy'),
                base_candidate_sha256='0'*64, aj_model_dir=str(AJ.parent), am_overlay_path=str(AM),
                am_overlay_sha256=r.AM_OVERLAY_SHA)
            (path/'manifest.json').write_text(json.dumps(manifest), encoding='utf-8')
            with self.assertRaisesRegex(ValueError, 'authority'):
                r.load_overlay(path/'delta.npz', path/'manifest.json')

    def test_learned_manifest_uses_external_base_and_requires_parent_v2(self):
        if not AM.is_file() or not AJ.is_file():
            self.skipTest('Pinned static source dependencies are not installed')
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory)
            ids = np.arange(14232, dtype=np.int64)
            base = np.full(14232, .275, np.float32)
            delta = np.zeros(14232, np.float32); delta[0] = np.float32(-.00275)
            np.save(path/'mask.npy', ids); np.save(path/'base.npy', base)
            np.savez(path/'delta.npz', edge_id=ids, plastic_delta_mv=delta)
            mask_sha = r.digest(path/'mask.npy')
            manifest = dict(schema='phase4an-overlay-v1', overlay_name='learned',
                overlay_sha256=r.digest(path/'delta.npz'), mask_path=str(path/'mask.npy'), mask_sha256=mask_sha,
                base_candidate_path=str(path/'base.npy'), base_candidate_sha256=r.digest(path/'base.npy'),
                aj_model_dir=str(AJ.parent), am_overlay_path=str(AM), am_overlay_sha256=r.AM_OVERLAY_SHA,
                parent_snapshot_sha256='a'*64, eligibility_v2_source_sha256=r.V2_SHA)
            (path/'manifest.json').write_text(json.dumps(manifest), encoding='utf-8')
            # Only the authority fixture is patched; production pin remains fixed.
            with patch.object(r, 'MASK_SHA', mask_sha):
                helper, plan, meta, inputs = r.load_overlay(path/'delta.npz', path/'manifest.json')
                self.assertEqual(plan['plastic_delta_mv'].tobytes(), delta.tobytes())
                self.assertFalse(plan['plastic_delta_mv'].flags.writeable)
                self.assertEqual(meta['parent_snapshot_sha256'], 'a'*64)
                unexpected = dict(manifest, within_trial_weight_updates=True)
                (path/'manifest.json').write_text(json.dumps(unexpected), encoding='utf-8')
                with self.assertRaisesRegex(ValueError, 'Unknown'):
                    r.load_overlay(path/'delta.npz', path/'manifest.json')
                manifest.pop('parent_snapshot_sha256')
                (path/'manifest.json').write_text(json.dumps(manifest), encoding='utf-8')
                with self.assertRaisesRegex(ValueError, 'parent'):
                    r.load_overlay(path/'delta.npz', path/'manifest.json')


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