//! This module defines the [`GeneticNode`] trait and related types for managing the state and behavior of nodes within a [`Gemla`] simulation.
//!
//! - [`GeneticNode`]: A trait for interacting with the internal state of nodes, supporting asynchronous initialization, simulation, mutation, and merging.
//! - [`GeneticState`]: Enum representing the current state of a node in the simulation lifecycle.
//! - [`GeneticNodeContext`]: Context struct passed to node methods, containing generation, node ID, and simulation context.
//! - [`GeneticNodeWrapper`]: Wrapper struct for managing state transitions and node lifecycle.
//!
//! [`Gemla`]: crate::Gemla

use crate::error::Error;

use anyhow::Context;
use async_trait::async_trait;
use log::info;
use serde::{de::DeserializeOwned, Deserialize, Serialize};
use std::fmt::Debug;
use uuid::Uuid;

/// An enum representing the current state of a [`GeneticNode`] in the simulation lifecycle.
///
/// Used to control the processing flow of a node, such as initialization, simulation, mutation, and completion.
#[derive(Debug, Serialize, Deserialize, PartialEq, Clone, Copy)]
pub enum GeneticState {
    /// The node and its data have not finished initializing.
    Initialize,
    /// The node is currently simulating a round against target data to determine the fitness of the population.
    Simulate,
    /// The node is currently mutating members of its population and breeding new members.
    Mutate,
    /// The node has finished processing for a given number of iterations.
    Finish,
}

/// Context information passed to [`GeneticNode`] trait methods. Most of the data originates from the [`GeneticNodeWrapper`].
///
/// Contains the current generation, node ID, and simulation context.
#[derive(Clone, Debug)]
pub struct GeneticNodeContext<S> {
    /// The current generation number. Generations start at 0 and increment after each mutation phase.
    pub generation: u64,

    /// The unique identifier for the node.
    pub id: Uuid,

    /// The user defined context specific to the simulation, such as parameters or environment data.
    pub gemla_context: S,
}

/// Trait for managing the state and behavior of nodes within a [`Gemla`] simulation
///
/// Implementors define how nodes are initialized, simulated, mutated, and merged asynchronously.
///
/// # Associated Types
/// - `Context`: The type of context data passed to node methods.
///
/// # Example
/// ```ignore
/// // Example implementation for a custom node type
/// #[async_trait]
/// impl GeneticNode for MyNode {
///     type Context = MyContext;
///     // ...
/// }
/// ```
/// 
/// [`Gemla`]: crate::Gemla
#[async_trait]
pub trait GeneticNode: Send {
    /// Custom type that provides a shared context across different nodes and simulations. Useful if you want to manage
    /// conncurrency or share data between nodes.
    /// 
    /// # Example
    /// ```ignore
    /// pub struct SharedContext {
    ///     pub shared_semaphore: Arc<Semaphore>,
    ///     pub visible_simulations: Arc<Semaphore>,
    /// }
    /// ```
    /// In this example, `SharedContext` could be used to limit the number of concurrent simulations
    /// and control visibility of simulations to users.
    type Context;

    /// Initializes a new instance of a node implementing [`GeneticNode`].
    ///
    /// # Arguments
    /// * `context` - The context for initialization, including generation and node ID.
    ///
    /// # Returns
    /// * A boxed instance of the initialized node.
    ///
    /// # Example
    /// ```rust
    /// # use gemla::core::genetic_node::{GeneticNode, GeneticNodeContext};
    /// # use async_trait::async_trait;
    /// # use serde::{Serialize, Deserialize};
    /// # use rand::prelude::*;
    /// # use uuid::Uuid;
    /// #[derive(Serialize, Deserialize, Debug, Clone)]
    /// struct TestState {
    ///     population: Vec<i64>,
    ///     max_generations: u64,
    /// }
    /// 
    /// #[async_trait]
    /// impl GeneticNode for TestState {
    ///     type Context = ();
    ///     async fn initialize(_context: GeneticNodeContext<Self::Context>) -> Result<Box<Self>, gemla::error::Error> {
    ///         let mut population = vec![];
    ///         for _ in 0..5 {
    ///             population.push(thread_rng().gen_range(0..100));
    ///         }
    ///         Ok(Box::new(TestState { population, max_generations: 10 }))
    ///     }
    ///     // ...
    /// }
    /// ```
    async fn initialize(context: GeneticNodeContext<Self::Context>) -> Result<Box<Self>, Error>;

