//! Python bindings for Supervisor (multi-agent orchestration)

use pyo3::prelude::*;
use std::sync::Arc;

use flowgentra_ai::core::state::SharedState;
use flowgentra_ai::core::state_graph::StateGraph;

use crate::error::to_py_err_generic;
use crate::get_runtime;
use crate::graph::PyStateGraph;
use crate::state::PySharedState;

// ─── PySupervisor ───────────────────────────────────────────────────────────

/// Multi-agent supervisor that orchestrates named agent graphs.
///
/// Routes work between agents based on a routing function.
///
/// Example:
///     def router(state):
///         task = state.get_string("task") or ""
///         if "research" in task:
///             return "researcher"
///         return "writer"
///
///     sup = Supervisor(router)
///     sup.add_agent("researcher", research_graph)
///     sup.add_agent("writer", writer_graph)
///     sup.max_rounds(5)
///     result = sup.run(SharedState({"task": "research AI"}))
#[pyclass(name = "Supervisor")]
pub struct PySupervisor {
    router: PyObject,
    agents: Vec<(String, Arc<StateGraph<SharedState>>)>,
    max_rounds_val: usize,
    finish_marker_val: String,
}

#[pymethods]
impl PySupervisor {
    /// Create a new supervisor with a routing function.
    ///
    /// The router receives a SharedState and returns the name of the
    /// agent to dispatch to, or the finish marker string to stop.
    #[new]
    fn new(router: PyObject) -> Self {
        PySupervisor {
            router,
            agents: Vec::new(),
            max_rounds_val: 10,
            finish_marker_val: "FINISH".to_string(),
        }
    }

    /// Add a named agent graph.
    fn add_agent(&mut self, name: &str, graph: &PyStateGraph) {
        self.agents.push((name.to_string(), graph.inner.clone()));
    }

    /// Set the maximum number of routing rounds.
    fn max_rounds(&mut self, rounds: usize) {
        self.max_rounds_val = rounds;
    }

    /// Set the finish marker string (default: "FINISH").
    fn finish_marker(&mut self, marker: &str) {
        self.finish_marker_val = marker.to_string();
    }

    /// Run the supervisor loop.
    fn run(&self, state: &PySharedState) -> PyResult<PySharedState> {
        let router = Python::with_gil(|py| self.router.clone_ref(py));
        let finish_marker = self.finish_marker_val.clone();
        let max_rounds = self.max_rounds_val;
        let mut current_state = state.inner.clone();
        let rt = get_runtime();

        for _round in 0..max_rounds {
            // Call the Python router to get next agent name
            let next_agent = Python::with_gil(|py| -> PyResult<String> {
                let py_state = PySharedState { inner: current_state.clone() };
                let result = router.call1(py, (py_state,))?;
                result.extract::<String>(py)
            })?;

            if next_agent == finish_marker {
                return Ok(PySharedState { inner: current_state });
            }

            // Find the agent graph and invoke it
            let agent_graph = self.agents.iter().find(|(n, _)| n == &next_agent);
            match agent_graph {
                Some((_, graph)) => {
                    current_state = rt
                        .block_on(graph.invoke(current_state))
                        .map_err(to_py_err_generic)?;
                }
                None => {
                    let names: Vec<&str> = self.agents.iter().map(|(n, _)| n.as_str()).collect();
                    return Err(pyo3::exceptions::PyRuntimeError::new_err(format!(
                        "Supervisor: agent '{}' not found. Available: {:?}",
                        next_agent, names
                    )));
                }
            }
        }

        Err(pyo3::exceptions::PyRuntimeError::new_err(format!(
            "Supervisor exceeded max rounds ({})",
            max_rounds
        )))
    }

    /// Get the list of registered agent names.
    fn agent_names(&self) -> Vec<String> {
        self.agents.iter().map(|(n, _)| n.clone()).collect()
    }

    fn __repr__(&self) -> String {
        let names: Vec<&str> = self.agents.iter().map(|(n, _)| n.as_str()).collect();
        format!("Supervisor(agents={:?})", names)
    }
}
