//! Storage-backed state watcher for PlasmaVMC.
//!
//! PlasmaVMC persists VM intent and node heartbeats in the shared metadata
//! store. This watcher polls that store and mirrors external changes into each
//! process-local cache so multiple control-plane or agent instances converge on
//! the same durable view.

use crate::storage::VmStore;
use plasmavmc_types::{Node, VirtualMachine, VmHandle};
use serde::Serialize;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::mpsc;
use tokio::time::MissedTickBehavior;
use tracing::{debug, info, warn};

type VmSnapshot = HashMap<VmKey, VirtualMachine>;
type NodeSnapshot = HashMap<String, Node>;

/// Event types from the state watcher
#[derive(Debug, Clone)]
pub enum StateEvent {
    /// A VM was created or updated
    VmUpdated {
        org_id: String,
        project_id: String,
        vm_id: String,
        vm: VirtualMachine,
    },
    /// A VM was deleted
    VmDeleted {
        org_id: String,
        project_id: String,
        vm_id: String,
    },
    /// A VM handle was updated
    HandleUpdated {
        org_id: String,
        project_id: String,
        vm_id: String,
        handle: VmHandle,
    },
    /// A VM handle was deleted
    HandleDeleted {
        org_id: String,
        project_id: String,
        vm_id: String,
    },
    /// A node was updated
    NodeUpdated { node_id: String, node: Node },
    /// A node was deleted
    NodeDeleted { node_id: String },
}

/// Configuration for the state watcher
#[derive(Debug, Clone)]
pub struct WatcherConfig {
    /// Poll interval for metadata refresh.
    pub poll_interval: Duration,
    /// Channel buffer size for events
    pub buffer_size: usize,
}

impl Default for WatcherConfig {
    fn default() -> Self {
        Self {
            poll_interval: Duration::from_millis(1000),
            buffer_size: 256,
        }
    }
}

/// State watcher that monitors the shared VM store for external changes.
pub struct StateWatcher {
    store: Arc<dyn VmStore>,
    config: WatcherConfig,
    event_tx: mpsc::Sender<StateEvent>,
}

impl StateWatcher {
    /// Create a new state watcher and return the event receiver.
    pub fn new(
        store: Arc<dyn VmStore>,
        config: WatcherConfig,
    ) -> (Self, mpsc::Receiver<StateEvent>) {
        let (event_tx, event_rx) = mpsc::channel(config.buffer_size);
        (
            Self {
                store,
                config,
                event_tx,
            },
            event_rx,
        )
    }

    /// Start watching for state changes.
    ///
    /// VM handles remain local process state because they point at ephemeral
    /// runtime artifacts such as QMP sockets and per-node directories. Durable
    /// VM intent and node heartbeats are the shared state synchronized here.
    pub async fn start(&self) -> Result<(), WatcherError> {
        info!(
            poll_interval_ms = self.config.poll_interval.as_millis(),
            "Starting storage-backed PlasmaVMC state watcher"
        );

        let mut vm_snapshot = self.load_vms().await?;
        let mut node_snapshot = self.load_nodes().await?;
        let mut ticker = tokio::time::interval(self.config.poll_interval);
        ticker.set_missed_tick_behavior(MissedTickBehavior::Delay);

        loop {
            ticker.tick().await;

            let next_vm_snapshot = match self.load_vms().await {
                Ok(snapshot) => snapshot,
                Err(error) => {
                    warn!(error = %error, "Failed to refresh PlasmaVMC VM snapshot");
                    continue;
                }
            };
            match diff_vm_snapshots(&vm_snapshot, &next_vm_snapshot) {
                Ok(events) => {
                    for event in events {
                        if self.event_tx.send(event).await.is_err() {
                            info!("State watcher receiver dropped, stopping");
                            return Ok(());
                        }
                    }
                    vm_snapshot = next_vm_snapshot;
                }
                Err(error) => {
                    warn!(error = %error, "Failed to diff PlasmaVMC VM snapshot");
                }
            }

            let next_node_snapshot = match self.load_nodes().await {
                Ok(snapshot) => snapshot,
                Err(error) => {
                    warn!(error = %error, "Failed to refresh PlasmaVMC node snapshot");
                    continue;
                }
            };
            match diff_node_snapshots(&node_snapshot, &next_node_snapshot) {
                Ok(events) => {
                    for event in events {
                        if self.event_tx.send(event).await.is_err() {
                            info!("State watcher receiver dropped, stopping");
                            return Ok(());
                        }
                    }
                    node_snapshot = next_node_snapshot;
                }
                Err(error) => {
                    warn!(error = %error, "Failed to diff PlasmaVMC node snapshot");
                }
            }
        }
    }

