//! GraphSnapshotBuilder Python wrapper

use pyo3::prelude::*;
use pyo3::types::{PyBool, PyDict};
use rustychickpeas_core::{GraphBuilder, ValueId};
use crate::graph_snapshot::GraphSnapshot;
use crate::rusty_chickpeas::RustyChickpeas;
use crate::utils::py_to_property_value;

/// Python wrapper for GraphBuilder
#[pyclass(name = "GraphSnapshotBuilder")]
pub struct GraphSnapshotBuilder {
    pub(crate) builder: GraphBuilder,
}

#[pymethods]
impl GraphSnapshotBuilder {
    #[new]
    #[pyo3(signature = (version=None, capacity_nodes=None, capacity_rels=None))]
    fn new(version: Option<String>, capacity_nodes: Option<usize>, capacity_rels: Option<usize>) -> Self {
        let builder = if let Some(v) = version {
            GraphBuilder::with_version(&v, capacity_nodes, capacity_rels)
        } else {
            GraphBuilder::new(capacity_nodes, capacity_rels)
        };
        Self { builder }
    }

    /// Add a node with labels
    /// 
    /// # Arguments
    /// * `labels` - List of label strings
    /// * `node_id` - Optional node ID. If None, auto-generates the next sequential ID.
    ///               If Some(id), uses that ID (must be u32)
    /// 
    /// # Returns
    /// The node ID that was used (either the provided ID or the auto-generated one)
    #[pyo3(signature = (labels, *, node_id = None))]
    fn add_node(&mut self, labels: Vec<String>, node_id: Option<u32>) -> PyResult<u32> {
        let label_refs: Vec<&str> = labels.iter().map(|s| s.as_str()).collect();
        Ok(self.builder.add_node(node_id, &label_refs))
    }

    /// Add a relationship
    /// 
    /// # Arguments
    /// * `u` - Start node ID (must be u32)
    /// * `v` - End node ID (must be u32)
    /// * `rel_type` - Relationship type string
    fn add_rel(&mut self, u: u32, v: u32, rel_type: String) -> PyResult<()> {
        self.builder.add_rel(u, v, &rel_type);
        Ok(())
    }

