pybevy 0.2.1

use std::sync::atomic::{AtomicUsize, Ordering};

use bevy::ecs::message::Messages;
use pybevy_core::registry::global_registry;
// Re-export base classes from pybevy_core
pub use pybevy_core::{PyMessage, PyMessageId};
use pyo3::{
    IntoPyObjectExt,
    exceptions::{PyStopIteration, PyTypeError},
    prelude::*,
    types::PyType,
};

use super::{
    messages::{
        CustomMessage1, CustomMessage2, CustomMessage3, CustomMessage4, CustomMessage5,
        CustomMessage6, CustomMessage7, CustomMessage8, CustomMessage9, CustomMessage10,
        CustomMessage11, CustomMessage12, CustomMessage13, CustomMessage14, CustomMessage15,
        CustomMessage16, CustomMessage17, CustomMessage18, CustomMessage19, CustomMessage20,
        MessageRegistry, MessageType, PyMessages,
    },
    world::PyWorld,
};
use crate::ecs::messages::PyMessageType;

// Macro to generate match arms for writing messages
macro_rules! write_message_arms {
    ($message_num:expr, $world:expr, $py:expr, $message:expr, $($idx:literal => $type_num:literal),+ $(,)?) => {
        paste::paste! {
            match $message_num {
                $(
                    $idx => {
                        let mut messages = $world.get_resource_or_insert_with(|| Messages::<[<CustomMessage $type_num>]>::default());
                        Box::new(messages.write([<CustomMessage $type_num>]($message.clone_ref($py))))
                    }
                )+
                _ => return Err(PyTypeError::new_err("Invalid message number")),
            }
        }
    };
}

#[derive(Debug, Clone, PartialEq)]
pub enum MessageClass {
    Reader,
    Writer,
}

#[pyclass(frozen)]
#[derive(Debug, Clone, PartialEq)]
pub struct MessageTypeParam {
    pub(crate) ty: MessageClass,
    pub(crate) message_type: MessageType,
}

/// Write-only handle to a message channel. Used as a system parameter to send messages.
///
/// Usage: `def my_system(writer: MessageWriter[DamageEvent])`
#[pyclass(name = "MessageWriter", frozen)]
pub struct PyMessageWriter {
    pub(crate) message_type: MessageType,
    pub(crate) world: PyWorld,
}

impl PyMessageWriter {
    // TODO: Add constructor when needed by system integration
}

#[pymethods]
impl PyMessageWriter {
    #[classmethod]
    #[pyo3(signature = (key, /))]
    pub fn __class_getitem__(
        _cls: &Bound<'_, PyType>,
        key: &Bound<'_, PyAny>,
    ) -> PyResult<Py<PyAny>> {
        let py = key.py();

        // Extract message type from key (e.g., AppExit from MessageWriter[AppExit])
        let type_obj = key.cast::<PyType>()?;
        let py_message_type = PyMessageType::from_message_type(&type_obj)?;

        // Create MessageTypeParam with Writer class
        let param = MessageTypeParam {
            ty: MessageClass::Writer,
            message_type: py_message_type.0,
        };

        param.into_py_any(py)
    }

