plexor-core 0.1.0-alpha.2

// Copyright 2025 Alecks Gates
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

//! Axon module for handling neuron-ganglion interactions with type erasure.
//!
//! This module provides the Axon trait and AxonImpl struct that manage the interface
//! between neurons and ganglia, handling type erasure and providing a clean API
//! for neuron operations.

use crate::codec::{Codec, CodecName};
use crate::erasure::payload::erase_payload;
use crate::erasure::reactant::ReactantErased;
use crate::ganglion::{GanglionError, GanglionInternal};
use crate::logging::TraceContext;
use crate::neuron::Neuron;
use crate::payload::Payload;
use crate::reactant::Reactant;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use std::time::Duration;
use thiserror::Error;
use tokio::sync::Mutex;
use tokio::time::timeout;
use tracing::Instrument;
use uuid::Uuid;

/// Error types for Axon operations
#[derive(Error, Debug)]
pub enum AxonError {
    #[error(
        "Neuron needs to be adapted to ganglion {ganglion_name} (id: {ganglion_id}): {neuron_name}"
    )]
    NeuronNotAdapted {
        neuron_name: String,
        ganglion_name: String,
        ganglion_id: Uuid,
    },
    #[error(
        "Failed to acquire lock on synapse: {neuron_name} (ganglion: {ganglion_name}, id: {ganglion_id})"
    )]
    SynapseLock {
        neuron_name: String,
        ganglion_name: String,
        ganglion_id: Uuid,
    },
    #[error(
        "Transmission error for neuron {neuron_name} (ganglion: {ganglion_name}, id: {ganglion_id}): {message}"
    )]
    Transmit {
        neuron_name: String,
        ganglion_name: String,
        ganglion_id: Uuid,
        message: String,
    },
    #[error("Encode error for neuron {neuron_name} (ganglion: {ganglion_name}, id: {ganglion_id})")]
    Encode {
        neuron_name: String,
        ganglion_name: String,
        ganglion_id: Uuid,
    },
    #[error("Decode error for neuron {neuron_name} (ganglion: {ganglion_name}, id: {ganglion_id})")]
    Decode {
        neuron_name: String,
        ganglion_name: String,
        ganglion_id: Uuid,
    },
    #[error("Transmission timeout")]
    TransmissionTimeout,
    #[error("Ganglion error: {0}")]
    GanglionError(#[from] GanglionError),
}

impl AxonError {
    /// Convert a GanglionError to AxonError, translating SynapseNotFound appropriately
    pub fn from_ganglion_error(error: GanglionError) -> Self {
        match error {
            GanglionError::SynapseNotFound {
                neuron_name,
                ganglion_name,
                ganglion_id,
            } => AxonError::NeuronNotAdapted {
                neuron_name,
                ganglion_name,
                ganglion_id,
            },
            GanglionError::SynapseLock {
                neuron_name,
                ganglion_name,
                ganglion_id,
            } => AxonError::SynapseLock {
                neuron_name,
                ganglion_name,
                ganglion_id,
            },
            GanglionError::Transmit {
                neuron_name,
                ganglion_name,
                ganglion_id,
                message,
            } => AxonError::Transmit {
                neuron_name,
                ganglion_name,
                ganglion_id,
                message,
            },
            GanglionError::Encode {
                neuron_name,
                ganglion_name,
                ganglion_id,
            } => AxonError::Encode {
                neuron_name,
                ganglion_name,
                ganglion_id,
            },
            GanglionError::Decode {
                neuron_name,
                ganglion_name,
                ganglion_id,
            } => AxonError::Decode {
                neuron_name,
                ganglion_name,
                ganglion_id,
            },
            GanglionError::Adapt { .. } | GanglionError::QueueFull { .. } => {
                // For adaptation and queue full errors, use the generic GanglionError wrapper
                AxonError::GanglionError(error)
            }
        }
    }
}

