//! SiFu (Signal Fully-connected flowing) Python bindings
//!
//! Phase-51 T51.3: PyO3 bindings for the SiFu signal propagation engine.

use std::sync::Arc;
use std::collections::HashSet;

use pyo3::prelude::*;
use pyo3::types::PyDict;

use aifastdb_core::storage::GraphStore as CoreGraphStore;
use aifastdb_core::query::sifu::{
    SiFuQueryEngine as CoreSiFuEngine,
    SiFuConfig as CoreSiFuConfig,
    ActivationFn as CoreActivationFn,
    EdgeWeightStore as CoreEdgeWeightStore,
    EdgeWeight as CoreEdgeWeight,
    TransformMatrix as CoreTransformMatrix,
};
use aifastdb_core::query::sifu_trainer::{
    HebbianConfig as CoreHebbianConfig,
    ContrastiveConfig as CoreContrastiveConfig,
    HebbianLearner as CoreHebbianLearner,
    ContrastiveFitter as CoreContrastiveFitter,
    QueryLog as CoreQueryLog,
    QueryLogEntry as CoreQueryLogEntry,
    PromotionConfig as CorePromotionConfig,
    run_promotion_scan as core_run_promotion_scan,
};

const DEFAULT_SEMANTIC_DIM: usize = 3;

fn extract_sifu_config(config: Option<&Bound<'_, PyDict>>) -> CoreSiFuConfig {
    let mut c = CoreSiFuConfig::default();
    if let Some(d) = config {
        if let Ok(Some(v)) = d.get_item("max_steps") {
            if let Ok(val) = v.extract::<usize>() { c.max_steps = val; }
        }
        if let Ok(Some(v)) = d.get_item("energy_threshold") {
            if let Ok(val) = v.extract::<f64>() { c.energy_threshold = val; }
        }
        if let Ok(Some(v)) = d.get_item("decay_factor") {
            if let Ok(val) = v.extract::<f64>() { c.decay_factor = val; }
        }
        if let Ok(Some(v)) = d.get_item("activation") {
            if let Ok(val) = v.extract::<String>() {
                c.activation = match val.as_str() {
                    "sigmoid" => CoreActivationFn::Sigmoid,
                    "tanh" => CoreActivationFn::Tanh,
                    "linear" => CoreActivationFn::Linear,
                    _ => CoreActivationFn::ReLU,
                };
            }
        }
        if let Ok(Some(v)) = d.get_item("top_k") {
            if let Ok(val) = v.extract::<usize>() { c.top_k = val; }
        }
    }
    c
}

#[pyclass(name = "SiFuEngine")]
pub struct PySiFuEngine {
    engine: Arc<CoreSiFuEngine>,
    edge_weight_store: Arc<parking_lot::RwLock<CoreEdgeWeightStore>>,
    #[allow(dead_code)]
    graph_store: Arc<CoreGraphStore>,
    query_log: Arc<CoreQueryLog>,
}

#[pymethods]
impl PySiFuEngine {
    #[new]
    #[pyo3(signature = (graph_path, semantic_dim=None))]
    fn new(graph_path: String, semantic_dim: Option<usize>) -> PyResult<Self> {
        let dim = semantic_dim.unwrap_or(DEFAULT_SEMANTIC_DIM);
        let graph_store = Arc::new(
            CoreGraphStore::open_or_create(&graph_path)
                .map_err(|e| pyo3::exceptions::PyIOError::new_err(format!("{}", e)))?
        );
        let edge_weight_store = Arc::new(parking_lot::RwLock::new(CoreEdgeWeightStore::new(dim)));
        let engine = Arc::new(CoreSiFuEngine::new(
            graph_store.clone(),
            edge_weight_store.clone(),
        ));
        let query_log = Arc::new(CoreQueryLog::new(10000));

        Ok(Self { engine, edge_weight_store, graph_store, query_log })
    }

    #[pyo3(signature = (seeds, config=None))]
    fn propagate(
        &self,
        py: Python<'_>,
        seeds: Vec<(String, f64)>,
        config: Option<&Bound<'_, PyDict>>,
    ) -> PyResult<Vec<PyObject>> {
        let core_config = extract_sifu_config(config);
        let results = self.engine.propagate(&seeds, &core_config);
        results
            .into_iter()
            .map(|r| {
                let dict = PyDict::new_bound(py);
                dict.set_item("node_id", r.node_id)?;
                dict.set_item("energy", r.energy)?;
                dict.set_item("steps", r.steps)?;
                Ok(dict.into())
            })
            .collect()
    }

