//! Python bindings for MessageGraphBuilder

use pyo3::prelude::*;

use flowgentra_ai::core::llm::Message;
use flowgentra_ai::core::state::PlainState;
use flowgentra_ai::core::state_graph::message_graph::MessageGraphBuilder;

use crate::error::to_py_err_generic;
use crate::get_runtime;
use crate::llm::PyMessage;

// ─── PyMessageGraph (compiled) ──────────────────────────────────────────────

/// A compiled message-centric graph.
#[pyclass(name = "MessageGraph")]
pub struct PyMessageGraph {
    inner: flowgentra_ai::core::state_graph::StateGraph<PlainState>,
}

#[pymethods]
impl PyMessageGraph {
    /// Invoke the graph with a list of initial messages.
    fn invoke(&self, messages: Vec<PyMessage>) -> PyResult<Vec<PyMessage>> {
        let msgs: Vec<Message> = messages.into_iter().map(|m| m.inner).collect();
        let state = MessageGraphBuilder::initial_state(msgs);
        let rt = get_runtime();
        let result = rt.block_on(self.inner.invoke(state)).map_err(to_py_err_generic)?;
        let out_msgs = MessageGraphBuilder::get_messages(&result);
        Ok(out_msgs.into_iter().map(|m| PyMessage { inner: m }).collect())
    }

    fn __repr__(&self) -> String {
        "MessageGraph(...)".to_string()
    }
}

// ─── PyMessageGraphBuilder ──────────────────────────────────────────────────

/// Convenience builder for chat-focused workflows.
///
/// Pre-configures a graph with message accumulation.
/// Nodes receive PlainState with a "messages" array.
///
/// Example:
///     builder = MessageGraphBuilder()
///     builder.add_node("echo", echo_fn)
///     builder.set_entry_point("echo")
///     builder.add_edge("echo", "__end__")
///     graph = builder.compile()
///     result = graph.invoke([Message.user("Hello")])
#[pyclass(name = "MessageGraphBuilder")]
pub struct PyMessageGraphBuilder {
    inner: Option<MessageGraphBuilder>,
}

#[pymethods]
impl PyMessageGraphBuilder {
    #[new]
    fn new() -> Self {
        PyMessageGraphBuilder {
            inner: Some(MessageGraphBuilder::new()),
        }
    }

    /// Add a node with a Python callable.
    ///
    /// The callable receives a list of Messages and must return a list of Messages.
    /// New messages in the returned list are appended to the conversation.
    fn add_node(&mut self, name: &str, func: PyObject) {
        let builder = self.inner.take().unwrap_or_else(MessageGraphBuilder::new);
        let func_clone = Python::with_gil(|py| func.clone_ref(py));
        let node_name = name.to_string();

        self.inner = Some(builder.add_fn(name, move |state: &PlainState| {
            let state_clone = state.clone();
            let func = Python::with_gil(|py| func_clone.clone_ref(py));
            let name = node_name.clone();

            Box::pin(async move {
                Python::with_gil(|py| -> Result<PlainState, flowgentra_ai::core::state_graph::StateGraphError> {
                    // Get messages from state
                    let messages = MessageGraphBuilder::get_messages(&state_clone);
                    let py_messages: Vec<PyMessage> = messages
                        .into_iter()
                        .map(|m| PyMessage { inner: m })
                        .collect();

                    // Call Python function with messages list
                    let result = func.call1(py, (py_messages,)).map_err(|e| {
                        flowgentra_ai::core::state_graph::StateGraphError::ExecutionError {
                            node: name.clone(),
                            reason: format!("{}", e),
                        }
                    })?;

                    // Extract returned messages
                    let returned_messages: Vec<PyMessage> = result.extract(py).map_err(|e| {
                        flowgentra_ai::core::state_graph::StateGraphError::ExecutionError {
                            node: name.clone(),
                            reason: format!("Node must return list of Messages: {}", e),
                        }
                    })?;

                    // Build new state with all returned messages
                    let rust_msgs: Vec<Message> = returned_messages.into_iter().map(|m| m.inner).collect();
                    Ok(MessageGraphBuilder::initial_state(rust_msgs))
                })
            })
        }));
    }

    /// Add a fixed edge.
    fn add_edge(&mut self, from_node: &str, to_node: &str) {
        let builder = self.inner.take().unwrap_or_else(MessageGraphBuilder::new);
        self.inner = Some(builder.add_edge(from_node, to_node));
    }

    /// Set the entry point.
    fn set_entry_point(&mut self, name: &str) {
        let builder = self.inner.take().unwrap_or_else(MessageGraphBuilder::new);
        self.inner = Some(builder.set_entry_point(name));
    }

    /// Compile the graph.
    fn compile(&mut self) -> PyResult<PyMessageGraph> {
        let builder = self.inner.take().unwrap_or_else(MessageGraphBuilder::new);
        let graph = builder.compile().map_err(to_py_err_generic)?;
        Ok(PyMessageGraph { inner: graph })
    }

    fn __repr__(&self) -> String {
        "MessageGraphBuilder(...)".to_string()
    }
}

// ─── Helper functions ───────────────────────────────────────────────────────

/// Create initial state with messages (helper for advanced use).
#[allow(dead_code)]
#[pyfunction]
pub fn py_message_graph_initial_state(messages: Vec<PyMessage>) -> PyResult<crate::state::PyPlainState> {
    let msgs: Vec<Message> = messages.into_iter().map(|m| m.inner).collect();
    let state = MessageGraphBuilder::initial_state(msgs);
    Ok(crate::state::PyPlainState { inner: state })
}
