//! Python bindings for StateGraph / StateGraphBuilder
//!
//! Exposes the code-driven graph API to Python, letting users build
//! workflows with Python callables as node functions.

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

use flowgentra_ai::core::state::SharedState;
use flowgentra_ai::core::state_graph::{
    StateGraph, StateGraphBuilder, StateGraphError,
    edge::END,
    node::Node,
    FileCheckpointer,
};

use crate::state::PySharedState;
use crate::get_runtime;

// ─── Subgraph wrapper node ─────────────────────────────────────────────────

/// Wraps an `Arc<StateGraph<SharedState>>` as a `Node<SharedState>`.
///
/// This avoids the `Arc::try_unwrap` ownership problem: Python holds a
/// reference to the compiled subgraph, so we cannot take ownership.
/// Instead we invoke the subgraph through the shared Arc.
struct ArcSubgraphNode {
    name: String,
    subgraph: Arc<StateGraph<SharedState>>,
}

#[async_trait::async_trait]
impl Node<SharedState> for ArcSubgraphNode {
    async fn execute(&self, state: &SharedState) -> Result<SharedState, StateGraphError> {
        self.subgraph.invoke(state.clone()).await
    }

    fn name(&self) -> &str {
        &self.name
    }
}

// ─── Python-callable Node ──────────────────────────────────────────────────

/// A node backed by a Python callable.
///
/// The callable receives a `SharedState` and must return a `SharedState`.
struct PyFunctionNode {
    name: String,
    func: PyObject,
}

#[async_trait::async_trait]
impl Node<SharedState> for PyFunctionNode {
    async fn execute(&self, state: &SharedState) -> Result<SharedState, StateGraphError> {
        let state_clone = state.clone();

        // Acquire the GIL to call into Python
        let result = Python::with_gil(|py| -> PyResult<SharedState> {
            let func = self.func.clone_ref(py);
            let py_state = PySharedState { inner: state_clone };
            let py_result = func.call1(py, (py_state,))?;
            let result_state: PySharedState = py_result.extract(py)?;
            Ok(result_state.inner)
        });

        result.map_err(|e| StateGraphError::ExecutionError {
            node: self.name.clone(),
            reason: format!("Python node error: {}", e),
        })
    }

    fn name(&self) -> &str {
        &self.name
    }
}

// ─── Python-callable Router ────────────────────────────────────────────────

/// Wraps a Python callable as a sync router function.
///
/// The callable receives a `SharedState` and must return a `str` (node name or "__end__").
fn make_router_fn(
    py_func: PyObject,
) -> Box<dyn Fn(&SharedState) -> Result<String, StateGraphError> + Send + Sync> {
    Box::new(move |state: &SharedState| {
        let state_clone = state.clone();

        Python::with_gil(|py| -> Result<String, StateGraphError> {
            let func = py_func.clone_ref(py);
            let py_state = PySharedState { inner: state_clone };
            let py_result = func
                .call1(py, (py_state,))
                .map_err(|e| StateGraphError::ExecutionError {
                    node: "router".into(),
                    reason: format!("Python router error: {}", e),
                })?;
            let next_node: String =
                py_result
                    .extract(py)
                    .map_err(|e| StateGraphError::ExecutionError {
                        node: "router".into(),
                        reason: format!("Router must return str: {}", e),
                    })?;
            Ok(next_node)
        })
    })
}

// ─── PyStateGraphBuilder ───────────────────────────────────────────────────

/// Build a state graph with Python callables as nodes.
///
/// Example:
///     def my_node(state):
///         state["output"] = state["input"].upper()
///         return state
///
///     builder = StateGraphBuilder()
///     builder.add_node("process", my_node)
///     builder.set_entry_point("process")
///     builder.add_edge("process", END)
///     graph = builder.compile()
///
///     state = SharedState({"input": "hello"})
///     result = graph.invoke(state)
///     print(result.to_dict())  # {"input": "hello", "output": "HELLO"}
#[pyclass(name = "StateGraphBuilder")]
pub struct PyStateGraphBuilder {
    nodes: Vec<(String, Arc<dyn Node<SharedState>>)>,
    edges: Vec<(String, String)>,
    conditional_edges: Vec<(String, PyObject)>,
    entry_point: Option<String>,
    max_steps: usize,
    interrupt_before: Vec<String>,
    interrupt_after: Vec<String>,
    subgraphs: Vec<(String, Arc<StateGraph<SharedState>>)>,
    checkpointer_path: Option<String>,
}