/// Trait for axon operations with typed neurons and data
pub trait Axon<T, C>: Send + Sync
where
    T: Send + Sync + 'static,
    C: Codec<T> + CodecName + Send + Sync + 'static,
{
        /// React to incoming reactants using the neuron
        fn react(
            &mut self,
            reactants: Vec<Box<dyn Reactant<T, C> + Send + Sync + 'static>>,
        ) -> Pin<Box<dyn Future<Output = Result<(), AxonError>> + Send + '_>>;
     
        /// Transmit data through the neuron and ganglion
        ///
        /// # Examples
        ///
        /// ```rust
        /// # use std::sync::Arc;
        /// # use tokio::sync::Mutex;
        /// # use plexor_core::axon::{Axon, AxonImpl};
        /// # use plexor_core::neuron::{Neuron, NeuronImpl};
        /// # use plexor_core::plexus::Plexus;
        /// # use plexor_core::namespace::NamespaceImpl;
        /// # use plexor_core::ganglion::Ganglion;
        /// # use plexor_core::codec::{Codec, CodecError, CodecName};
        /// #
        /// # #[derive(Debug, Clone)]
        /// # struct Dummy;
        /// # impl CodecName for Dummy { fn name() -> &'static str { "dummy" } }
        /// # impl Codec<Dummy> for Dummy {
        /// #     fn encode(_: &Dummy) -> Result<Vec<u8>, CodecError> { Ok(vec![]) }
        /// #     fn decode(_: &[u8]) -> Result<Dummy, CodecError> { Ok(Dummy) }
        /// # }
        /// #
        /// # tokio::runtime::Runtime::new().unwrap().block_on(async {
        /// let ns = Arc::new(NamespaceImpl { delimiter: ".", parts: vec!["test"] });
        /// let neuron = Arc::new(NeuronImpl::<Dummy, Dummy>::new(ns));
        /// let plexus = Plexus::new_shared(vec![], vec![]).await;
        /// plexus.lock().await.adapt(neuron.clone()).await.unwrap();
        ///
        /// let mut axon = AxonImpl::new(neuron, plexus);
        /// axon.transmit(Dummy).await.unwrap();
        /// # });
        /// ```
        fn transmit(
            &mut self,
            data: T,
        ) -> Pin<Box<dyn Future<Output = Result<Vec<()>, AxonError>> + Send + '_>>;
    /// Transmit data with explicit tracing context.
    fn transmit_with_trace(
        &mut self,
        data: T,
        trace_context: TraceContext,
    ) -> Pin<Box<dyn Future<Output = Result<Vec<()>, AxonError>> + Send + '_>>;

    /// Transmit data through the neuron and ganglion using type erasure
    fn transmit_erased(
        &mut self,
        data: T,
    ) -> Pin<Box<dyn Future<Output = Result<Vec<()>, AxonError>> + Send + '_>>;

    /// Transmit data using type erasure with explicit tracing context.
    fn transmit_erased_with_trace(
        &mut self,
        data: T,
        trace_context: TraceContext,
    ) -> Pin<Box<dyn Future<Output = Result<Vec<()>, AxonError>> + Send + '_>>;

    /// Get the neuron name
    fn neuron_name(&self) -> String;

    /// Transmit data with a timeout
    fn transmit_with_timeout(
        &mut self,
        data: T,
        timeout_duration: Duration,
    ) -> Pin<Box<dyn Future<Output = Result<Vec<()>, AxonError>> + Send + '_>> {
        Box::pin(async move {
            match timeout(timeout_duration, self.transmit(data)).await {
                Ok(result) => result,
                Err(_) => Err(AxonError::TransmissionTimeout),
            }
        })
    }

    /// Transmit data with a trace and timeout
    fn transmit_with_trace_with_timeout(
        &mut self,
        data: T,
        trace_context: TraceContext,
        timeout_duration: Duration,
    ) -> Pin<Box<dyn Future<Output = Result<Vec<()>, AxonError>> + Send + '_>> {
        Box::pin(async move {
            match timeout(timeout_duration, self.transmit_with_trace(data, trace_context)).await {
                Ok(result) => result,
                Err(_) => Err(AxonError::TransmissionTimeout),
            }
        })
    }

    /// Transmit multiple data items in batch
    #[allow(clippy::type_complexity)]
    fn transmit_batch(
        &mut self,
        data_items: Vec<T>,
    ) -> Pin<Box<dyn Future<Output = Result<Vec<Vec<()>>, AxonError>> + Send + '_>> {
        Box::pin(async move {
            let mut results = Vec::new();
            for item in data_items {
                match self.transmit(item).await {
                    Ok(result) => results.push(result),
                    Err(e) => return Err(e),
                }
            }
            Ok(results)
        })
    }

    /// Transmit multiple data items in batch with trace
    #[allow(clippy::type_complexity)]
    fn transmit_batch_with_trace(
        &mut self,
        data_items: Vec<T>,
        trace_context: TraceContext,
    ) -> Pin<Box<dyn Future<Output = Result<Vec<Vec<()>>, AxonError>> + Send + '_>> {
        Box::pin(async move {
            let mut results = Vec::new();
            for item in data_items {
                match self.transmit_with_trace(item, trace_context).await {
                    Ok(result) => results.push(result),
                    Err(e) => return Err(e),
                }
            }
            Ok(results)
        })
    }

    /// Check if the axon is ready for transmission (convenience method)
    fn is_ready(&self) -> bool {
        !self.neuron_name().is_empty()
    }

    /// Get a formatted status string for debugging
    fn status(&self) -> String {
        format!("Axon[neuron: {}]", self.neuron_name())
    }
}

/// Implementation of Axon that stores a neuron and ganglion with type erasure handling
///
/// # Examples
///
/// ```rust
/// # use std::sync::Arc;
/// # use tokio::sync::Mutex;
/// # use plexor_core::axon::{Axon, AxonImpl};
/// # use plexor_core::neuron::{Neuron, NeuronImpl};
/// # use plexor_core::ganglion::{Ganglion, GanglionInprocess};
/// # use plexor_core::namespace::NamespaceImpl;
/// # use plexor_core::codec::{Codec, CodecError, CodecName};
/// #
/// # #[derive(Debug, Clone)]
/// # struct Dummy;
/// # impl CodecName for Dummy { fn name() -> &'static str { "dummy" } }
/// # impl Codec<Dummy> for Dummy {
/// #     fn encode(_: &Dummy) -> Result<Vec<u8>, CodecError> { Ok(vec![]) }
/// #     fn decode(_: &[u8]) -> Result<Dummy, CodecError> { Ok(Dummy) }
/// # }
/// #
/// # tokio::runtime::Runtime::new().unwrap().block_on(async {
/// let ns = Arc::new(NamespaceImpl { delimiter: ".", parts: vec!["test"] });
/// let neuron = Arc::new(NeuronImpl::<Dummy, Dummy>::new(ns));
/// let ganglion = Arc::new(Mutex::new(GanglionInprocess::new()));
///
/// let mut axon = AxonImpl::new(neuron, ganglion);
/// assert!(axon.is_ready());
/// # });
/// ```
pub struct AxonImpl<T, C>
where
    T: Send + Sync + 'static,
    C: Codec<T> + CodecName + Send + Sync + 'static,
{
    /// The neuron this axon manages
    neuron: Arc<dyn Neuron<T, C> + Send + Sync + 'static>,

    /// The ganglion this axon uses for operations
    ganglion: Arc<Mutex<dyn GanglionInternal + Send + Sync + 'static>>,
}