    /// Simulates a round for the node, updating its state and returning whether to continue.
    ///
    /// # Arguments
    /// * `context` - The context for simulation, including generation and node ID.
    ///
    /// # Returns
    /// * `Ok(true)` if the node should continue to the next phase, `Ok(false)` if finished.
    ///
    /// # Example
    /// ```rust
    /// # use gemla::core::genetic_node::{GeneticNode, GeneticNodeContext};
    /// # use async_trait::async_trait;
    /// # use serde::{Serialize, Deserialize};
    /// # use rand::prelude::*;
    /// # use uuid::Uuid;
    /// #[derive(Serialize, Deserialize, Debug, Clone)]
    /// struct TestState {
    ///     population: Vec<i64>,
    ///     max_generations: u64,
    /// }
    /// 
    /// #[async_trait]
    /// impl GeneticNode for TestState {
    ///     type Context = ();
    ///     async fn simulate(&mut self, context: GeneticNodeContext<Self::Context>) -> Result<bool, gemla::error::Error> {
    ///         let mut rng = thread_rng();
    ///         self.population = self.population.iter().map(|p| p.saturating_add(rng.gen_range(-1..2))).collect();
    ///         if context.generation >= self.max_generations {
    ///             Ok(false)
    ///         } else {
    ///             Ok(true)
    ///         }
    ///     }
    ///     // ...
    /// }
    /// ```
    async fn simulate(&mut self, context: GeneticNodeContext<Self::Context>) -> Result<bool, Error>;

    /// Mutates members in a population and/or crossbreeds them to produce new offspring.
    ///
    /// # Arguments
    /// * `context` - The context for mutation, including generation and node ID.
    ///
    /// # Returns
    /// * `Ok(())` if mutation was successful.
    /// * `Err(Error)` if an error occurred during mutation.
    ///
    /// # Example
    /// ```rust
    /// # use gemla::core::genetic_node::{GeneticNode, GeneticNodeContext};
    /// # use async_trait::async_trait;
    /// # use serde::{Serialize, Deserialize};
    /// # use rand::prelude::*;
    /// # use uuid::Uuid;
    /// #[derive(Serialize, Deserialize, Debug, Clone)]
    /// struct TestState {
    ///     population: Vec<i64>,
    ///     max_generations: u64,
    /// }
    /// 
    /// #[async_trait]
    /// impl GeneticNode for TestState {
    ///     type Context = ();
    ///     async fn mutate(&mut self, _context: GeneticNodeContext<Self::Context>) -> Result<(), gemla::error::Error> {
    ///         let mut rng = thread_rng();
    ///         let mut v = self.population.clone();
    ///         v.sort_unstable();
    ///         v.reverse();
    ///         self.population = v[0..3].to_vec();
    ///         while self.population.len() < 5 {
    ///             let i = rng.gen_range(0..self.population.len());
    ///             let j = loop {
    ///                 let idx = rng.gen_range(0..self.population.len());
    ///                 if idx != i { break idx; }
    ///             };
    ///             let mut new_ind = self.population[i];
    ///             let cross = self.population[j];
    ///             new_ind = (new_ind.saturating_add(cross) / 2).saturating_add(rng.gen_range(-1..2));
    ///             self.population.push(new_ind);
    ///         }
    ///         Ok(())
    ///     }
    ///     // ...
    /// }
    /// ```
    async fn mutate(&mut self, context: GeneticNodeContext<Self::Context>) -> Result<(), Error>;