#[pymethods]
impl PyStateGraphBuilder {
    /// Create a new empty graph builder.
    #[new]
    fn new() -> Self {
        PyStateGraphBuilder {
            nodes: Vec::new(),
            edges: Vec::new(),
            conditional_edges: Vec::new(),
            entry_point: None,
            max_steps: 1000,
            interrupt_before: Vec::new(),
            interrupt_after: Vec::new(),
            subgraphs: Vec::new(),
            checkpointer_path: None,
        }
    }

    /// Add a node backed by a Python callable.
    ///
    /// The callable signature is: `(state: SharedState) -> SharedState`
    fn add_node(&mut self, name: &str, func: PyObject) {
        let node = Arc::new(PyFunctionNode {
            name: name.to_string(),
            func,
        }) as Arc<dyn Node<SharedState>>;
        self.nodes.push((name.to_string(), node));
    }

    /// Add a fixed edge (from → to).
    ///
    /// Use `END` (or `"__end__"`) as the `to` value to terminate the graph.
    fn add_edge(&mut self, from: &str, to: &str) {
        self.edges.push((from.to_string(), to.to_string()));
    }

    /// Add a conditional edge.
    ///
    /// The router callable receives a `SharedState` and must return
    /// a `str` — either the name of the next node or `"__end__"`.
    fn add_conditional_edge(&mut self, from: &str, router: PyObject) {
        self.conditional_edges.push((from.to_string(), router));
    }

    /// Set the entry point node (first node executed after START).
    fn set_entry_point(&mut self, name: &str) {
        self.entry_point = Some(name.to_string());
    }

    /// Set the maximum number of execution steps (default: 1000).
    fn set_max_steps(&mut self, max_steps: usize) {
        self.max_steps = max_steps;
    }

    /// Mark a node to pause execution BEFORE it runs (for human-in-the-loop).
    fn interrupt_before(&mut self, name: &str) {
        self.interrupt_before.push(name.to_string());
    }

    /// Mark a node to pause execution AFTER it runs.
    fn interrupt_after(&mut self, name: &str) {
        self.interrupt_after.push(name.to_string());
    }

    /// Add a compiled subgraph as a node.
    ///
    /// The subgraph will execute as a single node in this graph.
    fn add_subgraph(&mut self, name: &str, subgraph: &PyStateGraph) {
        self.subgraphs.push((name.to_string(), subgraph.inner.clone()));
    }

    /// Set a file checkpointer for persistent state across invocations.
    ///
    /// Args:
    ///     base_dir: Directory path where checkpoint files will be stored.
    fn set_checkpointer(&mut self, base_dir: &str) -> PyResult<()> {
        self.checkpointer_path = Some(base_dir.to_string());
        Ok(())
    }

    /// Compile the builder into a runnable StateGraph.
    ///
    /// Raises RuntimeError if validation fails (e.g. missing entry point,
    /// edges to non-existent nodes, etc.)
    fn compile(&self, py: Python<'_>) -> PyResult<PyStateGraph> {
        let mut builder = StateGraphBuilder::<SharedState>::new();

        for (name, node) in &self.nodes {
            builder = builder.add_node(name.clone(), node.clone());
        }

        for (from, to) in &self.edges {
            builder = builder.add_edge(from.clone(), to.clone());
        }

        for (from, router) in &self.conditional_edges {
            let router_clone = router.clone_ref(py);
            builder = builder.add_conditional_edge(from.clone(), make_router_fn(router_clone));
        }

        if let Some(ref ep) = self.entry_point {
            builder = builder.set_entry_point(ep.clone());
        }

        builder = builder.set_max_steps(self.max_steps);

        for name in &self.interrupt_before {
            builder = builder.interrupt_before(name.clone());
        }
        for name in &self.interrupt_after {
            builder = builder.interrupt_after(name.clone());
        }

        // Subgraphs are added as regular nodes via ArcSubgraphNode, which
        // holds an Arc<StateGraph> and delegates execute() to invoke().
        // This avoids the Arc::try_unwrap ownership problem.
        for (name, subgraph) in &self.subgraphs {
            let node = Arc::new(ArcSubgraphNode {
                name: name.clone(),
                subgraph: subgraph.clone(),
            }) as Arc<dyn Node<SharedState>>;
            builder = builder.add_node(name.clone(), node);
        }

        if let Some(ref path) = self.checkpointer_path {
            let cp = FileCheckpointer::new(path)
                .map_err(|e| PyRuntimeError::new_err(format!("Failed to create checkpointer: {}", e)))?;
            builder = builder.set_checkpointer(Arc::new(cp));
        }

        let graph = builder
            .compile()
            .map_err(|e| PyRuntimeError::new_err(format!("{}", e)))?;

        Ok(PyStateGraph {
            inner: Arc::new(graph),
        })
    }