    pub fn write(&self, py: Python, message: Py<PyAny>) -> PyResult<PyMessageId> {
        let bound_message = message.bind(py);

        match &self.message_type {
            MessageType::KeyboardInput => Err(PyTypeError::new_err(
                "KeyboardInput events are read-only (generated by input system)",
            )),
            MessageType::GamepadRumbleRequest => Err(PyTypeError::new_err(
                "GamepadRumbleRequest events are write-only via rumble API",
            )),
            MessageType::WindowEvent => Err(PyTypeError::new_err(
                "WindowEvent events are read-only (generated by window system)",
            )),
            MessageType::SceneInstanceReady => Err(PyTypeError::new_err(
                "SceneInstanceReady events are read-only (generated by scene system)",
            )),
            MessageType::AssetEventImage => Err(PyTypeError::new_err(
                "AssetEvent<Image> not yet implemented",
            )),
            MessageType::AssetEventMesh => {
                Err(PyTypeError::new_err("AssetEvent<Mesh> not yet implemented"))
            }
            MessageType::Custom(py_type) => {
                // Verify the message is an instance of the registered type
                let msg_type = bound_message.get_type();
                let expected_type = py_type.bind(py);
                if !msg_type.is(&expected_type) {
                    return Err(PyTypeError::new_err(format!(
                        "Expected message of type {}, got {}",
                        expected_type.name()?,
                        msg_type.name()?,
                    )));
                }

                // Look up the message number from registry
                let world = self.world.world_mut()?;
                let type_ptr = expected_type.as_type_ptr();
                let message_num = {
                    let registry = world.resource::<MessageRegistry>();
                    registry
                        .get(type_ptr)
                        .ok_or_else(|| PyTypeError::new_err("Message type not registered"))?
                        .message_num
                };

                // Write to the appropriate CustomMessageN resource
                let event_id: Box<dyn std::any::Any + Send + Sync> = write_message_arms!(message_num, world, py, message,
                    0=>1, 1=>2, 2=>3, 3=>4, 4=>5, 5=>6, 6=>7, 7=>8, 8=>9, 9=>10,
                    10=>11, 11=>12, 12=>13, 13=>14, 14=>15, 15=>16, 16=>17, 17=>18, 18=>19, 19=>20
                );

                Ok(PyMessageId::from_boxed(event_id))
            }
            MessageType::Dynamic(type_ptr) => {
                // Use global registry for dynamic dispatch (most message types use this)
                let bridge =
                    global_registry::get_message_bridge_by_py_type(*type_ptr).ok_or_else(|| {
                        PyTypeError::new_err("Message type not registered in global registry")
                    })?;

                if bridge.is_read_only() {
                    return Err(PyTypeError::new_err(format!(
                        "{} messages are read-only (generated by system)",
                        bridge.name()
                    )));
                }

                let world = self.world.world_mut()?;
                let event_id = bridge.write_message(py, world, &bound_message)?;
                Ok(PyMessageId::from_boxed(event_id))
            }
        }
    }

    pub fn write_batch(&self, py: Python, messages: Vec<Py<PyAny>>) -> PyResult<Vec<PyMessageId>> {
        let mut message_ids = Vec::new();
        for message in messages {
            let message_id = self.write(py, message)?;
            message_ids.push(message_id);
        }
        Ok(message_ids)
    }

    pub fn write_default(&self, py: Python) -> PyResult<PyMessageId> {
        match &self.message_type {
            MessageType::KeyboardInput => Err(PyTypeError::new_err(
                "KeyboardInput events are read-only (generated by input system)",
            )),
            MessageType::GamepadRumbleRequest => Err(PyTypeError::new_err(
                "GamepadRumbleRequest events are read-only (generated by gamepad system)",
            )),
            MessageType::WindowEvent => Err(PyTypeError::new_err(
                "WindowEvent events are read-only (generated by window system)",
            )),
            MessageType::SceneInstanceReady => Err(PyTypeError::new_err(
                "SceneInstanceReady events are read-only (generated by scene system)",
            )),
            MessageType::AssetEventImage => Err(PyTypeError::new_err(
                "AssetEvent<Image> has no default value",
            )),
            MessageType::AssetEventMesh => Err(PyTypeError::new_err(
                "AssetEvent<Mesh> has no default value",
            )),
            MessageType::Custom(type_obj) => {
                // Call the type with no arguments to create a default instance
                let instance = type_obj.bind(py).call0()?;
                self.write(py, instance.unbind())
            }
            MessageType::Dynamic(type_ptr) => {
                // Use global registry for dynamic dispatch
                let bridge =
                    global_registry::get_message_bridge_by_py_type(*type_ptr).ok_or_else(|| {
                        PyTypeError::new_err("Message type not registered in global registry")
                    })?;

                if bridge.is_read_only() {
                    return Err(PyTypeError::new_err(format!(
                        "{} messages are read-only (generated by system)",
                        bridge.name()
                    )));
                }

                // Dynamic types don't support write_default - they need explicit values
                Err(PyTypeError::new_err(format!(
                    "{} messages don't support write_default",
                    bridge.name()
                )))
            }
        }
    }
}

/// Read-only handle to a message channel. Used as a system parameter to receive messages.
///
/// Usage: `def my_system(reader: MessageReader[DamageEvent])`
#[pyclass(name = "MessageReader", frozen)]
pub struct PyMessageReader {
    #[allow(dead_code)] // retained for potential future direct-world access
    pub(crate) world: PyWorld,
    pub(crate) messages: PyMessages,
}

