//! Node operations: writes, sync, routing, and Merkle operations.

use super::{rand_hash, ChaosNode};
use crate::error::Result;
use std::collections::HashSet;
use stemedb_cluster::sharding::MetaRange;
use stemedb_core::serde::serialize;
use stemedb_core::testing::AssertionBuilder;
use stemedb_merkle::MerkleTree;
use stemedb_storage::crdt::AssertionTransfer;
use tracing::{debug, instrument, warn};

impl ChaosNode {
    /// Initializes sharding with the given node list.
    ///
    /// # Arguments
    ///
    /// * `nodes` - All node IDs in the cluster
    /// * `num_shards` - Number of shards to create
    /// * `replication_factor` - Number of replicas per shard
    pub fn init_shards(
        &self,
        nodes: &[stemedb_cluster::membership::NodeId],
        num_shards: u32,
        replication_factor: u32,
    ) {
        let meta = MetaRange::with_initial_shards(num_shards, nodes, replication_factor);
        self.router.update_meta_range(meta);
    }

    /// Writes an assertion to this node.
    ///
    /// Uses the skewed HLC clock for timestamp generation.
    ///
    /// # Arguments
    ///
    /// * `subject` - Assertion subject
    /// * `predicate` - Assertion predicate
    /// * `logical_time` - Logical timestamp (for backward compatibility)
    ///
    /// # Returns
    ///
    /// The assertion hash (BLAKE3).
    #[instrument(skip(self), fields(node = self.index))]
    pub async fn write_assertion(
        &mut self,
        subject: &str,
        predicate: &str,
        logical_time: u64,
    ) -> Result<[u8; 32]> {
        if !self.alive {
            return Err(crate::error::ChaosError::Network("Node is down".to_string()));
        }

        // Generate HLC timestamp with potential skew
        let hlc_timestamp = self.hlc.new_timestamp();

        let assertion = AssertionBuilder::new()
            .subject(subject)
            .predicate(predicate)
            .hlc_timestamp(hlc_timestamp)
            .timestamp(logical_time)
            .source_hash(rand_hash())
            .build();

        let data = serialize(&assertion).map_err(|e| {
            crate::error::ChaosError::Sync(format!("Failed to serialize assertion: {e}"))
        })?;

        let hash = self.crdt_store.put_assertion(subject, &data).await?;

        self.merkle_tree.insert(hash)?;
        self.hash_to_data.insert(hash, (subject.to_string(), data));

        debug!(
            node = self.index,
            subject,
            hash = %hex::encode(&hash[..8]),
            "Wrote assertion"
        );

        Ok(hash)
    }

    /// Checks if this node can accept a write for the given subject.
    ///
    /// Based on shard routing and replica assignment.
    pub fn can_accept_write(&self, subject: &str) -> bool {
        if !self.alive {
            return false;
        }

        let shard_id = match self.router.route_subject(subject) {
            Ok(id) => id,
            Err(_) => return false,
        };

        match self.router.get_replicas(shard_id) {
            Ok(replicas) => replicas.contains(&self.id),
            Err(_) => false,
        }
    }

    /// Returns all assertion hashes (Merkle tree leaves).
    pub fn leaves(&self) -> Vec<[u8; 32]> {
        self.merkle_tree.leaves().to_vec()
    }

    /// Returns the number of assertions stored.
    pub fn assertion_count(&self) -> usize {
        self.merkle_tree.leaves().len()
    }

    /// Computes the canonical Merkle root for convergence verification.
    ///
    /// The canonical root is computed from sorted leaves, ensuring that
    /// nodes with the same assertions produce the same root regardless
    /// of insertion order.
    pub fn canonical_merkle_root(&self) -> Option<[u8; 32]> {
        let mut sorted_leaves = self.merkle_tree.leaves().to_vec();
        if sorted_leaves.is_empty() {
            return None;
        }
        sorted_leaves.sort();

        let mut canonical = MerkleTree::new();
        for leaf in sorted_leaves {
            if canonical.insert(leaf).is_err() {
                return None;
            }
        }
        canonical.root().ok()
    }

    /// Syncs data from another node, respecting fault injection settings.
    ///
    /// This method checks for:
    /// 1. Network partitions (blocks sync if partitioned)
    /// 2. Network latency (delays sync)
    /// 3. Node liveness (fails if either node is down)
    ///
    /// # Arguments
    ///
    /// * `other` - The node to sync from
    ///
    /// # Returns
    ///
    /// Number of assertions merged.
    #[instrument(skip(self, other), fields(from = other.index, to = self.index))]
    pub async fn sync_from(&mut self, other: &ChaosNode) -> Result<usize> {
        // Check if either node is down
        if !self.alive {
            warn!(node = self.index, "Cannot sync - node is down");
            return Err(crate::error::ChaosError::Network("Target node is down".to_string()));
        }
        if !other.alive {
            warn!(node = other.index, "Cannot sync - source node is down");
            return Err(crate::error::ChaosError::Network("Source node is down".to_string()));
        }

        // Check for network partition
        if self.network_controller.should_drop(other.index, self.index) {
            debug!(from = other.index, to = self.index, "Sync blocked by network partition/drop");
            return Ok(0);
        }

        // Apply network latency
        self.network_controller.apply_effects(other.index, self.index).await;

        // Perform sync
        let my_leaves: HashSet<_> = self.leaves().into_iter().collect();
        let mut merged_count = 0;

        for hash in other.leaves() {
            if !my_leaves.contains(&hash) {
                if let Some((subject, data)) = other.hash_to_data.get(&hash) {
                    let transfer = AssertionTransfer { hash, data: data.clone() };

                    if self
                        .crdt_store
                        .merge_with_data(subject, std::slice::from_ref(&transfer))
                        .await
                        .is_ok()
                    {
                        self.merkle_tree.insert(hash)?;
                        self.hash_to_data.insert(hash, (subject.clone(), data.clone()));
                        merged_count += 1;
                    }
                }
            }
        }

        debug!(from = other.index, to = self.index, merged_count, "Sync completed");

        Ok(merged_count)
    }

    /// Returns all assertion data for a specific subject.
    pub fn get_assertions_for_subject(&self, subject: &str) -> Vec<([u8; 32], Vec<u8>)> {
        self.hash_to_data
            .iter()
            .filter(|(_, (s, _))| s == subject)
            .map(|(h, (_, d))| (*h, d.clone()))
            .collect()
    }
}