    fn __repr__(&self) -> String {
        let node_names: Vec<&str> = self.nodes.iter().map(|(n, _)| n.as_str()).collect();
        format!(
            "StateGraphBuilder(nodes={:?}, entry={:?})",
            node_names, self.entry_point
        )
    }
}

// ─── PyStateGraph ──────────────────────────────────────────────────────────

/// A compiled, runnable state graph.
///
/// Execute it with `invoke()` or `invoke_with_thread()`.
///
/// Example:
///     result = graph.invoke(SharedState({"input": "hello"}))
///     print(result.to_dict())
#[pyclass(name = "StateGraph")]
#[derive(Clone)]
pub struct PyStateGraph {
    pub(crate) inner: Arc<StateGraph<SharedState>>,
}

#[pymethods]
impl PyStateGraph {
    /// Execute the graph with the given initial state.
    ///
    /// Returns the final state after all nodes have executed.
    fn invoke(&self, state: &PySharedState) -> PyResult<PySharedState> {
        let rt = get_runtime();
        let result = rt
            .block_on(self.inner.invoke(state.inner.clone()))
            .map_err(|e| PyRuntimeError::new_err(format!("{}", e)))?;
        Ok(PySharedState { inner: result })
    }

    /// Execute the graph with a thread ID for checkpointing.
    ///
    /// Use this when you need to resume interrupted graphs.
    fn invoke_with_thread(
        &self,
        thread_id: &str,
        state: &PySharedState,
    ) -> PyResult<PySharedState> {
        let rt = get_runtime();
        let result = rt
            .block_on(
                self.inner
                    .invoke_with_id(thread_id.to_string(), state.inner.clone()),
            )
            .map_err(|e| PyRuntimeError::new_err(format!("{}", e)))?;
        Ok(PySharedState { inner: result })
    }

    /// Resume a previously interrupted graph from its checkpoint.
    fn resume(&self, thread_id: &str) -> PyResult<PySharedState> {
        let rt = get_runtime();
        let result = rt
            .block_on(self.inner.resume(thread_id))
            .map_err(|e| PyRuntimeError::new_err(format!("{}", e)))?;
        Ok(PySharedState { inner: result })
    }

    /// Resume with injected state updates (human-in-the-loop).
    fn resume_with_state(
        &self,
        thread_id: &str,
        updates: &PySharedState,
    ) -> PyResult<PySharedState> {
        let rt = get_runtime();
        let result = rt
            .block_on(
                self.inner
                    .resume_with_state(thread_id, updates.inner.clone()),
            )
            .map_err(|e| PyRuntimeError::new_err(format!("{}", e)))?;
        Ok(PySharedState { inner: result })
    }

    /// Get the list of node names in this graph.
    fn node_names(&self) -> Vec<String> {
        self.inner.node_names()
    }

    /// Get the entry point node name.
    fn entry_point(&self) -> String {
        self.inner.entry_point().to_string()
    }

    /// Export the graph as a Graphviz DOT string.
    fn to_dot(&self) -> String {
        self.inner.to_dot()
    }

    /// Export the graph as a Mermaid diagram string.
    fn to_mermaid(&self) -> String {
        self.inner.to_mermaid()
    }

    /// Export the graph structure as JSON.
    fn to_json(&self, py: Python<'_>) -> PyResult<PyObject> {
        let val = self.inner.to_json();
        crate::json_to_py(py, &val)
    }

    fn __repr__(&self) -> String {
        format!(
            "StateGraph(entry='{}', nodes={:?})",
            self.inner.entry_point(),
            self.inner.node_names()
        )
    }
}

// ─── Constants ─────────────────────────────────────────────────────────────

/// The END sentinel — use as edge target to terminate the graph.
pub const PY_END: &str = END;

/// Free function to get the END constant.
#[pyfunction]
pub fn graph_end() -> String {
    END.to_string()
}