    async fn load_vms(&self) -> Result<VmSnapshot, WatcherError> {
        let vms = self
            .store
            .list_all_vms()
            .await
            .map_err(|error| WatcherError::Storage(error.to_string()))?;
        let mut snapshot = HashMap::with_capacity(vms.len());
        for vm in vms {
            snapshot.insert(VmKey::from_vm(&vm), vm);
        }
        Ok(snapshot)
    }

    async fn load_nodes(&self) -> Result<NodeSnapshot, WatcherError> {
        let nodes = self
            .store
            .list_nodes()
            .await
            .map_err(|error| WatcherError::Storage(error.to_string()))?;
        let mut snapshot = HashMap::with_capacity(nodes.len());
        for node in nodes {
            snapshot.insert(node.id.to_string(), node);
        }
        Ok(snapshot)
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
struct VmKey {
    org_id: String,
    project_id: String,
    vm_id: String,
}

impl VmKey {
    fn from_vm(vm: &VirtualMachine) -> Self {
        Self {
            org_id: vm.org_id.clone(),
            project_id: vm.project_id.clone(),
            vm_id: vm.id.to_string(),
        }
    }
}

fn diff_vm_snapshots(
    previous: &VmSnapshot,
    current: &VmSnapshot,
) -> Result<Vec<StateEvent>, WatcherError> {
    let mut events = Vec::new();

    for (key, vm) in current {
        let changed = match previous.get(key) {
            Some(old_vm) => values_differ(old_vm, vm)?,
            None => true,
        };
        if changed {
            events.push(StateEvent::VmUpdated {
                org_id: key.org_id.clone(),
                project_id: key.project_id.clone(),
                vm_id: key.vm_id.clone(),
                vm: vm.clone(),
            });
        }
    }

    for key in previous.keys() {
        if !current.contains_key(key) {
            events.push(StateEvent::VmDeleted {
                org_id: key.org_id.clone(),
                project_id: key.project_id.clone(),
                vm_id: key.vm_id.clone(),
            });
        }
    }

    Ok(events)
}

fn diff_node_snapshots(
    previous: &NodeSnapshot,
    current: &NodeSnapshot,
) -> Result<Vec<StateEvent>, WatcherError> {
    let mut events = Vec::new();

    for (node_id, node) in current {
        let changed = match previous.get(node_id) {
            Some(old_node) => values_differ(old_node, node)?,
            None => true,
        };
        if changed {
            events.push(StateEvent::NodeUpdated {
                node_id: node_id.clone(),
                node: node.clone(),
            });
        }
    }

    for node_id in previous.keys() {
        if !current.contains_key(node_id) {
            events.push(StateEvent::NodeDeleted {
                node_id: node_id.clone(),
            });
        }
    }

    Ok(events)
}

fn values_differ<T: Serialize>(lhs: &T, rhs: &T) -> Result<bool, WatcherError> {
    let lhs =
        serde_json::to_value(lhs).map_err(|error| WatcherError::Serialize(error.to_string()))?;
    let rhs =
        serde_json::to_value(rhs).map_err(|error| WatcherError::Serialize(error.to_string()))?;
    Ok(lhs != rhs)
}

/// Watcher errors
#[derive(Debug, thiserror::Error)]
pub enum WatcherError {
    #[error("Storage error: {0}")]
    Storage(String),
    #[error("Serialization error: {0}")]
    Serialize(String),
}

/// State synchronizer that applies watch events to local state
pub struct StateSynchronizer<S: StateSink> {
    sink: Arc<S>,
}

/// Trait for applying state changes
pub trait StateSink: Send + Sync {
    /// Called when a VM is updated externally
    fn on_vm_updated(&self, org_id: &str, project_id: &str, vm_id: &str, vm: VirtualMachine);
    /// Called when a VM is deleted externally
    fn on_vm_deleted(&self, org_id: &str, project_id: &str, vm_id: &str);
    /// Called when a handle is updated externally
    fn on_handle_updated(&self, org_id: &str, project_id: &str, vm_id: &str, handle: VmHandle);
    /// Called when a handle is deleted externally
    fn on_handle_deleted(&self, org_id: &str, project_id: &str, vm_id: &str);
    /// Called when a node is updated externally
    fn on_node_updated(&self, node_id: &str, node: Node);
    /// Called when a node is deleted externally
    fn on_node_deleted(&self, node_id: &str);
}

impl<S: StateSink> StateSynchronizer<S> {
    /// Create a new state synchronizer
    pub fn new(sink: Arc<S>) -> Self {
        Self { sink }
    }

    /// Process events from the watcher
    pub async fn run(&self, mut event_rx: mpsc::Receiver<StateEvent>) {
        info!("Starting state synchronizer");

        while let Some(event) = event_rx.recv().await {
            match event {
                StateEvent::VmUpdated {
                    org_id,
                    project_id,
                    vm_id,
                    vm,
                } => {
                    debug!(org_id, project_id, vm_id, "External VM update received");
                    self.sink.on_vm_updated(&org_id, &project_id, &vm_id, vm);
                }
                StateEvent::VmDeleted {
                    org_id,
                    project_id,
                    vm_id,
                } => {
                    debug!(org_id, project_id, vm_id, "External VM deletion received");
                    self.sink.on_vm_deleted(&org_id, &project_id, &vm_id);
                }
                StateEvent::HandleUpdated {
                    org_id,
                    project_id,
                    vm_id,
                    handle,
                } => {
                    debug!(org_id, project_id, vm_id, "External handle update received");
                    self.sink
                        .on_handle_updated(&org_id, &project_id, &vm_id, handle);
                }
                StateEvent::HandleDeleted {
                    org_id,
                    project_id,
                    vm_id,
                } => {
                    debug!(
                        org_id,
                        project_id, vm_id, "External handle deletion received"
                    );
                    self.sink.on_handle_deleted(&org_id, &project_id, &vm_id);
                }
                StateEvent::NodeUpdated { node_id, node } => {
                    debug!(node_id, "External node update received");
                    self.sink.on_node_updated(&node_id, node);
                }
                StateEvent::NodeDeleted { node_id } => {
                    debug!(node_id, "External node deletion received");
                    self.sink.on_node_deleted(&node_id);
                }
            }
        }

        info!("State synchronizer stopped");
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use plasmavmc_types::{VmSpec, VmState};

    fn sample_vm() -> VirtualMachine {
        VirtualMachine::new("vm-a", "org-a", "project-a", VmSpec::default())
    }

    fn sample_node() -> Node {
        Node::new("node-a")
    }

    #[test]
    fn vm_snapshot_diff_emits_update_for_new_and_changed_vms() {
        let mut previous = VmSnapshot::new();
        let mut current = VmSnapshot::new();
        let vm = sample_vm();
        let key = VmKey::from_vm(&vm);

        current.insert(key.clone(), vm.clone());
        let events = diff_vm_snapshots(&previous, &current).unwrap();
        assert!(matches!(
            events.as_slice(),
            [StateEvent::VmUpdated { vm_id, .. }] if vm_id == &vm.id.to_string()
        ));

        previous = current.clone();
        let mut updated_vm = vm.clone();
        updated_vm.state = VmState::Running;
        updated_vm.status.actual_state = VmState::Running;
        current.insert(key, updated_vm.clone());
        let events = diff_vm_snapshots(&previous, &current).unwrap();
        assert!(matches!(
            events.as_slice(),
            [StateEvent::VmUpdated { vm, .. }] if vm.state == VmState::Running
        ));
    }

    #[test]
    fn vm_snapshot_diff_emits_delete_for_removed_vms() {
        let vm = sample_vm();
        let key = VmKey::from_vm(&vm);
        let previous = HashMap::from([(key, vm.clone())]);
        let current = VmSnapshot::new();

        let events = diff_vm_snapshots(&previous, &current).unwrap();
        assert!(matches!(
            events.as_slice(),
            [StateEvent::VmDeleted { vm_id, .. }] if vm_id == &vm.id.to_string()
        ));
    }

    #[test]
    fn node_snapshot_diff_emits_update_and_delete() {
        let node = sample_node();
        let mut previous = NodeSnapshot::new();
        let mut current = NodeSnapshot::from([(node.id.to_string(), node.clone())]);

        let events = diff_node_snapshots(&previous, &current).unwrap();
        assert!(matches!(
            events.as_slice(),
            [StateEvent::NodeUpdated { node_id, .. }] if node_id == &node.id.to_string()
        ));

        previous = current.clone();
        let mut updated_node = node.clone();
        updated_node.last_heartbeat = 42;
        current.insert(updated_node.id.to_string(), updated_node);
        let events = diff_node_snapshots(&previous, &current).unwrap();
        assert!(matches!(
            events.as_slice(),
            [StateEvent::NodeUpdated { node, .. }] if node.last_heartbeat == 42
        ));

        let events = diff_node_snapshots(&current, &NodeSnapshot::new()).unwrap();
        assert!(matches!(
            events.as_slice(),
            [StateEvent::NodeDeleted { node_id }] if node_id == &node.id.to_string()
        ));
    }
}
