//! Distributed simulation builder using nested ProtocolMachines.

use std::collections::{BTreeMap, BTreeSet};

use telltale_machine::coroutine::Value;
use telltale_machine::model::effects::{EffectFailure, EffectHandler, EffectResult};
use telltale_machine::runtime::loader::CodeImage;
use telltale_machine::runtime::runner::NestedProtocolMachineHandler;
use telltale_machine::{ProtocolMachine, ProtocolMachineConfig, ProtocolMachineError};

/// Builder for distributed simulations with nested inner ProtocolMachines.
pub struct DistributedSimBuilder {
    sites: BTreeMap<String, Vec<CodeImage>>,
    inter_site: Option<CodeImage>,
    outer_concurrency: usize,
    inner_rounds_per_step: usize,
}

impl DistributedSimBuilder {
    /// Create an empty builder.
    #[must_use]
    pub fn new() -> Self {
        Self {
            sites: BTreeMap::new(),
            inter_site: None,
            outer_concurrency: 1,
            inner_rounds_per_step: 1,
        }
    }

    /// Add a site with its local protocol images.
    #[must_use]
    pub fn add_site(mut self, name: impl Into<String>, protocols: Vec<CodeImage>) -> Self {
        self.sites.insert(name.into(), protocols);
        self
    }

    /// Set the inter-site routing protocol (outer ProtocolMachine).
    #[must_use]
    pub fn inter_site(mut self, protocol: CodeImage) -> Self {
        self.inter_site = Some(protocol);
        self
    }

    /// Set outer ProtocolMachine scheduler concurrency.
    #[must_use]
    pub fn outer_concurrency(mut self, concurrency: usize) -> Self {
        self.outer_concurrency = concurrency.max(1);
        self
    }

    /// Set inner ProtocolMachine rounds attempted per outer handler callback.
    #[must_use]
    pub fn inner_rounds_per_step(mut self, rounds: usize) -> Self {
        self.inner_rounds_per_step = rounds.max(1);
        self
    }

    /// Build with a default no-op inner handler.
    ///
    /// # Errors
    ///
    /// Returns an error if the inter-site protocol is missing or loading fails.
    pub fn build(self, config: &ProtocolMachineConfig) -> Result<DistributedSimulation, String> {
        self.build_with(config, |_| Box::new(NoOpHandler))
    }

    /// Build with a per-site handler factory.
    ///
    /// # Errors
    ///
    /// Returns an error if the inter-site protocol is missing or loading fails.
    pub fn build_with<F>(
        self,
        config: &ProtocolMachineConfig,
        mut handler_factory: F,
    ) -> Result<DistributedSimulation, String>
    where
        F: FnMut(&str) -> Box<dyn EffectHandler>,
    {
        let inter_site = self
            .inter_site
            .ok_or_else(|| "missing inter-site protocol".to_string())?;

        let site_names: BTreeSet<String> = self.sites.keys().cloned().collect();
        let outer_roles: BTreeSet<String> = inter_site.roles().into_iter().collect();
        if site_names != outer_roles {
            return Err(format!(
                "outer protocol roles {outer_roles:?} do not match sites {site_names:?}"
            ));
        }

        let mut outer_vm = ProtocolMachine::new(config.clone());
        outer_vm
            .load_choreography(&inter_site)
            .map_err(|e| format!("outer load error: {e}"))?;

        let mut nested =
            NestedProtocolMachineHandler::new().with_rounds_per_step(self.inner_rounds_per_step);

        for (site, protocols) in self.sites {
            let mut inner_vm = ProtocolMachine::new(config.clone());
            for image in protocols {
                inner_vm
                    .load_choreography(&image)
                    .map_err(|e| format!("inner load error for {site}: {e}"))?;
            }
            let handler = handler_factory(&site);
            nested.add_site(site, inner_vm, handler);
        }

        Ok(DistributedSimulation {
            outer_vm,
            handler: nested,
            outer_concurrency: self.outer_concurrency,
        })
    }
}

impl Default for DistributedSimBuilder {
    fn default() -> Self {
        Self::new()
    }
}

/// A distributed simulation with an outer ProtocolMachine and nested handler.
pub struct DistributedSimulation {
    outer_vm: ProtocolMachine,
    handler: NestedProtocolMachineHandler,
    outer_concurrency: usize,
}

impl DistributedSimulation {
    /// Run the outer ProtocolMachine for a fixed number of rounds.
    ///
    /// # Errors
    ///
    /// Returns a `ProtocolMachineError` if the outer ProtocolMachine faults.
    pub fn run(&mut self, max_rounds: usize) -> Result<(), ProtocolMachineError> {
        self.outer_vm
            .run_concurrent(&self.handler, max_rounds, self.outer_concurrency)
            .map(|_| ())
    }

    /// Access the outer ProtocolMachine.
    #[must_use]
    pub fn outer(&self) -> &ProtocolMachine {
        &self.outer_vm
    }

    /// Access the nested handler.
    #[must_use]
    pub fn handler(&self) -> &NestedProtocolMachineHandler {
        &self.handler
    }

    /// Configured outer ProtocolMachine scheduler concurrency.
    #[must_use]
    pub fn outer_concurrency(&self) -> usize {
        self.outer_concurrency
    }
}

struct NoOpHandler;

impl EffectHandler for NoOpHandler {
    fn handle_send(
        &self,
        _role: &str,
        _partner: &str,
        _label: &str,
        _state: &[Value],
    ) -> EffectResult<Value> {
        EffectResult::success(Value::Unit)
    }

    fn handle_recv(
        &self,
        _role: &str,
        _partner: &str,
        _label: &str,
        _state: &mut Vec<Value>,
        _payload: &Value,
    ) -> EffectResult<()> {
        EffectResult::success(())
    }

    fn handle_choose(
        &self,
        _role: &str,
        _partner: &str,
        labels: &[String],
        _state: &[Value],
    ) -> EffectResult<String> {
        match labels.first().cloned() {
            Some(label) => EffectResult::success(label),
            None => EffectResult::failure(EffectFailure::invalid_input("no labels available")),
        }
    }

    fn step(&self, _role: &str, _state: &mut Vec<Value>) -> EffectResult<()> {
        EffectResult::success(())
    }
}