    #[pyo3(signature = (seeds, config=None))]
    fn trace_flow(
        &self,
        py: Python<'_>,
        seeds: Vec<(String, f64)>,
        config: Option<&Bound<'_, PyDict>>,
    ) -> PyResult<PyObject> {
        let core_config = extract_sifu_config(config);
        let (results, trace) = self.engine.propagate_with_trace(&seeds, &core_config, true);

        let results_list: Vec<PyObject> = results
            .into_iter()
            .map(|r| {
                let dict = PyDict::new_bound(py);
                dict.set_item("node_id", r.node_id).unwrap();
                dict.set_item("energy", r.energy).unwrap();
                dict.set_item("steps", r.steps).unwrap();
                dict.into()
            })
            .collect();

        let out = PyDict::new_bound(py);
        out.set_item("results", results_list)?;

        if let Some(t) = trace {
            let trace_dict = PyDict::new_bound(py);
            let hops: Vec<PyObject> = t.hops.iter().map(|h| {
                let d = PyDict::new_bound(py);
                d.set_item("from_id", &h.from_id).unwrap();
                d.set_item("to_id", &h.to_id).unwrap();
                d.set_item("relation_type", &h.relation_type).unwrap();
                d.set_item("energy_in", h.energy_in).unwrap();
                d.set_item("energy_out", h.energy_out).unwrap();
                d.set_item("signal_strength", h.signal_strength).unwrap();
                d.set_item("step", h.step).unwrap();
                d.into()
            }).collect();
            trace_dict.set_item("hops", hops)?;
            trace_dict.set_item("total_steps", t.total_steps)?;
            trace_dict.set_item("nodes_visited", t.final_energies.len())?;
            out.set_item("trace", trace_dict)?;
        }

        Ok(out.into())
    }

    #[pyo3(signature = (node_id, top_k=None))]
    fn recommend_related(
        &self,
        py: Python<'_>,
        node_id: String,
        top_k: Option<usize>,
    ) -> PyResult<Vec<PyObject>> {
        let k = top_k.unwrap_or(10);
        let config = CoreSiFuConfig { top_k: k, ..Default::default() };
        let seeds = vec![(node_id, 1.0)];
        self.engine.propagate(&seeds, &config)
            .into_iter()
            .map(|r| {
                let dict = PyDict::new_bound(py);
                dict.set_item("node_id", r.node_id)?;
                dict.set_item("energy", r.energy)?;
                dict.set_item("steps", r.steps)?;
                Ok(dict.into())
            })
            .collect()
    }

    fn discover_clusters(
        &self,
        py: Python<'_>,
        seed_node_ids: Vec<String>,
        num_clusters: Option<usize>,
    ) -> PyResult<Vec<PyObject>> {
        let k = num_clusters.unwrap_or(5);
        let config = CoreSiFuConfig { top_k: k, ..Default::default() };
        let seed_refs: Vec<&str> = seed_node_ids.iter().map(|s| s.as_str()).collect();
        let results = self.engine.discover_clusters(&seed_refs, &config);
        results
            .into_iter()
            .map(|r| {
                let dict = PyDict::new_bound(py);
                dict.set_item("node_id", r.node_id)?;
                dict.set_item("energy", r.energy)?;
                dict.set_item("steps", r.steps)?;
                Ok(dict.into())
            })
            .collect()
    }

    #[pyo3(signature = (source_id, target_id, signal_strength, co_activation_count=None, relation_type=None))]
    fn set_edge_weight(
        &self,
        source_id: String,
        target_id: String,
        signal_strength: f64,
        co_activation_count: Option<u64>,
        relation_type: Option<String>,
    ) {
        let rel_type = relation_type.unwrap_or_else(|| "custom".to_string());
        let mut ew = CoreEdgeWeight::new(&source_id, &target_id, &rel_type)
            .with_signal_strength(signal_strength);
        ew.co_activation_count = co_activation_count.unwrap_or(0);
        self.edge_weight_store.write().set_edge_weight(ew);
    }

    fn get_edge_weight(
        &self,
        py: Python<'_>,
        source_id: String,
        target_id: String,
    ) -> PyResult<Option<PyObject>> {
        let store = self.edge_weight_store.read();
        Ok(store.get_edge_weight(&source_id, &target_id).map(|ew| {
            let dict = PyDict::new_bound(py);
            dict.set_item("signal_strength", ew.signal_strength).unwrap();
            dict.set_item("co_activation_count", ew.co_activation_count).unwrap();
            dict.into()
        }))
    }