    /// Merges two nodes into a new node, using the provided context and ID. This occurs after
    /// two nodes have finished simulating and the populations need to be combined.
    ///
    /// # Arguments
    /// * `left` - The left node to merge.
    /// * `right` - The right node to merge.
    /// * `id` - The ID for the new merged node.
    /// * `context` - The context for merging.
    ///
    /// # Example
    /// ```rust
    /// # use gemla::core::genetic_node::{GeneticNode, GeneticNodeContext};
    /// # use async_trait::async_trait;
    /// # use serde::{Serialize, Deserialize};
    /// # use uuid::Uuid;
    /// #[derive(Serialize, Deserialize, Debug, Clone)]
    /// struct TestState {
    ///     population: Vec<i64>,
    ///     max_generations: u64,
    /// }
    /// 
    /// #[async_trait]
    /// impl GeneticNode for TestState {
    ///     type Context = ();
    ///     async fn merge(left: &TestState, right: &TestState, id: &Uuid, gemla_context: Self::Context) -> Result<Box<TestState>, gemla::error::Error> {
    ///         let mut v = left.population.clone();
    ///         v.append(&mut right.population.clone());
    ///         v.sort_by(|a, b| a.partial_cmp(b).unwrap());
    ///         v.reverse();
    ///         v = v[..3].to_vec();
    ///         let mut result = TestState { population: v, max_generations: 10 };
    ///         result.mutate(GeneticNodeContext { id: *id, generation: 0, gemla_context }).await?;
    ///         Ok(Box::new(result))
    ///     }
    ///     // ...
    /// }
    /// ```
    async fn merge(
        left: &Self,
        right: &Self,
        id: &Uuid,
        context: Self::Context,
    ) -> Result<Box<Self>, Error>;
}

/// Wrapper for a node implementing [`GeneticNode`], managing state transitions and lifecycle.
///
/// Used externally to process state transitions and signal recovery. State transitions are managed using [`GeneticState`].
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
pub struct GeneticNodeWrapper<T>
where
    T: Clone,
{
    /// The wrapped node instance, if initialized.
    node: Option<T>,

    /// The current state of the node in the simulation lifecycle.
    state: GeneticState,

    /// The current generation number. Generations start at 0 and increment after each mutation phase.
    generation: u64,

    /// The unique identifier for the node.
    id: Uuid,
}

impl<T> Default for GeneticNodeWrapper<T>
where
    T: Clone,
{
    fn default() -> Self {
        GeneticNodeWrapper {
            node: None,
            state: GeneticState::Initialize,
            generation: 1,
            id: Uuid::new_v4(),
        }
    }
}

