use std::collections::BTreeMap;

use serde::{Deserialize, Serialize};
use thiserror::Error;

use crate::{Digest, NodeId, TaskId};

#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub struct VfsPath(String);

impl VfsPath {
    pub fn new(path: impl Into<String>) -> Result<Self, VfsError> {
        let path = path.into();
        if !path.starts_with("/vfs/") {
            return Err(VfsError::InvalidPath(path));
        }
        Ok(Self(path))
    }

    pub fn as_str(&self) -> &str {
        &self.0
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct VfsObject {
    pub path: VfsPath,
    pub digest: Digest,
    pub size: u64,
    pub producer: TaskId,
    pub node: NodeId,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct VfsManifest {
    pub epoch: u64,
    pub producer: TaskId,
    pub node: NodeId,
    pub objects: BTreeMap<VfsPath, VfsObject>,
    pub large_bytes_uploaded: bool,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum SyncPolicy {
    MetadataOnly,
    ExplicitNode(NodeId),
    ExplicitStore(String),
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum VfsSyncDecision {
    NoBytesMoved,
    MoveBytesToNode(NodeId),
    MoveBytesToStore(String),
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub enum ReuseDecision {
    SameNodeZeroCopy,
    NeedsTransfer { from: NodeId, to: NodeId },
    Unavailable,
}

#[derive(Clone, Debug, Error, PartialEq, Eq)]
pub enum VfsError {
    #[error("VFS path must start with /vfs/: {0}")]
    InvalidPath(String),
    #[error("path is not visible in the published VFS manifest: {0}")]
    NotVisible(String),
}

#[derive(Clone, Debug)]
pub struct VfsOverlay {
    task: TaskId,
    node: NodeId,
    epoch: u64,
    pending: BTreeMap<VfsPath, VfsObject>,
    published: BTreeMap<VfsPath, VfsObject>,
}

impl VfsOverlay {
    pub fn new(task: TaskId, node: NodeId) -> Self {
        Self {
            task,
            node,
            epoch: 0,
            pending: BTreeMap::new(),
            published: BTreeMap::new(),
        }
    }

    pub fn write(&mut self, path: VfsPath, digest: Digest, size: u64) -> VfsObject {
        let object = VfsObject {
            path: path.clone(),
            digest,
            size,
            producer: self.task.clone(),
            node: self.node.clone(),
        };
        self.pending.insert(path, object.clone());
        object
    }

    pub fn flush(&mut self) -> VfsManifest {
        self.epoch += 1;
        self.published.append(&mut self.pending);
        VfsManifest {
            epoch: self.epoch,
            producer: self.task.clone(),
            node: self.node.clone(),
            objects: self.published.clone(),
            large_bytes_uploaded: false,
        }
    }

    pub fn sync(&self, policy: SyncPolicy) -> VfsSyncDecision {
        match policy {
            SyncPolicy::MetadataOnly => VfsSyncDecision::NoBytesMoved,
            SyncPolicy::ExplicitNode(node) => VfsSyncDecision::MoveBytesToNode(node),
            SyncPolicy::ExplicitStore(store) => VfsSyncDecision::MoveBytesToStore(store),
        }
    }

    pub fn read_published<'a>(
        manifest: &'a VfsManifest,
        path: &VfsPath,
    ) -> Result<&'a VfsObject, VfsError> {
        manifest
            .objects
            .get(path)
            .ok_or_else(|| VfsError::NotVisible(path.as_str().to_owned()))
    }

    pub fn reuse_for_consumer(
        manifest: &VfsManifest,
        path: &VfsPath,
        consumer_node: &NodeId,
    ) -> ReuseDecision {
        let Some(object) = manifest.objects.get(path) else {
            return ReuseDecision::Unavailable;
        };
        if &object.node == consumer_node {
            ReuseDecision::SameNodeZeroCopy
        } else {
            ReuseDecision::NeedsTransfer {
                from: object.node.clone(),
                to: consumer_node.clone(),
            }
        }
    }

    pub fn discard_unflushed(&mut self) {
        self.pending.clear();
    }

    pub fn pending_len(&self) -> usize {
        self.pending.len()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn path() -> VfsPath {
        VfsPath::new("/vfs/artifacts/app").unwrap()
    }

    #[test]
    fn flush_publishes_manifest_without_large_byte_upload() {
        let mut overlay = VfsOverlay::new(TaskId::from("task"), NodeId::from("node-a"));
        overlay.write(path(), Digest::sha256("binary"), 6);

        let manifest = overlay.flush();

        assert_eq!(manifest.epoch, 1);
        assert!(!manifest.large_bytes_uploaded);
        assert!(manifest.objects.contains_key(&path()));
    }

    #[test]
    fn downstream_task_can_read_after_flush_but_not_before() {
        let mut overlay = VfsOverlay::new(TaskId::from("task"), NodeId::from("node-a"));
        overlay.write(path(), Digest::sha256("binary"), 6);

        let empty = VfsManifest {
            epoch: 0,
            producer: TaskId::from("task"),
            node: NodeId::from("node-a"),
            objects: BTreeMap::new(),
            large_bytes_uploaded: false,
        };
        assert!(VfsOverlay::read_published(&empty, &path()).is_err());

        let manifest = overlay.flush();
        assert!(VfsOverlay::read_published(&manifest, &path()).is_ok());
    }

    #[test]
    fn sync_is_explicit_and_policy_driven() {
        let overlay = VfsOverlay::new(TaskId::from("task"), NodeId::from("node-a"));

        assert_eq!(
            overlay.sync(SyncPolicy::MetadataOnly),
            VfsSyncDecision::NoBytesMoved
        );
        assert_eq!(
            overlay.sync(SyncPolicy::ExplicitStore("s3://bucket/app".to_owned())),
            VfsSyncDecision::MoveBytesToStore("s3://bucket/app".to_owned())
        );
    }

    #[test]
    fn same_node_reuse_avoids_transfer() {
        let mut overlay = VfsOverlay::new(TaskId::from("task"), NodeId::from("node-a"));
        overlay.write(path(), Digest::sha256("binary"), 6);
        let manifest = overlay.flush();

        assert_eq!(
            VfsOverlay::reuse_for_consumer(&manifest, &path(), &NodeId::from("node-a")),
            ReuseDecision::SameNodeZeroCopy
        );
        assert_eq!(
            VfsOverlay::reuse_for_consumer(&manifest, &path(), &NodeId::from("node-b")),
            ReuseDecision::NeedsTransfer {
                from: NodeId::from("node-a"),
                to: NodeId::from("node-b")
            }
        );
    }

    #[test]
    fn unflushed_task_local_changes_can_be_discarded() {
        let mut overlay = VfsOverlay::new(TaskId::from("task"), NodeId::from("node-a"));
        overlay.write(path(), Digest::sha256("binary"), 6);

        overlay.discard_unflushed();

        assert_eq!(overlay.pending_len(), 0);
    }
}