impl<T, C> Clone for AxonImpl<T, C>
where
    T: Send + Sync + 'static,
    C: Codec<T> + CodecName + Send + Sync + 'static,
{
    fn clone(&self) -> Self {
        Self {
            neuron: self.neuron.clone(),
            ganglion: self.ganglion.clone(),
        }
    }
}

impl<T, C> AxonImpl<T, C>
where
    T: Send + Sync + 'static,
    C: Codec<T> + CodecName + Send + Sync + 'static,
{
    /// Create a new AxonImpl with the given neuron and ganglion
    pub fn new(
        neuron: Arc<dyn Neuron<T, C> + Send + Sync + 'static>,
        ganglion: Arc<Mutex<dyn GanglionInternal + Send + Sync + 'static>>,
    ) -> Self {
        Self { neuron, ganglion }
    }

    /// Create a builder for constructing an AxonImpl
    pub fn builder() -> AxonBuilder<T, C> {
        AxonBuilder::new()
    }

    /// Get a reference to the underlying neuron
    pub fn neuron(&self) -> &Arc<dyn Neuron<T, C> + Send + Sync + 'static> {
        &self.neuron
    }

    /// Get a reference to the underlying ganglion
    pub fn ganglion(&self) -> &Arc<Mutex<dyn GanglionInternal + Send + Sync + 'static>> {
        &self.ganglion
    }

    /// Clone the neuron for sharing
    pub fn clone_neuron(&self) -> Arc<dyn Neuron<T, C> + Send + Sync + 'static> {
        self.neuron.clone()
    }

    /// Clone the ganglion for sharing
    pub fn clone_ganglion(&self) -> Arc<Mutex<dyn GanglionInternal + Send + Sync + 'static>> {
        self.ganglion.clone()
    }

    /// Get the unique ID of the underlying ganglion
    pub async fn ganglion_id(&self) -> Uuid {
        self.ganglion.lock().await.unique_id()
    }

    /// Check if the axon components are properly configured
    pub fn validate(&self) -> Result<(), String> {
        if self.neuron.name().is_empty() {
            return Err("Neuron name is empty".to_string());
        }
        Ok(())
    }
}

/// Builder for constructing AxonImpl instances with a fluent API
pub struct AxonBuilder<T, C>
where
    T: Send + Sync + 'static,
    C: Codec<T> + CodecName + Send + Sync + 'static,
{
    neuron: Option<Arc<dyn Neuron<T, C> + Send + Sync + 'static>>,
    ganglion: Option<Arc<Mutex<dyn GanglionInternal + Send + Sync + 'static>>>,
}

impl<T, C> AxonBuilder<T, C>
where
    T: Send + Sync + 'static,
    C: Codec<T> + CodecName + Send + Sync + 'static,
{
    /// Create a new AxonBuilder
    pub fn new() -> Self {
        Self {
            neuron: None,
            ganglion: None,
        }
    }

    /// Set the neuron for the axon
    pub fn with_neuron(mut self, neuron: Arc<dyn Neuron<T, C> + Send + Sync + 'static>) -> Self {
        self.neuron = Some(neuron);
        self
    }

    /// Set the ganglion for the axon
    pub fn with_ganglion(
        mut self,
        ganglion: Arc<Mutex<dyn GanglionInternal + Send + Sync + 'static>>,
    ) -> Self {
        self.ganglion = Some(ganglion);
        self
    }

    /// Build the AxonImpl instance
    pub fn build(self) -> Result<AxonImpl<T, C>, String> {
        let neuron = self.neuron.ok_or("Neuron is required")?;
        let ganglion = self.ganglion.ok_or("Ganglion is required")?;

        let axon = AxonImpl::new(neuron, ganglion);
        axon.validate()?;
        Ok(axon)
    }

    /// Build the AxonImpl instance without validation
    pub fn build_unchecked(self) -> Result<AxonImpl<T, C>, String> {
        let neuron = self.neuron.ok_or("Neuron is required")?;
        let ganglion = self.ganglion.ok_or("Ganglion is required")?;

        Ok(AxonImpl::new(neuron, ganglion))
    }
}

impl<T, C> Default for AxonBuilder<T, C>
where
    T: Send + Sync + 'static,
    C: Codec<T> + CodecName + Send + Sync + 'static,
{
    fn default() -> Self {
        Self::new()
    }
}