impl<T> GeneticNodeWrapper<T>
where
    T: GeneticNode + Debug + Send + Clone,
    T::Context: Send + Sync + Clone + Debug + Serialize + DeserializeOwned + 'static + Default,
{
    /// Creates a new instance of [`GeneticNodeWrapper`]. Calls [`Default::default()`] of the inner type.
    pub fn new() -> Self {
        GeneticNodeWrapper::<T> {
            ..Default::default()
        }
    }

    /// Creates a new instance from the given node data and ID.
    pub fn from(data: T, id: Uuid) -> Self {
        GeneticNodeWrapper {
            node: Some(data),
            state: GeneticState::Simulate,
            generation: 1,
            id,
        }
    }

    /// Returns a reference to the wrapped node, if available.
    pub fn as_ref(&self) -> Option<&T> {
        self.node.as_ref()
    }

    /// Takes the wrapped node, consuming the wrapper.
    pub fn take(&mut self) -> Option<T> {
        self.node.take()
    }

    /// Returns the ID of the node.
    pub fn id(&self) -> Uuid {
        self.id
    }

    /// Returns the current generation number.
    pub fn generation(&self) -> u64 {
        self.generation
    }

    /// Returns the current state of the node.
    pub fn state(&self) -> GeneticState {
        self.state
    }

    /// Processes the node for the current generation, updating its state and transitioning to the next phase as needed.
    ///
    /// # Arguments
    /// * `gemla_context` - The user-defined context for the simulation, passed to node methods.
    ///
    /// # Returns
    /// * The updated state of the node after processing.
    pub async fn process_node(&mut self, gemla_context: T::Context) -> Result<GeneticState, Error> {
        let context = GeneticNodeContext {
            generation: self.generation,
            id: self.id,
            gemla_context,
        };

        match (self.state, &mut self.node) {
            (GeneticState::Initialize, _) => {
                self.node = Some(*T::initialize(context.clone()).await?);
                self.state = GeneticState::Simulate;
            }
            (GeneticState::Simulate, Some(n)) => {
                let next_generation = n
                    .simulate(context.clone())
                    .await
                    .with_context(|| format!("Error simulating node: {:?}", self))?;

                info!("Simulation complete and continuing: {:?}", next_generation);

                self.state = if next_generation {
                    GeneticState::Mutate
                } else {
                    GeneticState::Finish
                };
            }
            (GeneticState::Mutate, Some(n)) => {
                n.mutate(context.clone())
                    .await
                    .with_context(|| format!("Error mutating node: {:?}", self))?;

                self.generation += 1;
                self.state = GeneticState::Simulate;
            }
            (GeneticState::Finish, Some(_)) => (),
            _ => panic!("Error processing node {:?}", self.node),
        }

        Ok(self.state)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::error::Error;
    use anyhow::anyhow;
    use async_trait::async_trait;

    #[derive(Deserialize, Serialize, Clone, Debug, PartialEq)]
    struct TestState {
        pub score: f64,
        pub max_generations: u64,
    }

    #[async_trait]
    impl GeneticNode for TestState {
        type Context = ();

        async fn simulate(
            &mut self,
            context: GeneticNodeContext<Self::Context>,
        ) -> Result<bool, Error> {
            self.score += 1.0;
            if context.generation >= self.max_generations {
                Ok(false)
            } else {
                Ok(true)
            }
        }

        async fn mutate(
            &mut self,
            _context: GeneticNodeContext<Self::Context>,
        ) -> Result<(), Error> {
            Ok(())
        }

        async fn initialize(
            _context: GeneticNodeContext<Self::Context>,
        ) -> Result<Box<TestState>, Error> {
            Ok(Box::new(TestState {
                score: 0.0,
                max_generations: 2,
            }))
        }

        async fn merge(
            _l: &TestState,
            _r: &TestState,
            _id: &Uuid,
            _: Self::Context,
        ) -> Result<Box<TestState>, Error> {
            Err(Error::Other(anyhow!("Unable to merge")))
        }
    }

    #[test]
    fn test_new() -> Result<(), Error> {
        let genetic_node = GeneticNodeWrapper::<TestState>::new();

        let other_genetic_node = GeneticNodeWrapper::<TestState> {
            node: None,
            state: GeneticState::Initialize,
            generation: 1,
            id: genetic_node.id(),
        };

        assert_eq!(genetic_node, other_genetic_node);

        Ok(())
    }

    #[test]
    fn test_from() -> Result<(), Error> {
        let val = TestState {
            score: 0.0,
            max_generations: 10,
        };
        let uuid = Uuid::new_v4();
        let genetic_node = GeneticNodeWrapper::from(val.clone(), uuid);

        let other_genetic_node = GeneticNodeWrapper::<TestState> {
            node: Some(val),
            state: GeneticState::Simulate,
            generation: 1,
            id: genetic_node.id(),
        };

        assert_eq!(genetic_node, other_genetic_node);

        Ok(())
    }

    #[test]
    fn test_as_ref() -> Result<(), Error> {
        let val = TestState {
            score: 3.0,
            max_generations: 10,
        };
        let uuid = Uuid::new_v4();
        let genetic_node = GeneticNodeWrapper::from(val.clone(), uuid);

        let ref_value = genetic_node.as_ref().unwrap();

        assert_eq!(*ref_value, val);

        Ok(())
    }

    #[test]
    fn test_id() -> Result<(), Error> {
        let val = TestState {
            score: 3.0,
            max_generations: 10,
        };
        let uuid = Uuid::new_v4();
        let genetic_node = GeneticNodeWrapper::from(val.clone(), uuid);

        let id_value = genetic_node.id();

        assert_eq!(id_value, uuid);

        Ok(())
    }

    #[test]
    fn test_state() -> Result<(), Error> {
        let val = TestState {
            score: 3.0,
            max_generations: 10,
        };
        let uuid = Uuid::new_v4();
        let genetic_node = GeneticNodeWrapper::from(val.clone(), uuid);

        let state = genetic_node.state();

        assert_eq!(state, GeneticState::Simulate);

        Ok(())
    }

    #[tokio::test]
    async fn test_process_node() -> Result<(), Error> {
        let mut genetic_node = GeneticNodeWrapper::<TestState>::new();

        assert_eq!(genetic_node.state(), GeneticState::Initialize);
        assert_eq!(genetic_node.process_node(()).await?, GeneticState::Simulate);
        assert_eq!(genetic_node.process_node(()).await?, GeneticState::Mutate);
        assert_eq!(genetic_node.process_node(()).await?, GeneticState::Simulate);
        assert_eq!(genetic_node.process_node(()).await?, GeneticState::Finish);
        assert_eq!(genetic_node.process_node(()).await?, GeneticState::Finish);

        Ok(())
    }
}