impl PyMessageReader {
    // TODO: Add constructor when needed by system integration
}

#[pymethods]
impl PyMessageReader {
    #[classmethod]
    #[pyo3(signature = (key, /))]
    pub fn __class_getitem__(
        _cls: &Bound<'_, PyType>,
        key: &Bound<'_, PyAny>,
    ) -> PyResult<Py<PyAny>> {
        let py = key.py();

        // Extract message type from key (e.g., AppExit from MessageReader[AppExit])
        let type_obj = key.cast::<PyType>()?;
        let py_message_type = PyMessageType::from_message_type(&type_obj)?;

        // Create MessageTypeParam with Reader class
        let param = MessageTypeParam {
            ty: MessageClass::Reader,
            message_type: py_message_type.0,
        };

        param.into_py_any(py)
    }

    pub fn clear(&self) -> PyResult<()> {
        self.messages.clear()
    }

    pub fn is_empty(&self) -> PyResult<bool> {
        self.messages.is_empty()
    }

    pub fn len(&self) -> PyResult<usize> {
        self.messages.len()
    }

    pub fn read(&self, py: Python) -> PyResult<PyMessageReaderIter> {
        // Collect all messages at iterator creation time using persistent cursor
        let cached_messages = self.messages.iter_to_python(py)?;

        Ok(PyMessageReaderIter {
            cursor: AtomicUsize::new(0),
            cached_messages,
        })
    }

    pub fn __iter__(&self, py: Python) -> PyResult<PyMessageReaderIter> {
        self.read(py)
    }
}

#[pyclass(name = "MessageReaderIter")]
pub struct PyMessageReaderIter {
    cursor: AtomicUsize,
    // Cache all messages at creation time for iteration
    cached_messages: Vec<Py<PyAny>>,
}

#[pymethods]
impl PyMessageReaderIter {
    pub fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
        slf
    }

    pub fn __next__(&self, py: Python) -> PyResult<Py<PyAny>> {
        let index = self.cursor.fetch_add(1, Ordering::Relaxed);

        if index < self.cached_messages.len() {
            Ok(self.cached_messages[index].clone_ref(py))
        } else {
            Err(PyStopIteration::new_err(""))
        }
    }
}

/// Write a Python message to the appropriate CustomMessage slot.
/// Called from external crates via `global_registry::write_python_message`.
fn write_custom_python_message(
    world: &mut bevy::ecs::world::World,
    py: Python,
    msg: &Bound<'_, PyAny>,
) -> PyResult<()> {
    use pyo3::types::PyTypeMethods;

    let type_ptr = msg.get_type().as_type_ptr();
    let message_num = {
        let registry = world
            .get_resource::<MessageRegistry>()
            .ok_or_else(|| PyTypeError::new_err("MessageRegistry not found in world"))?;
        registry
            .get(type_ptr)
            .ok_or_else(|| PyTypeError::new_err("Message type not registered"))?
            .message_num
    };

    // Write to the appropriate CustomMessageN slot
    let msg_py: Py<PyAny> = msg.clone().unbind();
    let _: Box<dyn std::any::Any + Send + Sync> = write_message_arms!(message_num, world, py, msg_py,
        0=>1, 1=>2, 2=>3, 3=>4, 4=>5, 5=>6, 6=>7, 7=>8, 8=>9, 9=>10,
        10=>11, 11=>12, 12=>13, 13=>14, 14=>15, 15=>16, 16=>17, 17=>18, 18=>19, 19=>20
    );
    Ok(())
}

/// Register the message write function in the global registry.
/// Called once at module init.
pub(crate) fn register_message_write_fn() {
    global_registry::register_message_write_fn(|world, py, msg| {
        write_custom_python_message(world, py, msg).map_err(|e| format!("{e}"))
    });
}

/// Register the run_system_once function in the global registry.
/// Called once at module init.
pub(crate) fn register_run_system_once_fn() {
    global_registry::register_run_system_once_fn(|world, py, func| {
        PyWorld::with_temporary(world, py, |py_world| py_world.run_system_once(func.clone()))
            .map_err(|e| format!("{e}"))
    });
}