impl<T, C> Axon<T, C> for AxonImpl<T, C>
where
    T: Send + Sync + 'static,
    C: Codec<T> + CodecName + Send + Sync + 'static,
{
    fn react(
        &mut self,
        reactants: Vec<Box<dyn Reactant<T, C> + Send + Sync + 'static>>,
    ) -> Pin<Box<dyn Future<Output = Result<(), AxonError>> + Send + '_>> {
        let neuron_name = self.neuron.name();
        let ganglion = self.ganglion.clone();

        Box::pin(async move {
            // Convert typed reactants to erased reactants
            let erased_reactants: Vec<Arc<dyn ReactantErased + Send + Sync + 'static>> = reactants
                .into_iter()
                .map(|reactant| reactant.erase())
                .collect();

            let future = {
                let mut ganglion_guard = ganglion.lock().await;
                ganglion_guard.react(neuron_name, erased_reactants, vec![])
            };
            match future.await {
                Ok(result) => Ok(result),
                Err(ganglion_error) => Err(AxonError::from_ganglion_error(ganglion_error)),
            }
        })
    }

    fn transmit(
        &mut self,
        data: T,
    ) -> Pin<Box<dyn Future<Output = Result<Vec<()>, AxonError>> + Send + '_>> {
        let neuron = self.neuron.clone();
        let ganglion = self.ganglion.clone();

        Box::pin(async move {
            // Create a payload
            let payload = Payload::new(data, neuron.clone());

            // Transmit through the ganglion
            let future = {
                let mut ganglion_guard = ganglion.lock().await;
                ganglion_guard.transmit(erase_payload(payload))
            };
            match future.await {
                Ok(result) => Ok(result),
                Err(ganglion_error) => Err(AxonError::from_ganglion_error(ganglion_error)),
            }
        })
    }

    fn transmit_with_trace(
        &mut self,
        data: T,
        trace_context: TraceContext,
    ) -> Pin<Box<dyn Future<Output = Result<Vec<()>, AxonError>> + Send + '_>> {
        let neuron = self.neuron.clone();
        let ganglion = self.ganglion.clone();

        Box::pin(async move {
            // Create a payload with explicit trace info
            let payload = Arc::new(Payload::from_parts(
                Arc::new(data),
                neuron.clone(),
                trace_context.new_child(),
            ));

            // Transmit through the ganglion
            let future = {
                let mut ganglion_guard = ganglion.lock().await;
                ganglion_guard.transmit(erase_payload(payload))
            };
            match future.await {
                Ok(result) => Ok(result),
                Err(ganglion_error) => Err(AxonError::from_ganglion_error(ganglion_error)),
            }
        })
    }

    fn transmit_erased(
        &mut self,
        data: T,
    ) -> Pin<Box<dyn Future<Output = Result<Vec<()>, AxonError>> + Send + '_>> {
        let neuron = self.neuron.clone();
        let ganglion = self.ganglion.clone();

        Box::pin(
            async move {
                // Create a payload
                let payload = Payload::new(data, neuron.clone());
                let erased_payload = erase_payload(payload);
                let span = erased_payload.span_debug("AxonImpl::transmit_erased");

                async move {
                    tracing::debug!("Starting erased transmission");
                    let future = {
                        let mut ganglion_guard = ganglion.lock().await;
                        ganglion_guard.transmit(erased_payload)
                    };
                    match future.await {
                        Ok(result) => Ok(result),
                        Err(ganglion_error) => Err(AxonError::from_ganglion_error(ganglion_error)),
                    }
                }
                .instrument(span)
                .await
            },
        )
    }

    fn transmit_erased_with_trace(
        &mut self,
        data: T,
        trace_context: TraceContext,
    ) -> Pin<Box<dyn Future<Output = Result<Vec<()>, AxonError>> + Send + '_>> {
        let neuron = self.neuron.clone();
        let ganglion = self.ganglion.clone();

        Box::pin(async move {
            // Create a payload with explicit trace info
            let payload = Arc::new(Payload::from_parts(
                Arc::new(data),
                neuron.clone(),
                trace_context.new_child(),
            ));

            // Convert to erased payload
            let erased_payload = erase_payload(payload);
            let span = erased_payload.span_debug("AxonImpl::transmit_erased_with_trace");

            async move {
                tracing::debug!("Starting erased transmission with trace");

                let future = {
                    let mut ganglion_guard = ganglion.lock().await;
                    ganglion_guard.transmit(erased_payload)
                };
                match future.await {
                    Ok(result) => Ok(result),
                    Err(ganglion_error) => Err(AxonError::from_ganglion_error(ganglion_error)),
                }
            }
            .instrument(span)
            .await
        })
    }

    fn neuron_name(&self) -> String {
        self.neuron.name()
    }
}

