"""sync_bodies pipeline without DDF patch — recompile, esopull guard, step count preservation.

Requires GPU."""

import unittest

import numpy as np
import pycuda.driver as cuda
import pycuda.autoinit

from CelerisLab.simulation import Simulation


# Use a small grid for fast compilation and test execution
NX, NY = 128, 64


class TestSyncBodiesSkeleton(unittest.TestCase):
    """Test sync_bodies() with real GPU -- skeleton without DDF patch."""

    def _make_sim(self) -> Simulation:
        """Create a Simulation with a small double_buffer D2Q9 grid."""
        return Simulation(device_id=0)

    def test_remove_body_sync_and_run(self):
        """Remove a body, sync, and continue running without crash."""
        sim = self._make_sim()
        sim.add_body("circle", center=(NX // 4, NY // 2), radius=8)
        sim.initialize()
        sim.run(50)

        # Remove the body and sync
        sim.remove_body(0)
        sim.sync_bodies()

        # Should be able to continue running
        sim.run(50)

        # No body left -- count should be 0
        self.assertEqual(sim.bodies.count, 0)
        sim.close()

    def test_add_body_after_initialize(self):
        """Add a body after initialize, sync, and run."""
        sim = self._make_sim()
        sim.initialize()
        sim.run(50)

        # Add a body (returns -1 since it's staged)
        result_id = sim.add_body("circle", center=(NX // 4, NY // 2), radius=8)
        self.assertEqual(result_id, -1)
        self.assertTrue(sim.bodies.has_pending_edit())

        # Sync commits the body
        sim.sync_bodies()
        self.assertFalse(sim.bodies.has_pending_edit())
        self.assertEqual(sim.bodies.count, 1)

        # Should be able to run with the new body
        sim.run(50)

        # Force readback should work
        force = sim.read_force(0)
        self.assertEqual(force.shape[0], 2)
        sim.close()

    def test_add_then_remove_body(self):
        """Add a body, then remove it, sync -- should result in zero bodies."""
        sim = self._make_sim()
        sim.add_body("circle", center=(NX // 4, NY // 2), radius=8)
        sim.initialize()
        sim.run(50)

        # Add another body and remove the original
        sim.add_body("circle", center=(NX // 2, NY // 2), radius=6)
        sim.remove_body(0)
        sim.sync_bodies()

        # One body remaining (the newly added one, now id=0)
        self.assertEqual(sim.bodies.count, 1)
        sim.run(50)
        sim.close()

    def test_sync_preserves_step_count(self):
        """sync_bodies() should not reset the step counter."""
        sim = self._make_sim()
        sim.add_body("circle", center=(NX // 4, NY // 2), radius=8)
        sim.initialize()
        sim.run(100)
        steps_before = sim.stepper.step_count

        sim.remove_body(0)
        sim.sync_bodies()

        steps_after = sim.stepper.step_count
        self.assertEqual(steps_after, steps_before)
        sim.close()

    def test_no_pending_edit_is_noop(self):
        """sync_bodies() with no pending edits should be a no-op."""
        sim = self._make_sim()
        sim.add_body("circle", center=(NX // 4, NY // 2), radius=8)
        sim.initialize()
        sim.run(50)

        # No edits -- sync should return immediately
        sim.sync_bodies()
        self.assertEqual(sim.bodies.count, 1)
        sim.close()

    def test_run_discards_pending_without_sync(self):
        """Running without sync_bodies() should discard pending edits."""
        sim = self._make_sim()
        sim.add_body("circle", center=(NX // 4, NY // 2), radius=8)
        sim.initialize()
        sim.run(50)

        # Stage a removal but don't sync
        sim.remove_body(0)
        self.assertTrue(sim.bodies.has_pending_edit())

        # run() auto-discards pending
        sim.run(50)
        self.assertFalse(sim.bodies.has_pending_edit())
        self.assertEqual(sim.bodies.count, 1)
        sim.close()

    def test_esopull_raises_not_implemented(self):
        """sync_bodies() with esopull should raise NotImplementedError."""
        # Create a sim with esopull streaming
        from CelerisLab.config import load_lbm_config
        cfg = load_lbm_config()
        cfg.streaming = "esopull"
        # We need to build the sim manually to override streaming
        sim = Simulation.__new__(Simulation)
        sim._stream = None
        sim.lbm_cfg = cfg
        from CelerisLab.config import BodyConfig
        sim.body_cfg = BodyConfig()
        from CelerisLab.cuda.context import CudaContext
        sim.ctx = CudaContext(0)
        from CelerisLab.cuda import compiler_v2 as compiler
        arch = sim._resolve_compile_arch = lambda: sim.ctx.sm_arch
        arch_val = CudaContext(0).sm_arch
        compiler.generate_config(cfg, n_objects=0)
        ptx_path = compiler.compile_kernel(arch=arch_val)
        module = compiler.load_module(ptx_path)
        sim._ptx_path = ptx_path
        sim._module = module
        from CelerisLab.lbm.field import LBMField
        sim.field = LBMField(cfg, module)
        from CelerisLab.lbm.stepper import LBMStepper
        sim.stepper = LBMStepper(sim.field, module, cfg)
        from CelerisLab.body.manager import ObjectManager
        sim.bodies = ObjectManager(
            cfg.nx, cfg.ny, cfg.nz, cfg.nq, cfg)
        sim._initialized = True

        sim.add_body("circle", center=(NX // 4, NY // 2), radius=8)
        with self.assertRaises(NotImplementedError):
            sim.sync_bodies()
        sim.close()


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