//! Distributed query router.
//!
//! Routes a Cypher query to the appropriate partition(s), executes it on
//! each, and merges the results.
//!
//! ## Execution model
//!
//! - **Single partition**: execute directly on the local backend (if owned)
//!   or via Arrow Flight (if remote).
//! - **Multi partition**: fan-out in parallel, merge result sets.
//! - **All partitions**: broadcast to all known partitions.
//!
//! ## Consistency model (Phase 1)
//! Eventual consistency: queries return the local snapshot; writes are
//! accepted immediately and propagated asynchronously.

#[cfg(feature = "distributed")]
use std::sync::Arc;

#[cfg(feature = "distributed")]
use tokio::sync::RwLock;

#[cfg(feature = "distributed")]
use crate::distributed::partition::location_service::PartitionLocationService;
#[cfg(feature = "distributed")]
use crate::distributed::partition::PartitionMap;
#[cfg(feature = "distributed")]
use crate::distributed::query::partition_analyzer::{PartitionAnalyzer, QueryDistribution};
#[cfg(feature = "distributed")]
use crate::error::CypherResult;

// ── Router ────────────────────────────────────────────────────────────────────

/// Routes distributed Cypher queries to the correct partition(s).
#[cfg(feature = "distributed")]
pub struct QueryRouter {
    location_svc:  Arc<RwLock<PartitionLocationService>>,
    partition_map: Arc<std::sync::RwLock<PartitionMap>>,
    analyzer:      PartitionAnalyzer,
}

#[cfg(feature = "distributed")]
impl QueryRouter {
    pub fn new(
        location_svc:  Arc<RwLock<PartitionLocationService>>,
        partition_map: Arc<std::sync::RwLock<PartitionMap>>,
    ) -> Self {
        let analyzer = {
            let pm = partition_map.read().expect("partition map poisoned");
            PartitionAnalyzer::from_partition_map(&pm)
        };
        Self { location_svc, partition_map, analyzer }
    }

    /// Rebuild the partition analyzer from the current partition map.
    /// Call this after partition membership changes (rebalancing, recovery).
    pub fn refresh_analyzer(&mut self) {
        let pm = self.partition_map.read().expect("partition map poisoned");
        self.analyzer = PartitionAnalyzer::from_partition_map(&pm);
    }

    /// Determine which partitions a query touches.
    pub fn distribution(&self, query: &str) -> QueryDistribution {
        let pm = self.partition_map.read().expect("partition map poisoned");
        self.analyzer.analyze(query, &pm)
    }

    /// Return the partition IDs that need to be queried for `query`.
    pub fn target_partitions(&self, query: &str) -> Vec<u32> {
        let pm = self.partition_map.read().expect("partition map poisoned");
        match self.analyzer.analyze(query, &pm) {
            QueryDistribution::SinglePartition(id) => vec![id],
            QueryDistribution::MultiPartition(ids) => ids,
            QueryDistribution::AllPartitions => {
                (0..pm.num_partitions).collect()
            }
        }
    }

    /// Return the Flight endpoint URL for `partition_id`.
    pub async fn endpoint_for_partition(&self, partition_id: u32) -> Option<String> {
        self.location_svc
            .read()
            .await
            .get_partition_location(partition_id)
            .await
            .map(|loc| {
                // Convert pod name to Flight endpoint
                // e.g. "ocg-partition-2" → "http://ocg-partition-2:8815"
                format!("http://{}:8815", loc.owner)
            })
    }

    /// Stub: execute a query on a list of partition endpoints and merge results.
    ///
    /// In production this fans out via `GraphFlightClient`, executes the query
    /// on each partition pod, and merges the `QueryResult`s. For Phase 11 this
    /// returns an empty merged result to indicate the routing decision.
    pub async fn route(&self, query: &str) -> CypherResult<RoutingDecision> {
        let targets = self.target_partitions(query);
        let mut endpoints = Vec::new();

        for pid in &targets {
            if let Some(ep) = self.endpoint_for_partition(*pid).await {
                endpoints.push((*pid, ep));
            }
        }

        Ok(RoutingDecision {
            query:     query.to_string(),
            targets:   targets.clone(),
            endpoints,
        })
    }
}

// ── Routing decision ──────────────────────────────────────────────────────────

/// The result of routing a query — which partitions and endpoints to use.
#[cfg(feature = "distributed")]
#[derive(Debug, Clone)]
pub struct RoutingDecision {
    pub query:     String,
    pub targets:   Vec<u32>,
    pub endpoints: Vec<(u32, String)>,
}

// ── Tests ─────────────────────────────────────────────────────────────────────

#[cfg(all(test, feature = "distributed"))]
mod tests {
    use super::*;
    use crate::distributed::partition::PartitionMap;

    fn make_partition_map_arc(num: u32) -> Arc<std::sync::RwLock<PartitionMap>> {
        let mut pm = PartitionMap::new(num);
        pm.add_label(0, "Person".to_string());
        pm.add_label(1, "Company".to_string());
        if num > 2 { pm.add_label(2, "Product".to_string()); }
        Arc::new(std::sync::RwLock::new(pm))
    }

    #[test]
    fn test_target_partitions_single_label() {
        let pm = make_partition_map_arc(2);
        // We need a location service; use a mock by testing distribution directly
        let pm_ref = pm.read().unwrap();
        let analyzer = PartitionAnalyzer::from_partition_map(&pm_ref);
        let dist = analyzer.analyze("MATCH (n:Person) RETURN n", &pm_ref);
        assert_eq!(dist, QueryDistribution::SinglePartition(0));
    }

    #[test]
    fn test_target_partitions_all_on_no_label() {
        let pm = make_partition_map_arc(3);
        let pm_ref = pm.read().unwrap();
        let analyzer = PartitionAnalyzer::from_partition_map(&pm_ref);
        let dist = analyzer.analyze("MATCH (n) RETURN count(n)", &pm_ref);
        assert_eq!(dist, QueryDistribution::AllPartitions);
    }

    #[test]
    fn test_routing_decision_fields() {
        let decision = RoutingDecision {
            query:     "MATCH (n) RETURN n".to_string(),
            targets:   vec![0, 1, 2],
            endpoints: vec![(0, "http://pod-0:8815".to_string())],
        };
        assert_eq!(decision.targets.len(), 3);
        assert_eq!(decision.endpoints.len(), 1);
    }
}