impl<T, C> std::fmt::Debug for AxonImpl<T, C>
where
    T: Send + Sync + 'static,
    C: Codec<T> + CodecName + Send + Sync + 'static,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("AxonImpl")
            .field("neuron", &self.neuron.name())
            .finish()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ganglion::{Ganglion, GanglionInprocess};
    use crate::neuron::NeuronImpl;
    use crate::payload::Payload;
    use crate::plexus::Plexus;
    use crate::test_utils::{DebugCodec, DebugStruct, TokioMpscReactant, test_namespace};
    use tokio::sync::{Mutex, mpsc::channel};
    use uuid::Uuid;

    #[tokio::test]
    async fn test_axon_creation() {
        let ns = test_namespace();
        let neuron: Arc<dyn Neuron<DebugStruct, DebugCodec> + Send + Sync> =
            Arc::new(NeuronImpl::new(ns));
        let ganglion: Arc<Mutex<dyn GanglionInternal + Send + Sync>> =
            Arc::new(Mutex::new(GanglionInprocess::new()));

        let axon = AxonImpl::new(neuron, ganglion);

        assert_eq!(axon.neuron_name(), "dev.plexo.DebugStruct.debug");
    }

    #[tokio::test]
    async fn test_axon_transmit_with_plexus() {
        let ns = test_namespace();
        let neuron_impl: NeuronImpl<DebugStruct, DebugCodec> = NeuronImpl::new(ns);
        let neuron_arc = neuron_impl.clone_to_arc();

        // Create a plexus to use as the ganglion
        let mut plexus = Plexus::new(vec![], vec![]).await;

        // Adapt the neuron to the plexus
        plexus
            .adapt(neuron_arc.clone())
            .await
            .expect("Failed to adapt neuron");

        // Create channels for receiving transmitted data
        let (tx1, mut rx1) = channel::<Arc<Payload<DebugStruct, DebugCodec>>>(2);
        let (tx2, mut rx2) = channel::<Arc<Payload<DebugStruct, DebugCodec>>>(2);

        // Create axon with the plexus as ganglion
        let plexus_ganglion: Arc<Mutex<dyn GanglionInternal + Send + Sync>> =
            Arc::new(Mutex::new(plexus));
        let mut axon = AxonImpl::new(neuron_arc.clone(), plexus_ganglion);

        let reactants: Vec<Box<dyn Reactant<DebugStruct, DebugCodec> + Send + Sync + 'static>> = vec![
            Box::new(TokioMpscReactant::new(tx1)),
            Box::new(TokioMpscReactant::new(tx2)),
        ];

        axon.react(reactants).await.expect("Failed to react");

        // Create test data to transmit
        let test_data = DebugStruct {
            foo: 42,
            bar: "test transmission".to_string(),
        };

        // Transmit data through the axon
        let trace_context = TraceContext::new();
        let correlation_id = trace_context.correlation_id;
        axon.transmit_with_trace(test_data.clone(), trace_context)
            .await
            .expect("Failed to transmit");

        // Verify that the data was received by both reactants
        let received1 = rx1.recv().await.expect("Failed to receive payload 1");
        let received2 = rx2.recv().await.expect("Failed to receive payload 2");

        assert_eq!(received1.value.foo, 42);
        assert_eq!(received1.value.bar, "test transmission");
        assert_eq!(received2.value.foo, 42);
        assert_eq!(received2.value.bar, "test transmission");
        assert_eq!(received1.correlation_id(), correlation_id);
        assert_eq!(received2.correlation_id(), correlation_id);
    }

    #[tokio::test]
    async fn test_axon_transmit_erased() {
        let ns = test_namespace();
        let neuron_impl: NeuronImpl<DebugStruct, DebugCodec> = NeuronImpl::new(ns);
        let neuron_arc = neuron_impl.clone_to_arc();

        // Create a plexus to use as the ganglion
        let mut plexus = Plexus::new(vec![], vec![]).await;

        // Adapt the neuron to the plexus
        plexus
            .adapt(neuron_arc.clone())
            .await
            .expect("Failed to adapt neuron");

        // Create channel for receiving transmitted data
        let (tx, mut rx) = channel::<Arc<Payload<DebugStruct, DebugCodec>>>(1);

        // Create axon with the plexus as ganglion
        let plexus_ganglion: Arc<Mutex<dyn GanglionInternal + Send + Sync>> =
            Arc::new(Mutex::new(plexus));
        let mut axon = AxonImpl::new(neuron_arc.clone(), plexus_ganglion);

        let reactants: Vec<Box<dyn Reactant<DebugStruct, DebugCodec> + Send + Sync + 'static>> =
            vec![Box::new(TokioMpscReactant::new(tx))];

        axon.react(reactants).await.expect("Failed to react");

        // Create test data to transmit
        let test_data = DebugStruct {
            foo: 99,
            bar: "erased test".to_string(),
        };

        // Transmit data through the axon using transmit_erased
        let trace_context = TraceContext::new();
        let correlation_id = trace_context.correlation_id;
        axon.transmit_erased_with_trace(test_data.clone(), trace_context)
            .await
            .expect("Failed to transmit erased");

        // Verify that the data was received
        let received = rx.recv().await.expect("Failed to receive payload");

        assert_eq!(received.value.foo, 99);
        assert_eq!(received.value.bar, "erased test");
        assert_eq!(received.correlation_id(), correlation_id);
    }

    #[tokio::test]
    async fn test_axon_transmit_without_correlation_id() {
        let ns = test_namespace();
        let neuron_impl: NeuronImpl<DebugStruct, DebugCodec> = NeuronImpl::new(ns);
        let neuron_arc = neuron_impl.clone_to_arc();

        // Create a plexus to use as the ganglion
        let mut plexus = Plexus::new(vec![], vec![]).await;

        // Adapt the neuron to the plexus
        plexus
            .adapt(neuron_arc.clone())
            .await
            .expect("Failed to adapt neuron");

        // Create channel for receiving transmitted data
        let (tx, mut rx) = channel::<Arc<Payload<DebugStruct, DebugCodec>>>(1);

        // Create axon with the plexus as ganglion
        let plexus_ganglion: Arc<Mutex<dyn GanglionInternal + Send + Sync>> =
            Arc::new(Mutex::new(plexus));
        let mut axon = AxonImpl::new(neuron_arc.clone(), plexus_ganglion);

        let reactants: Vec<Box<dyn Reactant<DebugStruct, DebugCodec> + Send + Sync + 'static>> =
            vec![Box::new(TokioMpscReactant::new(tx))];

        axon.react(reactants).await.expect("Failed to react");

        // Create test data to transmit
        let test_data = DebugStruct {
            foo: 123,
            bar: "auto correlation".to_string(),
        };

        // Transmit data through the axon without providing correlation_id
        axon.transmit(test_data.clone())
            .await
            .expect("Failed to transmit");

        // Verify that the data was received and correlation_id was auto-generated
        let received = rx.recv().await.expect("Failed to receive payload");

        assert_eq!(received.value.foo, 123);
        assert_eq!(received.value.bar, "auto correlation");
        // Correlation ID should be auto-generated (not empty)
        assert_ne!(received.correlation_id(), Uuid::nil());
    }

    #[tokio::test]
    async fn test_axon_multiple_transmissions() {
        let ns = test_namespace();
        let neuron_impl: NeuronImpl<DebugStruct, DebugCodec> = NeuronImpl::new(ns);
        let neuron_arc = neuron_impl.clone_to_arc();

        // Create a plexus to use as the ganglion
        let mut plexus = Plexus::new(vec![], vec![]).await;

        // Adapt the neuron to the plexus
        plexus
            .adapt(neuron_arc.clone())
            .await
            .expect("Failed to adapt neuron");

        // Create channel for receiving transmitted data
        let (tx, mut rx) = channel::<Arc<Payload<DebugStruct, DebugCodec>>>(10);

        // Create axon with the plexus as ganglion
        let plexus_ganglion: Arc<Mutex<dyn GanglionInternal + Send + Sync>> =
            Arc::new(Mutex::new(plexus));
        let mut axon = AxonImpl::new(neuron_arc.clone(), plexus_ganglion);

        let reactants: Vec<Box<dyn Reactant<DebugStruct, DebugCodec> + Send + Sync + 'static>> =
            vec![Box::new(TokioMpscReactant::new(tx))];

        axon.react(reactants).await.expect("Failed to react");

        // Transmit multiple pieces of data
        for i in 0..5 {
            let test_data = DebugStruct {
                foo: i,
                bar: format!("transmission {i}"),
            };

            axon.transmit(test_data)
                .await
                .expect("Failed to transmit");
        }

        // Verify that all transmissions were received
        for i in 0..5 {
            let received = rx
                .recv()
                .await
                .unwrap_or_else(|| panic!("Failed to receive payload {i}"));

            assert_eq!(received.value.foo, i);
            assert_eq!(received.value.bar, format!("transmission {i}"));
        }
    }

    #[tokio::test]
    async fn test_axon_transmit_with_plexus_external_ganglion() {
        use crate::erasure::reactant::erase_reactant_raw;
        use crate::ganglion::GanglionExternal;
        use crate::neuron::NeuronImpl;
        use crate::payload::PayloadRaw;
        use crate::plexus::Plexus;
        use crate::test_utils::{
            DebugCodec, DebugStruct, GanglionExternalInprocess, TokioMpscReactantRaw,
            test_namespace,
        };
        use tokio::sync::mpsc::channel;
        use tokio::time::{Duration, sleep};

        let ns = test_namespace();
        let neuron_impl: NeuronImpl<DebugStruct, DebugCodec> = NeuronImpl::new(ns);
        let neuron_arc = neuron_impl.clone_to_arc();

        // Create a plexus to use as the ganglion
        let mut plexus = Plexus::new(vec![], vec![]).await;

        // Adapt the neuron to the plexus
        plexus
            .adapt(neuron_arc.clone())
            .await
            .expect("Failed to adapt neuron");

        // Create channel for receiving raw payloads from external ganglion
        let (tx_raw, mut rx_raw) = channel::<Arc<PayloadRaw<DebugStruct, DebugCodec>>>(2);

        let mut external_ganglion = GanglionExternalInprocess::new();

        // Adapt the neuron to the external ganglion
        external_ganglion
            .adapt(neuron_arc.clone())
            .await
            .expect("Failed to adapt neuron to external ganglion");

        // Create raw reactants and add them using react
        let raw_reactants = vec![erase_reactant_raw::<DebugStruct, DebugCodec, _>(Box::new(
            TokioMpscReactantRaw::new(tx_raw),
        ))];
        external_ganglion
            .react(neuron_arc.name(), vec![], raw_reactants, vec![])
            .await
            .expect("Failed to react with external ganglion");

        // Add the external ganglion to the plexus
        let external_ganglion_arc = Arc::new(Mutex::new(external_ganglion));
        let _ = plexus
            .infuse_external_ganglion(external_ganglion_arc.clone())
            .await;

        // Create axon with the plexus as ganglion
        let plexus_ganglion: Arc<Mutex<dyn GanglionInternal + Send + Sync>> =
            Arc::new(Mutex::new(plexus));
        let mut axon = AxonImpl::new(neuron_arc.clone(), plexus_ganglion);

        // No need to add reactants to the axon since we're testing the external ganglion
        let reactants: Vec<Box<dyn Reactant<DebugStruct, DebugCodec> + Send + Sync + 'static>> =
            vec![];
        axon.react(reactants).await.expect("Failed to react");

        // Create test data to transmit
        let test_data = DebugStruct {
            foo: 99,
            bar: "external ganglion test".to_string(),
        };

        // Transmit data through the axon
        let trace_context = TraceContext::new();
        let correlation_id = trace_context.correlation_id;
        axon.transmit_with_trace(test_data.clone(), trace_context)
            .await
            .expect("Failed to transmit");

        // Give some time for async transmission to complete
        sleep(Duration::from_millis(100)).await;

        // Verify that the external ganglion received the payloadraw
        assert_eq!(
            rx_raw.len(),
            1,
            "External ganglion should have received exactly one payload"
        );
        let received_raw = rx_raw
            .recv()
            .await
            .expect("Failed to receive raw payload from external ganglion");

        // Decode the raw payload to verify it contains the correct data
        let decoded_data = DebugCodec::decode(&received_raw.value).expect("Failed to decode");
        assert_eq!(decoded_data.foo, 99);
        assert_eq!(decoded_data.bar, "external ganglion test");
        assert_eq!(received_raw.correlation_id(), correlation_id);
    }

    #[tokio::test]
    async fn test_axon_transmit_simple() {
        let ns = test_namespace();
        let neuron_impl: NeuronImpl<DebugStruct, DebugCodec> = NeuronImpl::new(ns);
        let neuron_arc = neuron_impl.clone_to_arc();

        let mut plexus = Plexus::new(vec![], vec![]).await;
        plexus
            .adapt(neuron_arc.clone())
            .await
            .expect("Failed to adapt neuron");

        let (tx, mut rx) = channel::<Arc<Payload<DebugStruct, DebugCodec>>>(1);
        let plexus_ganglion: Arc<Mutex<dyn GanglionInternal + Send + Sync>> =
            Arc::new(Mutex::new(plexus));
        let mut axon = AxonImpl::new(neuron_arc.clone(), plexus_ganglion);

        let reactants: Vec<Box<dyn Reactant<DebugStruct, DebugCodec> + Send + Sync + 'static>> =
            vec![Box::new(TokioMpscReactant::new(tx))];
        axon.react(reactants).await.expect("Failed to react");

        let test_data = DebugStruct {
            foo: 555,
            bar: "simple transmission".to_string(),
        };

        // Test transmit convenience method
        axon.transmit(test_data.clone())
            .await
            .expect("Failed to transmit simple");

        let received = rx.recv().await.expect("Failed to receive payload");
        assert_eq!(received.value.foo, 555);
        assert_eq!(received.value.bar, "simple transmission");
        assert_ne!(received.correlation_id(), Uuid::nil());
    }

    #[tokio::test]
    async fn test_axon_transmit_with_timeout() {
        let ns = test_namespace();
        let neuron_impl: NeuronImpl<DebugStruct, DebugCodec> = NeuronImpl::new(ns);
        let neuron_arc = neuron_impl.clone_to_arc();

        let mut plexus = Plexus::new(vec![], vec![]).await;
        plexus
            .adapt(neuron_arc.clone())
            .await
            .expect("Failed to adapt neuron");

        let (tx, mut rx) = channel::<Arc<Payload<DebugStruct, DebugCodec>>>(1);
        let plexus_ganglion: Arc<Mutex<dyn GanglionInternal + Send + Sync>> =
            Arc::new(Mutex::new(plexus));
        let mut axon = AxonImpl::new(neuron_arc.clone(), plexus_ganglion);

        let reactants: Vec<Box<dyn Reactant<DebugStruct, DebugCodec> + Send + Sync + 'static>> =
            vec![Box::new(TokioMpscReactant::new(tx))];
        axon.react(reactants).await.expect("Failed to react");

        let test_data = DebugStruct {
            foo: 777,
            bar: "timeout test".to_string(),
        };

        // Test transmit_with_timeout with a reasonable timeout
        let timeout_duration = Duration::from_secs(1);
        axon.transmit_with_timeout(test_data.clone(), timeout_duration)
            .await
            .expect("Failed to transmit with timeout");

        let received = rx.recv().await.expect("Failed to receive payload");
        assert_eq!(received.value.foo, 777);
        assert_eq!(received.value.bar, "timeout test");
    }

    #[tokio::test]
    async fn test_axon_transmit_batch() {
        let ns = test_namespace();
        let neuron_impl: NeuronImpl<DebugStruct, DebugCodec> = NeuronImpl::new(ns);
        let neuron_arc = neuron_impl.clone_to_arc();

        let mut plexus = Plexus::new(vec![], vec![]).await;
        plexus
            .adapt(neuron_arc.clone())
            .await
            .expect("Failed to adapt neuron");

        let (tx, mut rx) = channel::<Arc<Payload<DebugStruct, DebugCodec>>>(10);
        let plexus_ganglion: Arc<Mutex<dyn GanglionInternal + Send + Sync>> =
            Arc::new(Mutex::new(plexus));
        let mut axon = AxonImpl::new(neuron_arc.clone(), plexus_ganglion);

        let reactants: Vec<Box<dyn Reactant<DebugStruct, DebugCodec> + Send + Sync + 'static>> =
            vec![Box::new(TokioMpscReactant::new(tx))];
        axon.react(reactants).await.expect("Failed to react");

        // Create batch data
        let batch_data = vec![
            DebugStruct {
                foo: 1,
                bar: "batch item 1".to_string(),
            },
            DebugStruct {
                foo: 2,
                bar: "batch item 2".to_string(),
            },
            DebugStruct {
                foo: 3,
                bar: "batch item 3".to_string(),
            },
        ];

        // Test transmit_batch
        let results = axon
            .transmit_batch(batch_data.clone())
            .await
            .expect("Failed to transmit batch");
        assert_eq!(results.len(), 3);

        // Verify all items were received
        for i in 0..3 {
            let received = rx
                .recv()
                .await
                .unwrap_or_else(|| panic!("Failed to receive batch item {i}"));
            assert_eq!(received.value.foo, (i + 1));
            assert_eq!(received.value.bar, format!("batch item {}", i + 1));
            // Check that the parent_id is set since this is a batch transmission which might reuse context or generate new
            // In the implementation, transmit_batch calls transmit_simple which calls transmit(data, None).
            // So parent_id should be None and correlation_id should be set.
            assert_ne!(received.correlation_id(), Uuid::nil());
        }
    }

    #[tokio::test]
    async fn test_axon_status_and_validation() {
        let ns = test_namespace();
        let neuron_impl: NeuronImpl<DebugStruct, DebugCodec> = NeuronImpl::new(ns);
        let neuron_arc = neuron_impl.clone_to_arc();

        let ganglion: Arc<Mutex<dyn GanglionInternal + Send + Sync>> =
            Arc::new(Mutex::new(GanglionInprocess::new()));
        let axon = AxonImpl::new(neuron_arc, ganglion);

        // Test is_ready
        assert!(axon.is_ready());

        // Test status
        let status = axon.status();
        assert!(status.contains("dev.plexo.DebugStruct.debug"));

        // Test validation
        assert!(axon.validate().is_ok());

        // Test neuron_name
        assert_eq!(axon.neuron_name(), "dev.plexo.DebugStruct.debug");
    }

    #[tokio::test]
    async fn test_axon_builder() {
        let ns = test_namespace();
        let neuron_impl: NeuronImpl<DebugStruct, DebugCodec> = NeuronImpl::new(ns);
        let neuron_arc = neuron_impl.clone_to_arc();

        let ganglion: Arc<Mutex<dyn GanglionInternal + Send + Sync>> =
            Arc::new(Mutex::new(GanglionInprocess::new()));

        // Test builder pattern
        let axon = AxonImpl::builder()
            .with_neuron(neuron_arc.clone())
            .with_ganglion(ganglion.clone())
            .build()
            .expect("Failed to build axon");

        assert_eq!(axon.neuron_name(), "dev.plexo.DebugStruct.debug");
        assert!(axon.is_ready());

        // Test builder with missing components
        let result = AxonImpl::<DebugStruct, DebugCodec>::builder()
            .with_neuron(neuron_arc.clone())
            .build();
        assert!(result.is_err());
        assert!(result.unwrap_err().contains("Ganglion is required"));

        // Test build_unchecked
        let axon_unchecked = AxonImpl::builder()
            .with_neuron(neuron_arc.clone())
            .with_ganglion(ganglion.clone())
            .build_unchecked()
            .expect("Failed to build unchecked axon");

        assert_eq!(axon_unchecked.neuron_name(), "dev.plexo.DebugStruct.debug");
    }

    #[tokio::test]
    async fn test_axon_additional_methods() {
        let ns = test_namespace();
        let neuron_impl: NeuronImpl<DebugStruct, DebugCodec> = NeuronImpl::new(ns);
        let neuron_arc = neuron_impl.clone_to_arc();

        let ganglion: Arc<Mutex<dyn GanglionInternal + Send + Sync>> =
            Arc::new(Mutex::new(GanglionInprocess::new()));
        let axon = AxonImpl::new(neuron_arc.clone(), ganglion.clone());

        // Test clone methods
        let cloned_neuron = axon.clone_neuron();
        assert_eq!(cloned_neuron.name(), neuron_arc.name());

        // Test ganglion_id
        let ganglion_id = axon.ganglion_id().await;
        assert_ne!(ganglion_id, Uuid::nil());

        // Test neuron and ganglion getters
        assert_eq!(axon.neuron().name(), "dev.plexo.DebugStruct.debug");
    }

    #[test]
    fn test_axon_error_variants_and_conversion() {
        let ganglion_id = Uuid::now_v7();
        let ganglion_name = "TestGanglion".to_string();
        let neuron_name = "test_neuron".to_string();

        // Test GanglionError::SynapseNotFound -> AxonError::NeuronNotAdapted conversion
        let ganglion_synapse_not_found = GanglionError::SynapseNotFound {
            neuron_name: neuron_name.clone(),
            ganglion_name: ganglion_name.clone(),
            ganglion_id,
        };
        let axon_neuron_not_adapted = AxonError::from_ganglion_error(ganglion_synapse_not_found);
        match axon_neuron_not_adapted {
            AxonError::NeuronNotAdapted {
                neuron_name: n,
                ganglion_name: g,
                ganglion_id: id,
            } => {
                assert_eq!(n, neuron_name);
                assert_eq!(g, ganglion_name);
                assert_eq!(id, ganglion_id);
            }
            _ => panic!("Expected NeuronNotAdapted variant"),
        }

        // Test GanglionError::SynapseLockError -> AxonError::SynapseLockError conversion
        let ganglion_synapse_lock = GanglionError::SynapseLock {
            neuron_name: neuron_name.clone(),
            ganglion_name: ganglion_name.clone(),
            ganglion_id,
        };
        let axon_synapse_lock = AxonError::from_ganglion_error(ganglion_synapse_lock);
        match axon_synapse_lock {
            AxonError::SynapseLock {
                neuron_name: n,
                ganglion_name: g,
                ganglion_id: id,
            } => {
                assert_eq!(n, neuron_name);
                assert_eq!(g, ganglion_name);
                assert_eq!(id, ganglion_id);
            }
            _ => panic!("Expected SynapseLockError variant"),
        }

        // Test GanglionError::Transmit -> AxonError::Transmit conversion
        let ganglion_transmit = GanglionError::Transmit {
            neuron_name: neuron_name.clone(),
            ganglion_name: ganglion_name.clone(),
            ganglion_id,
            message: "test transmit error".to_string(),
        };
        let axon_transmit = AxonError::from_ganglion_error(ganglion_transmit);
        match axon_transmit {
            AxonError::Transmit {
                neuron_name: n,
                ganglion_name: g,
                ganglion_id: id,
                message: m,
            } => {
                assert_eq!(n, neuron_name);
                assert_eq!(g, ganglion_name);
                assert_eq!(id, ganglion_id);
                assert_eq!(m, "test transmit error");
            }
            _ => panic!("Expected Transmit variant"),
        }

        // Test direct AxonError creation
        let direct_neuron_not_adapted = AxonError::NeuronNotAdapted {
            neuron_name: neuron_name.clone(),
            ganglion_name: ganglion_name.clone(),
            ganglion_id,
        };
        assert!(direct_neuron_not_adapted.to_string().contains(&neuron_name));
        assert!(
            direct_neuron_not_adapted
                .to_string()
                .contains(&ganglion_name)
        );

        let direct_transmit = AxonError::Transmit {
            neuron_name: neuron_name.clone(),
            ganglion_name: ganglion_name.clone(),
            ganglion_id,
            message: "direct message".to_string(),
        };
        assert!(direct_transmit.to_string().contains(&neuron_name));
        assert!(direct_transmit.to_string().contains(&ganglion_name));
        assert!(direct_transmit.to_string().contains("direct message"));
    }
}