    fn set_type_default(&self, relation_type: String, signal_strength: f64) {
        let dim = DEFAULT_SEMANTIC_DIM;
        let mut matrix = CoreTransformMatrix::identity(dim);
        for i in 0..dim {
            matrix.data[i * dim + i] = signal_strength as f32;
        }
        self.edge_weight_store.write().set_type_default(&relation_type, matrix);
    }

    fn edge_count(&self) -> usize {
        self.edge_weight_store.read().edge_count()
    }

    fn on_query(&self, result_ids: Vec<String>) {
        let entry = CoreQueryLogEntry::new("hybrid", result_ids.clone());
        self.query_log.record(entry);
        let hebbian = CoreHebbianLearner::new(
            self.edge_weight_store.clone(),
            CoreHebbianConfig::default(),
        );
        hebbian.on_query_result(&result_ids);
    }

    #[pyo3(signature = (epochs=None))]
    fn fit_all(&self, py: Python<'_>, epochs: Option<usize>) -> PyResult<PyObject> {
        let _epoch_count = epochs.unwrap_or(5);
        let hebbian = CoreHebbianLearner::new(
            self.edge_weight_store.clone(),
            CoreHebbianConfig::default(),
        );
        hebbian.decay_scan();

        let contrastive = CoreContrastiveFitter::new(
            self.edge_weight_store.clone(),
            CoreContrastiveConfig::default(),
        );
        let log_entries = self.query_log.entries();
        let all_node_ids: Vec<String> = log_entries.iter()
            .flat_map(|e| e.result_ids.clone())
            .collect::<HashSet<_>>()
            .into_iter()
            .collect();
        let samples = contrastive.generate_samples_from_log(&log_entries, &all_node_ids);
        let result = contrastive.fit(&samples);

        let dict = PyDict::new_bound(py);
        dict.set_item("edges_updated", result.edges_updated)?;
        dict.set_item("final_loss", result.final_loss)?;
        dict.set_item("epochs_completed", result.epochs_completed)?;
        Ok(dict.into())
    }

    #[pyo3(signature = (config=None))]
    fn run_promotion_scan(&self, py: Python<'_>, config: Option<&Bound<'_, PyDict>>) -> PyResult<PyObject> {
        let mut pc = CorePromotionConfig::default();
        if let Some(d) = config {
            if let Ok(Some(v)) = d.get_item("promote_to_type_threshold") {
                if let Ok(val) = v.extract::<u64>() { pc.promote_to_type_threshold = val; }
            }
            if let Ok(Some(v)) = d.get_item("demote_inactive_ms") {
                if let Ok(val) = v.extract::<i64>() { pc.demote_inactive_ms = val; }
            }
            if let Ok(Some(v)) = d.get_item("max_individual_edges") {
                if let Ok(val) = v.extract::<usize>() { pc.max_individual_edges = val; }
            }
        }

        let result = core_run_promotion_scan(&self.edge_weight_store, &pc);

        let dict = PyDict::new_bound(py);
        dict.set_item("promoted", result.promoted)?;
        dict.set_item("demoted", result.demoted)?;
        dict.set_item("pruned", result.pruned)?;
        Ok(dict.into())
    }

    fn backend_info(&self, py: Python<'_>) -> PyResult<PyObject> {
        let info = self.engine.backend_info();
        let dict = PyDict::new_bound(py);
        dict.set_item("name", info.name)?;
        dict.set_item("is_gpu", info.is_gpu)?;
        dict.set_item("max_efficient_dim", info.max_efficient_dim)?;
        Ok(dict.into())
    }

    #[staticmethod]
    fn mobile_config(py: Python<'_>) -> PyResult<PyObject> {
        let c = CoreSiFuConfig::mobile();
        let dict = PyDict::new_bound(py);
        dict.set_item("max_steps", c.max_steps)?;
        dict.set_item("energy_threshold", c.energy_threshold)?;
        dict.set_item("decay_factor", c.decay_factor)?;
        dict.set_item("activation", "relu")?;
        dict.set_item("top_k", c.top_k)?;
        Ok(dict.into())
    }

    #[staticmethod]
    fn server_config(py: Python<'_>) -> PyResult<PyObject> {
        let c = CoreSiFuConfig::server();
        let dict = PyDict::new_bound(py);
        dict.set_item("max_steps", c.max_steps)?;
        dict.set_item("energy_threshold", c.energy_threshold)?;
        dict.set_item("decay_factor", c.decay_factor)?;
        dict.set_item("activation", "relu")?;
        dict.set_item("top_k", c.top_k)?;
        Ok(dict.into())
    }
}

pub fn register_sifu(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<PySiFuEngine>()?;
    Ok(())
}