    /// Set property with automatic type detection
    /// Automatically calls the correct type-specific method based on the value type
    /// 
    /// # Arguments
    /// * `node_id` - Node ID (must be u32)
    /// * `key` - Property key string
    /// * `value` - Property value (str, int, float, or bool)
    fn set_prop(&mut self, node_id: u32, key: String, value: &PyAny) -> PyResult<()> {
        // Check bool first, as True/False can be extracted as int
        if let Ok(b) = value.extract::<bool>() {
            self.builder.set_prop_bool(node_id, &key, b);
        } else if let Ok(s) = value.extract::<String>() {
            self.builder.set_prop_str(node_id, &key, &s);
        } else if let Ok(i) = value.extract::<i64>() {
            self.builder.set_prop_i64(node_id, &key, i);
        } else if let Ok(f) = value.extract::<f64>() {
            self.builder.set_prop_f64(node_id, &key, f);
        } else {
            return Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(
                "Property value must be str, int, float, or bool"
            ));
        }
        Ok(())
    }

    /// Set multiple properties at once from a dictionary
    /// 
    /// This is more efficient than calling set_prop() multiple times because
    /// it reduces FFI overhead by batching all property updates in a single call.
    /// 
    /// # Arguments
    /// * `node_id` - Node ID (must be u32)
    /// * `properties` - Dictionary of property key-value pairs
    /// 
    /// # Example
    /// ```python
    /// builder.set_props(1, {"name": "Alice", "age": 30, "active": True})
    /// ```
    fn set_props(&mut self, node_id: u32, properties: &PyDict) -> PyResult<()> {
        for (key_obj, value_obj) in properties {
            let key: String = key_obj.extract()?;
            let value: &PyAny = value_obj;

            // Check bool first, as True/False can be extracted as int
            if let Ok(b) = value.extract::<bool>() {
                self.builder.set_prop_bool(node_id, &key, b);
            } else if let Ok(s) = value.extract::<String>() {
                self.builder.set_prop_str(node_id, &key, &s);
            } else if let Ok(i) = value.extract::<i64>() {
                self.builder.set_prop_i64(node_id, &key, i);
            } else if let Ok(f) = value.extract::<f64>() {
                self.builder.set_prop_f64(node_id, &key, f);
            } else {
                return Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(
                    format!("Property value for key '{}' must be str, int, float, or bool", key)
                ));
            }
        }
        Ok(())
    }

    /// Set string property
    fn set_prop_str(&mut self, node_id: u32, key: String, value: String) -> PyResult<()> {
        self.builder.set_prop_str(node_id, &key, &value);
        Ok(())
    }

    /// Set i64 property
    fn set_prop_i64(&mut self, node_id: u32, key: String, value: i64) -> PyResult<()> {
        self.builder.set_prop_i64(node_id, &key, value);
        Ok(())
    }

    /// Set f64 property
    fn set_prop_f64(&mut self, node_id: u32, key: String, value: f64) -> PyResult<()> {
        self.builder.set_prop_f64(node_id, &key, value);
        Ok(())
    }

    /// Set boolean property
    fn set_prop_bool(&mut self, node_id: u32, key: String, value: bool) -> PyResult<()> {
        self.builder.set_prop_bool(node_id, &key, value);
        Ok(())
    }

    /// Set string property on a relationship
    /// Finds the relationship by (u, v, rel_type) and sets the property
    fn set_rel_prop_str(&mut self, u: u32, v: u32, rel_type: String, key: String, value: String) -> PyResult<()> {
        self.builder.set_rel_prop_str(u, v, &rel_type, &key, &value);
        Ok(())
    }

    /// Set i64 property on a relationship
    fn set_rel_prop_i64(&mut self, u: u32, v: u32, rel_type: String, key: String, value: i64) -> PyResult<()> {
        self.builder.set_rel_prop_i64(u, v, &rel_type, &key, value);
        Ok(())
    }

    /// Set f64 property on a relationship
    fn set_rel_prop_f64(&mut self, u: u32, v: u32, rel_type: String, key: String, value: f64) -> PyResult<()> {
        self.builder.set_rel_prop_f64(u, v, &rel_type, &key, value);
        Ok(())
    }

    /// Set boolean property on a relationship
    fn set_rel_prop_bool(&mut self, u: u32, v: u32, rel_type: String, key: String, value: bool) -> PyResult<()> {
        self.builder.set_rel_prop_bool(u, v, &rel_type, &key, value);
        Ok(())
    }

    /// Set property on a relationship with automatic type detection
    fn set_rel_prop(&mut self, u: u32, v: u32, rel_type: String, key: String, value: &PyAny) -> PyResult<()> {
        // Check bool first, as True/False can be extracted as int
        if let Ok(b) = value.extract::<bool>() {
            self.builder.set_rel_prop_bool(u, v, &rel_type, &key, b);
        } else if let Ok(s) = value.extract::<String>() {
            self.builder.set_rel_prop_str(u, v, &rel_type, &key, &s);
        } else if let Ok(i) = value.extract::<i64>() {
            self.builder.set_rel_prop_i64(u, v, &rel_type, &key, i);
        } else if let Ok(f) = value.extract::<f64>() {
            self.builder.set_rel_prop_f64(u, v, &rel_type, &key, f);
        } else {
            return Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(
                "Property value must be str, int, float, or bool"
            ));
        }
        Ok(())
    }

    /// Load nodes from a Parquet file into the builder
    fn load_nodes_from_parquet(
        &mut self,
        path: String,
        node_id_column: Option<String>,
        label_columns: Option<Vec<String>>,
        property_columns: Option<Vec<String>>,
        unique_properties: Option<Vec<String>>,
    ) -> PyResult<Vec<u32>> {
        let label_cols = label_columns.as_ref().map(|cols| cols.iter().map(|s| s.as_str()).collect());
        let prop_cols = property_columns.as_ref().map(|cols| cols.iter().map(|s| s.as_str()).collect());
        let unique_props = unique_properties.as_ref().map(|cols| cols.iter().map(|s| s.as_str()).collect());
        
        self.builder
            .load_nodes_from_parquet(
                &path,
                node_id_column.as_deref(),
                label_cols,
                prop_cols,
                unique_props,
            )
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
    }

    /// Load relationships from a Parquet file into the builder
    fn load_relationships_from_parquet(
        &mut self,
        path: String,
        start_node_column: String,
        end_node_column: String,
        rel_type_column: Option<String>,
        property_columns: Option<Vec<String>>,
        fixed_rel_type: Option<String>,
        deduplication: Option<String>,
    ) -> PyResult<Vec<(u32, u32)>> {
        let prop_cols = property_columns.as_ref().map(|cols| cols.iter().map(|s| s.as_str()).collect());
        let dedup = deduplication.as_ref().and_then(|s| match s.as_str() {
            "CreateAll" => Some(rustychickpeas_core::types::RelationshipDeduplication::CreateAll),
            "CreateUniqueByRelType" => Some(rustychickpeas_core::types::RelationshipDeduplication::CreateUniqueByRelType),
            "CreateUniqueByRelTypeAndKeyProperties" => Some(rustychickpeas_core::types::RelationshipDeduplication::CreateUniqueByRelTypeAndKeyProperties),
            _ => None,
        });
        
        self.builder
            .load_relationships_from_parquet(
                &path,
                &start_node_column,
                &end_node_column,
                rel_type_column.as_deref(),
                prop_cols,
                fixed_rel_type.as_deref(),
                dedup,
            )
            .map_err(|e| PyErr::new::<pyo3::exceptions::PyRuntimeError, _>(e.to_string()))
    }

    /// Get property value for a node (before finalization)
    /// Note: Uses get_property to avoid conflict with Python's property builtin
    fn get_property(&self, node_id: u32, key: String) -> PyResult<Option<PyObject>> {
        let value_id = self.builder.prop(node_id, &key);
        
        Python::with_gil(|py| {
            if let Some(vid) = value_id {
                match vid {
                    ValueId::Str(sid) => {
                        // Resolve string ID from builder's interner
                        let s = self.builder.resolve_string(sid);
                        Ok(Some(s.to_object(py)))
                    }
                    ValueId::I64(i) => Ok(Some(i.to_object(py))),
                    ValueId::F64(bits) => Ok(Some(f64::from_bits(bits).to_object(py))),
                    ValueId::Bool(b) => Ok(Some(PyBool::new(py, b).into_py(py))),
                }
            } else {
                Ok(None)
            }
        })
    }

    /// Update property with automatic type detection
    /// Automatically calls the correct type-specific method based on the value type
    fn update_prop(&mut self, node_id: u32, key: String, value: &PyAny) -> PyResult<()> {
        // Check if property exists first
        if self.builder.prop(node_id, &key).is_none() {
            return Err(PyErr::new::<pyo3::exceptions::PyValueError, _>(
                format!("Property key '{}' not found", key)
            ));
        }
        
        // Check bool first, as True/False can be extracted as int
        if let Ok(b) = value.extract::<bool>() {
            self.builder.update_prop_bool(node_id, &key, b);
        } else if let Ok(s) = value.extract::<String>() {
            self.builder.update_prop_str(node_id, &key, &s);
        } else if let Ok(i) = value.extract::<i64>() {
            self.builder.update_prop_i64(node_id, &key, i);
        } else if let Ok(f) = value.extract::<f64>() {
            self.builder.update_prop_f64(node_id, &key, f);
        } else {
            return Err(PyErr::new::<pyo3::exceptions::PyTypeError, _>(
                "Property value must be str, int, float, or bool"
            ));
        }
        Ok(())
    }

    /// Update string property (removes old, sets new)
    fn update_prop_str(&mut self, node_id: u32, key: String, value: String) -> PyResult<()> {
        self.builder.update_prop_str(node_id, &key, &value);
        Ok(())
    }

    /// Update i64 property
    fn update_prop_i64(&mut self, node_id: u32, key: String, value: i64) -> PyResult<()> {
        self.builder.update_prop_i64(node_id, &key, value);
        Ok(())
    }

    /// Update f64 property
    fn update_prop_f64(&mut self, node_id: u32, key: String, value: f64) -> PyResult<()> {
        self.builder.update_prop_f64(node_id, &key, value);
        Ok(())
    }

    /// Update boolean property
    fn update_prop_bool(&mut self, node_id: u32, key: String, value: bool) -> PyResult<()> {
        self.builder.update_prop_bool(node_id, &key, value);
        Ok(())
    }

    /// Get nodes with a specific property value, scoped by label (before finalization)
    /// 
    /// # Arguments
    /// * `label` - The label to scope the query to
    /// * `key` - The property key
    /// * `value` - The property value to search for
    #[pyo3(signature = (label, key, value))]
    fn nodes_with_property(&self, label: String, key: String, value: &PyAny) -> PyResult<Vec<u32>> {
        let prop_value = py_to_property_value(value)?;
        let value_id = match prop_value {
            rustychickpeas_core::PropertyValue::String(_s) => {
                // Need to intern the string to get ID
                // For now, we can't easily do this without exposing interner
                // TODO: Add helper to convert PropertyValue to ValueId
                return Err(PyErr::new::<pyo3::exceptions::PyNotImplementedError, _>(
                    "String property queries not yet supported in GraphBuilder"
                ));
            }
            rustychickpeas_core::PropertyValue::Integer(i) => ValueId::I64(i),
            rustychickpeas_core::PropertyValue::Float(f) => ValueId::from_f64(f),
            rustychickpeas_core::PropertyValue::Boolean(b) => ValueId::Bool(b),
            rustychickpeas_core::PropertyValue::InternedString(_) => {
                return Err(PyErr::new::<pyo3::exceptions::PyNotImplementedError, _>(
                    "InternedString not supported in GraphBuilder queries"
                ));
            }
        };
        
        let node_ids = self.builder.nodes_with_property(&label, &key, value_id);
        Ok(node_ids)
    }

    /// Get node labels (before finalization)
    fn node_labels(&self, node_id: u32) -> PyResult<Vec<String>> {
        Ok(self.builder.node_labels(node_id))
    }

    /// Get neighbors of a node (before finalization)
    /// Returns (outgoing, incoming) as tuple of lists of node IDs
    fn neighbor_ids(&self, node_id: u32) -> PyResult<(Vec<u32>, Vec<u32>)> {
        let (out, inc) = self.builder.neighbor_ids(node_id);
        Ok((out, inc))
    }

    /// Set the version for this snapshot
    /// This version will be stored with the snapshot when finalized
    fn set_version(&mut self, version: String) -> PyResult<()> {
        self.builder.set_version(&version);
        Ok(())
    }

    /// Finalize the builder into a GraphSnapshot
    /// 
    /// # Arguments
    /// * `index_properties` - Optional list of property key names to index during finalization.
    ///   If provided, these properties will be indexed upfront (faster queries, more memory).
    ///   If None, all properties will be indexed lazily on first access (saves memory).
    #[pyo3(signature = (index_properties=None))]
    fn finalize(&mut self, index_properties: Option<Vec<String>>) -> PyResult<GraphSnapshot> {
        // We need to take ownership to finalize, so we'll use replace
        let builder = std::mem::replace(&mut self.builder, GraphBuilder::new(None, None));
        
        // Convert Vec<String> to Vec<&str> for the Rust API
        let keys_to_index: Option<Vec<&str>> = index_properties.as_ref().map(|names| {
            names.iter().map(|s| s.as_str()).collect()
        });
        
        let snapshot = builder.finalize(keys_to_index.as_deref());
        Ok(GraphSnapshot::new(snapshot))
    }

    /// Finalize the builder into a GraphSnapshot and add it to the manager
    /// 
    /// This is a convenience method that finalizes the builder and automatically
    /// adds the snapshot to the manager. Equivalent to:
    /// ```python
    /// snapshot = builder.finalize()
    /// manager.add_snapshot(snapshot)
    /// ```
    /// 
    /// # Arguments
    /// * `index_properties` - Optional list of property key names to index during finalization.
    ///   If provided, these properties will be indexed upfront (faster queries, more memory).
    ///   If None, all properties will be indexed lazily on first access (saves memory).
    #[pyo3(signature = (manager, index_properties=None))]
    fn finalize_into(&mut self, manager: &RustyChickpeas, index_properties: Option<Vec<String>>) -> PyResult<()> {
        let builder = std::mem::replace(&mut self.builder, GraphBuilder::new(None, None));
        
        // Convert Vec<String> to Vec<&str> for the Rust API
        let keys_to_index: Option<Vec<&str>> = index_properties.as_ref().map(|names| {
            names.iter().map(|s| s.as_str()).collect()
        });
        
        let snapshot = builder.finalize(keys_to_index.as_deref());
        manager.manager.add_snapshot(snapshot);
        Ok(())
    }
}

