use std::collections::{BTreeMap, BTreeSet};

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

use crate::{
    auth::same_tenant_project, Actor, ArtifactId, AuthContext, Digest, LimitError, LimitKind,
    NodeId, ProcessId, ProjectId, ResourceLimits, ResourceMeter, Scope, TaskId, TenantId,
};

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

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct RetentionPolicy {
    pub best_effort_node_retention: bool,
    pub max_download_bytes: u64,
}

impl Default for RetentionPolicy {
    fn default() -> Self {
        Self {
            best_effort_node_retention: true,
            max_download_bytes: 256 * 1024 * 1024,
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct DownloadPolicy {
    pub max_bytes: u64,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct DownloadAction {
    pub artifact: ArtifactId,
    pub source: StorageLocation,
    pub scoped_token_subject: String,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct DownloadLink {
    pub artifact: ArtifactId,
    pub source: StorageLocation,
    pub url_path: String,
    pub scoped_token_digest: Digest,
    pub expires_at_epoch_seconds: u64,
    pub tenant: TenantId,
    pub project: ProjectId,
    pub process: ProcessId,
    pub actor: Actor,
    pub max_bytes: u64,
    pub policy_context_digest: Digest,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct IssuedDownloadLink {
    pub link: DownloadLink,
    pub revoked: bool,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ArtifactDownloadStream {
    pub link: DownloadLink,
    pub streamed_bytes: u64,
}

#[derive(Clone, Debug, Error, PartialEq, Eq)]
pub enum DownloadError {
    #[error("artifact does not exist")]
    NotFound,
    #[error("artifact is unavailable from current retention or explicit storage")]
    Unavailable,
    #[error("artifact download direct connectivity unavailable: {0}")]
    DirectConnectivityUnavailable(String),
    #[error("artifact download denied: {0}")]
    Unauthorized(String),
    #[error("artifact size {size} exceeds download limit {limit}")]
    LimitExceeded { size: u64, limit: u64 },
    #[error("download link token is invalid for this scoped artifact link")]
    InvalidToken,
    #[error("download link has expired")]
    Expired,
    #[error("download link has been revoked")]
    Revoked,
    #[error("download usage limit failed: {0}")]
    Usage(String),
}

#[derive(Clone, Debug, Error, PartialEq, Eq)]
#[error("artifact is unavailable because node-local unsynced bytes were lost")]
pub struct ArtifactUnavailable;

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ArtifactMetadata {
    pub id: ArtifactId,
    pub tenant: TenantId,
    pub project: ProjectId,
    pub process: ProcessId,
    pub producer_task: TaskId,
    pub producer_node: NodeId,
    pub digest: Digest,
    pub size: u64,
    pub flushed_epoch: u64,
    pub retaining_nodes: BTreeSet<NodeId>,
    pub explicit_locations: Vec<String>,
    pub coordinator_has_large_bytes: bool,
}

#[derive(Clone, Debug, Default)]
pub struct ArtifactRegistry {
    artifacts: BTreeMap<ArtifactId, ArtifactMetadata>,
    issued_download_links: BTreeMap<Digest, IssuedDownloadLink>,
    next_epoch: u64,
}

impl ArtifactRegistry {
    pub fn flush_metadata(
        &mut self,
        id: ArtifactId,
        tenant: TenantId,
        project: ProjectId,
        process: ProcessId,
        producer_task: TaskId,
        retaining_node: NodeId,
        digest: Digest,
        size: u64,
    ) -> ArtifactMetadata {
        self.next_epoch += 1;
        let metadata = ArtifactMetadata {
            id: id.clone(),
            tenant,
            project,
            process,
            producer_task,
            producer_node: retaining_node.clone(),
            digest,
            size,
            flushed_epoch: self.next_epoch,
            retaining_nodes: BTreeSet::from([retaining_node]),
            explicit_locations: Vec::new(),
            coordinator_has_large_bytes: false,
        };
        self.artifacts.insert(id, metadata.clone());
        metadata
    }

    pub fn sync_to_explicit_store(
        &mut self,
        artifact: &ArtifactId,
        location: impl Into<String>,
    ) -> Result<(), ArtifactUnavailable> {
        let metadata = self
            .artifacts
            .get_mut(artifact)
            .ok_or(ArtifactUnavailable)?;
        metadata.explicit_locations.push(location.into());
        Ok(())
    }

    pub fn garbage_collect_node(&mut self, node: &NodeId) {
        for metadata in self.artifacts.values_mut() {
            metadata.retaining_nodes.remove(node);
        }
    }

    pub fn metadata(&self, artifact: &ArtifactId) -> Option<&ArtifactMetadata> {
        self.artifacts.get(artifact)
    }

    pub fn download_action(
        &self,
        context: &AuthContext,
        artifact: &ArtifactId,
        policy: &DownloadPolicy,
    ) -> Result<DownloadAction, DownloadError> {
        let metadata = self
            .artifacts
            .get(artifact)
            .ok_or(DownloadError::NotFound)?;
        let scope = Scope {
            tenant: metadata.tenant.clone(),
            project: metadata.project.clone(),
            process: Some(metadata.process.clone()),
            task: Some(metadata.producer_task.clone()),
            node: None,
            artifact: Some(metadata.id.clone()),
        };
        let authz = same_tenant_project(context, &scope);
        if !authz.allowed {
            return Err(DownloadError::Unauthorized(authz.reason));
        }
        if metadata.size > policy.max_bytes {
            return Err(DownloadError::LimitExceeded {
                size: metadata.size,
                limit: policy.max_bytes,
            });
        }
        let source = metadata
            .retaining_nodes
            .iter()
            .next()
            .cloned()
            .map(StorageLocation::RetainedNode)
            .or_else(|| {
                metadata
                    .explicit_locations
                    .first()
                    .cloned()
                    .map(StorageLocation::ExplicitStore)
            })
            .ok_or(DownloadError::Unavailable)?;

        Ok(DownloadAction {
            artifact: artifact.clone(),
            source,
            scoped_token_subject: format!(
                "{}/{}/{}/{}",
                metadata.tenant, metadata.project, metadata.process, metadata.id
            ),
        })
    }

    pub fn downloadable_size(
        &self,
        context: &AuthContext,
        artifact: &ArtifactId,
        policy: &DownloadPolicy,
    ) -> Result<u64, DownloadError> {
        self.download_action(context, artifact, policy)?;
        let metadata = self
            .artifacts
            .get(artifact)
            .ok_or(DownloadError::NotFound)?;
        Ok(metadata.size)
    }

    pub fn create_download_link(
        &mut self,
        context: &AuthContext,
        artifact: &ArtifactId,
        policy: &DownloadPolicy,
        token_nonce: &str,
        now_epoch_seconds: u64,
        ttl_seconds: u64,
    ) -> Result<DownloadLink, DownloadError> {
        let action = self.download_action(context, artifact, policy)?;
        let metadata = self
            .artifacts
            .get(artifact)
            .ok_or(DownloadError::NotFound)?;
        let expires_at_epoch_seconds = now_epoch_seconds.saturating_add(ttl_seconds);
        let policy_context_digest =
            download_policy_context_digest(metadata, &action.source, policy);
        let scoped_token_digest = Digest::from_parts([
            b"artifact-download-token:v2".as_slice(),
            action.scoped_token_subject.as_bytes(),
            actor_subject(&context.actor).as_bytes(),
            token_nonce.as_bytes(),
            metadata.digest.as_str().as_bytes(),
            metadata.size.to_string().as_bytes(),
            policy_context_digest.as_str().as_bytes(),
            expires_at_epoch_seconds.to_string().as_bytes(),
        ]);
        let link = DownloadLink {
            artifact: artifact.clone(),
            source: action.source,
            url_path: format!(
                "/artifacts/{}/{}/{}/{}",
                metadata.tenant, metadata.project, metadata.process, metadata.id
            ),
            scoped_token_digest,
            expires_at_epoch_seconds,
            tenant: metadata.tenant.clone(),
            project: metadata.project.clone(),
            process: metadata.process.clone(),
            actor: context.actor.clone(),
            max_bytes: policy.max_bytes,
            policy_context_digest,
        };
        self.issued_download_links.insert(
            link.scoped_token_digest.clone(),
            IssuedDownloadLink {
                link: link.clone(),
                revoked: false,
            },
        );
        Ok(link)
    }

    pub fn revoke_download_link(
        &mut self,
        context: &AuthContext,
        artifact: &ArtifactId,
        presented_token_digest: &Digest,
    ) -> Result<DownloadLink, DownloadError> {
        let issued = self
            .issued_download_links
            .get(presented_token_digest)
            .ok_or(DownloadError::InvalidToken)?;
        if issued.link.artifact != *artifact || issued.link.actor != context.actor {
            return Err(DownloadError::InvalidToken);
        }
        self.download_action(
            context,
            artifact,
            &DownloadPolicy {
                max_bytes: issued.link.max_bytes,
            },
        )?;

        let issued = self
            .issued_download_links
            .get_mut(presented_token_digest)
            .ok_or(DownloadError::InvalidToken)?;
        issued.revoked = true;
        Ok(issued.link.clone())
    }

    pub fn open_download_stream(
        &self,
        context: &AuthContext,
        artifact: &ArtifactId,
        policy: &DownloadPolicy,
        presented_token_digest: &Digest,
        now_epoch_seconds: u64,
        limits: &ResourceLimits,
        meter: &mut ResourceMeter,
    ) -> Result<ArtifactDownloadStream, DownloadError> {
        let issued = self
            .issued_download_links
            .get(presented_token_digest)
            .ok_or(DownloadError::InvalidToken)?;
        if issued.link.artifact != *artifact
            || issued.link.max_bytes != policy.max_bytes
            || issued.link.actor != context.actor
        {
            return Err(DownloadError::InvalidToken);
        }
        if issued.revoked {
            return Err(DownloadError::Revoked);
        }
        if now_epoch_seconds > issued.link.expires_at_epoch_seconds {
            return Err(DownloadError::Expired);
        }
        let action = self.download_action(context, artifact, policy)?;
        if action.source != issued.link.source {
            return Err(DownloadError::Unavailable);
        }
        let metadata = self
            .artifacts
            .get(artifact)
            .ok_or(DownloadError::NotFound)?;
        if download_policy_context_digest(metadata, &action.source, policy)
            != issued.link.policy_context_digest
        {
            return Err(DownloadError::InvalidToken);
        }
        meter
            .charge(limits, LimitKind::ArtifactDownloadBytes, 0)
            .map_err(download_usage_error)?;
        Ok(ArtifactDownloadStream {
            link: issued.link.clone(),
            streamed_bytes: 0,
        })
    }

    pub fn stream_download_chunk(
        &self,
        stream: &mut ArtifactDownloadStream,
        limits: &ResourceLimits,
        meter: &mut ResourceMeter,
        bytes: u64,
    ) -> Result<(), DownloadError> {
        let metadata = self
            .artifacts
            .get(&stream.link.artifact)
            .ok_or(DownloadError::NotFound)?;
        if !source_is_available(metadata, &stream.link.source) {
            return Err(DownloadError::Unavailable);
        }
        stream.stream_chunk(limits, meter, bytes)
    }
}

impl ArtifactDownloadStream {
    pub fn stream_chunk(
        &mut self,
        limits: &ResourceLimits,
        meter: &mut ResourceMeter,
        bytes: u64,
    ) -> Result<(), DownloadError> {
        if self.streamed_bytes.saturating_add(bytes) > self.link.max_bytes {
            return Err(DownloadError::LimitExceeded {
                size: self.streamed_bytes.saturating_add(bytes),
                limit: self.link.max_bytes,
            });
        }
        meter
            .charge(limits, LimitKind::ArtifactDownloadBytes, bytes)
            .map_err(download_usage_error)?;
        self.streamed_bytes += bytes;
        Ok(())
    }
}

fn download_usage_error(error: LimitError) -> DownloadError {
    DownloadError::Usage(error.to_string())
}

fn download_policy_context_digest(
    metadata: &ArtifactMetadata,
    source: &StorageLocation,
    policy: &DownloadPolicy,
) -> Digest {
    Digest::from_parts([
        b"artifact-download-policy-context:v1".as_slice(),
        metadata.tenant.as_str().as_bytes(),
        metadata.project.as_str().as_bytes(),
        metadata.process.as_str().as_bytes(),
        metadata.id.as_str().as_bytes(),
        metadata.digest.as_str().as_bytes(),
        metadata.size.to_string().as_bytes(),
        storage_location_key(source).as_bytes(),
        policy.max_bytes.to_string().as_bytes(),
    ])
}

fn actor_subject(actor: &Actor) -> String {
    match actor {
        Actor::User(id) => format!("user:{id}"),
        Actor::Agent(id) => format!("agent:{id}"),
        Actor::Node(id) => format!("node:{id}"),
        Actor::Task(id) => format!("task:{id}"),
    }
}

fn storage_location_key(source: &StorageLocation) -> String {
    match source {
        StorageLocation::RetainedNode(node) => format!("retained-node:{node}"),
        StorageLocation::ExplicitStore(location) => format!("explicit-store:{location}"),
    }
}

fn source_is_available(metadata: &ArtifactMetadata, source: &StorageLocation) -> bool {
    match source {
        StorageLocation::RetainedNode(node) => metadata.retaining_nodes.contains(node),
        StorageLocation::ExplicitStore(location) => metadata.explicit_locations.contains(location),
    }
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeMap;

    use crate::{Actor, LimitKind, ResourceLimits, ResourceMeter, UserId};

    use super::*;

    fn registry_with_artifact() -> ArtifactRegistry {
        let mut registry = ArtifactRegistry::default();
        registry.flush_metadata(
            ArtifactId::from("artifact"),
            TenantId::from("tenant"),
            ProjectId::from("project"),
            ProcessId::from("process"),
            TaskId::from("task"),
            NodeId::from("node"),
            Digest::sha256("bytes"),
            32,
        );
        registry
    }

    #[test]
    fn flush_publishes_metadata_without_coordinator_bytes() {
        let registry = registry_with_artifact();
        let metadata = registry.metadata(&ArtifactId::from("artifact")).unwrap();

        assert!(!metadata.coordinator_has_large_bytes);
        assert_eq!(metadata.id, ArtifactId::from("artifact"));
        assert_eq!(metadata.tenant, TenantId::from("tenant"));
        assert_eq!(metadata.project, ProjectId::from("project"));
        assert_eq!(metadata.process, ProcessId::from("process"));
        assert_eq!(metadata.producer_task, TaskId::from("task"));
        assert_eq!(metadata.producer_node, NodeId::from("node"));
        assert_eq!(metadata.digest, Digest::sha256("bytes"));
        assert_eq!(metadata.size, 32);
        assert_eq!(metadata.flushed_epoch, 1);
        assert!(metadata.retaining_nodes.contains(&NodeId::from("node")));
        assert!(metadata.explicit_locations.is_empty());
    }

    #[test]
    fn unsynced_node_loss_surfaces_as_unavailable() {
        let mut registry = registry_with_artifact();
        registry.garbage_collect_node(&NodeId::from("node"));
        let context = AuthContext {
            tenant: TenantId::from("tenant"),
            project: ProjectId::from("project"),
            actor: Actor::User(UserId::from("user")),
        };

        let error = registry
            .download_action(
                &context,
                &ArtifactId::from("artifact"),
                &DownloadPolicy { max_bytes: 100 },
            )
            .unwrap_err();

        assert_eq!(error, DownloadError::Unavailable);
    }

    #[test]
    fn explicit_user_storage_location_survives_node_retention_loss() {
        let mut registry = registry_with_artifact();
        registry
            .sync_to_explicit_store(&ArtifactId::from("artifact"), "s3://bucket/app")
            .unwrap();
        registry.garbage_collect_node(&NodeId::from("node"));
        let context = AuthContext {
            tenant: TenantId::from("tenant"),
            project: ProjectId::from("project"),
            actor: Actor::User(UserId::from("user")),
        };

        let action = registry
            .download_action(
                &context,
                &ArtifactId::from("artifact"),
                &DownloadPolicy { max_bytes: 100 },
            )
            .unwrap();

        assert_eq!(
            action.source,
            StorageLocation::ExplicitStore("s3://bucket/app".to_owned())
        );
        assert!(
            !registry
                .metadata(&ArtifactId::from("artifact"))
                .unwrap()
                .coordinator_has_large_bytes
        );
    }

    #[test]
    fn default_retention_policy_is_best_effort_node_retention() {
        let policy = RetentionPolicy::default();

        assert!(policy.best_effort_node_retention);
        assert_eq!(policy.max_download_bytes, 256 * 1024 * 1024);
    }

    #[test]
    fn cross_tenant_download_is_denied_even_with_known_artifact_id() {
        let registry = registry_with_artifact();
        let context = AuthContext {
            tenant: TenantId::from("other"),
            project: ProjectId::from("project"),
            actor: Actor::User(UserId::from("user")),
        };

        let error = registry
            .download_action(
                &context,
                &ArtifactId::from("artifact"),
                &DownloadPolicy { max_bytes: 100 },
            )
            .unwrap_err();

        assert!(matches!(error, DownloadError::Unauthorized(_)));
    }

    #[test]
    fn cross_project_download_is_denied_even_with_known_artifact_id() {
        let registry = registry_with_artifact();
        let context = AuthContext {
            tenant: TenantId::from("tenant"),
            project: ProjectId::from("other-project"),
            actor: Actor::User(UserId::from("user")),
        };

        let error = registry
            .download_action(
                &context,
                &ArtifactId::from("artifact"),
                &DownloadPolicy { max_bytes: 100 },
            )
            .unwrap_err();

        assert!(matches!(error, DownloadError::Unauthorized(_)));
    }

    #[test]
    fn download_link_is_not_created_when_artifact_is_unavailable_or_too_large() {
        let mut registry = registry_with_artifact();
        registry.garbage_collect_node(&NodeId::from("node"));
        let context = AuthContext {
            tenant: TenantId::from("tenant"),
            project: ProjectId::from("project"),
            actor: Actor::User(UserId::from("user")),
        };

        assert_eq!(
            registry
                .create_download_link(
                    &context,
                    &ArtifactId::from("artifact"),
                    &DownloadPolicy { max_bytes: 100 },
                    "nonce",
                    10,
                    60,
                )
                .unwrap_err(),
            DownloadError::Unavailable
        );

        let mut registry = registry_with_artifact();
        assert!(matches!(
            registry.create_download_link(
                &context,
                &ArtifactId::from("artifact"),
                &DownloadPolicy { max_bytes: 1 },
                "nonce",
                10,
                60,
            ),
            Err(DownloadError::LimitExceeded { .. })
        ));
    }

    #[test]
    fn download_link_is_authenticated_scoped_and_not_guessable() {
        let mut registry = registry_with_artifact();
        let context = AuthContext {
            tenant: TenantId::from("tenant"),
            project: ProjectId::from("project"),
            actor: Actor::User(UserId::from("user")),
        };
        let link = registry
            .create_download_link(
                &context,
                &ArtifactId::from("artifact"),
                &DownloadPolicy { max_bytes: 100 },
                "nonce-a",
                10,
                60,
            )
            .unwrap();
        let other = registry
            .create_download_link(
                &context,
                &ArtifactId::from("artifact"),
                &DownloadPolicy { max_bytes: 100 },
                "nonce-b",
                10,
                60,
            )
            .unwrap();

        assert_eq!(link.tenant, TenantId::from("tenant"));
        assert_eq!(link.project, ProjectId::from("project"));
        assert_eq!(link.process, ProcessId::from("process"));
        assert_eq!(link.actor, Actor::User(UserId::from("user")));
        assert_eq!(link.max_bytes, 100);
        assert!(link.policy_context_digest.is_valid_sha256());
        assert_eq!(link.expires_at_epoch_seconds, 70);
        assert!(link
            .url_path
            .contains("/artifacts/tenant/project/process/artifact"));
        assert_ne!(link.scoped_token_digest, other.scoped_token_digest);
        assert!(matches!(link.source, StorageLocation::RetainedNode(_)));
    }

    #[test]
    fn download_link_is_bound_to_actor_and_policy_context() {
        let mut registry = registry_with_artifact();
        let context = AuthContext {
            tenant: TenantId::from("tenant"),
            project: ProjectId::from("project"),
            actor: Actor::User(UserId::from("user")),
        };
        let other_actor = AuthContext {
            tenant: TenantId::from("tenant"),
            project: ProjectId::from("project"),
            actor: Actor::User(UserId::from("other-user")),
        };
        let policy = DownloadPolicy { max_bytes: 100 };
        let link = registry
            .create_download_link(
                &context,
                &ArtifactId::from("artifact"),
                &policy,
                "nonce",
                10,
                60,
            )
            .unwrap();
        let limits = ResourceLimits {
            limits: BTreeMap::from([(LimitKind::ArtifactDownloadBytes, 32)]),
        };
        let mut meter = ResourceMeter::default();

        assert_eq!(
            registry
                .open_download_stream(
                    &other_actor,
                    &ArtifactId::from("artifact"),
                    &policy,
                    &link.scoped_token_digest,
                    11,
                    &limits,
                    &mut meter,
                )
                .unwrap_err(),
            DownloadError::InvalidToken
        );
        assert_eq!(
            registry
                .open_download_stream(
                    &context,
                    &ArtifactId::from("artifact"),
                    &DownloadPolicy { max_bytes: 99 },
                    &link.scoped_token_digest,
                    11,
                    &limits,
                    &mut meter,
                )
                .unwrap_err(),
            DownloadError::InvalidToken
        );
        assert_eq!(
            registry
                .revoke_download_link(
                    &other_actor,
                    &ArtifactId::from("artifact"),
                    &link.scoped_token_digest,
                )
                .unwrap_err(),
            DownloadError::InvalidToken
        );
    }

    #[test]
    fn download_link_expires_and_can_be_revoked() {
        let mut registry = registry_with_artifact();
        let context = AuthContext {
            tenant: TenantId::from("tenant"),
            project: ProjectId::from("project"),
            actor: Actor::User(UserId::from("user")),
        };
        let policy = DownloadPolicy { max_bytes: 100 };
        let limits = ResourceLimits {
            limits: BTreeMap::from([(LimitKind::ArtifactDownloadBytes, 32)]),
        };
        let mut meter = ResourceMeter::default();
        let expired = registry
            .create_download_link(
                &context,
                &ArtifactId::from("artifact"),
                &policy,
                "expired",
                10,
                5,
            )
            .unwrap();

        let error = registry
            .open_download_stream(
                &context,
                &ArtifactId::from("artifact"),
                &policy,
                &expired.scoped_token_digest,
                16,
                &limits,
                &mut meter,
            )
            .unwrap_err();
        assert_eq!(error, DownloadError::Expired);

        let active = registry
            .create_download_link(
                &context,
                &ArtifactId::from("artifact"),
                &policy,
                "active",
                20,
                60,
            )
            .unwrap();
        let revoked = registry
            .revoke_download_link(
                &context,
                &ArtifactId::from("artifact"),
                &active.scoped_token_digest,
            )
            .unwrap();
        assert_eq!(revoked.scoped_token_digest, active.scoped_token_digest);

        let error = registry
            .open_download_stream(
                &context,
                &ArtifactId::from("artifact"),
                &policy,
                &active.scoped_token_digest,
                21,
                &limits,
                &mut meter,
            )
            .unwrap_err();
        assert_eq!(error, DownloadError::Revoked);
    }

    #[test]
    fn download_stream_accounts_usage_before_and_during_streaming() {
        let mut registry = registry_with_artifact();
        let context = AuthContext {
            tenant: TenantId::from("tenant"),
            project: ProjectId::from("project"),
            actor: Actor::User(UserId::from("user")),
        };
        let policy = DownloadPolicy { max_bytes: 100 };
        let link = registry
            .create_download_link(
                &context,
                &ArtifactId::from("artifact"),
                &policy,
                "nonce",
                10,
                60,
            )
            .unwrap();
        let limits = ResourceLimits {
            limits: BTreeMap::from([(LimitKind::ArtifactDownloadBytes, 32)]),
        };
        let mut meter = ResourceMeter::default();
        let mut stream = registry
            .open_download_stream(
                &context,
                &ArtifactId::from("artifact"),
                &policy,
                &link.scoped_token_digest,
                11,
                &limits,
                &mut meter,
            )
            .unwrap();

        stream.stream_chunk(&limits, &mut meter, 16).unwrap();
        stream.stream_chunk(&limits, &mut meter, 16).unwrap();
        assert!(matches!(
            stream.stream_chunk(&limits, &mut meter, 1),
            Err(DownloadError::Usage(_))
        ));
        assert_eq!(meter.used(&LimitKind::ArtifactDownloadBytes), 32);
    }

    #[test]
    fn download_stream_fails_honestly_when_source_disappears_mid_stream() {
        let mut registry = registry_with_artifact();
        let context = AuthContext {
            tenant: TenantId::from("tenant"),
            project: ProjectId::from("project"),
            actor: Actor::User(UserId::from("user")),
        };
        let policy = DownloadPolicy { max_bytes: 100 };
        let link = registry
            .create_download_link(
                &context,
                &ArtifactId::from("artifact"),
                &policy,
                "nonce",
                10,
                60,
            )
            .unwrap();
        let limits = ResourceLimits {
            limits: BTreeMap::from([(LimitKind::ArtifactDownloadBytes, 32)]),
        };
        let mut meter = ResourceMeter::default();
        let mut stream = registry
            .open_download_stream(
                &context,
                &ArtifactId::from("artifact"),
                &policy,
                &link.scoped_token_digest,
                11,
                &limits,
                &mut meter,
            )
            .unwrap();

        registry
            .stream_download_chunk(&mut stream, &limits, &mut meter, 16)
            .unwrap();
        registry.garbage_collect_node(&NodeId::from("node"));

        assert_eq!(
            registry
                .stream_download_chunk(&mut stream, &limits, &mut meter, 1)
                .unwrap_err(),
            DownloadError::Unavailable
        );
        assert_eq!(stream.streamed_bytes, 16);
    }

    #[test]
    fn guessed_download_token_is_rejected() {
        let registry = registry_with_artifact();
        let context = AuthContext {
            tenant: TenantId::from("tenant"),
            project: ProjectId::from("project"),
            actor: Actor::User(UserId::from("user")),
        };
        let limits = ResourceLimits {
            limits: BTreeMap::from([(LimitKind::ArtifactDownloadBytes, 32)]),
        };
        let mut meter = ResourceMeter::default();

        let error = registry
            .open_download_stream(
                &context,
                &ArtifactId::from("artifact"),
                &DownloadPolicy { max_bytes: 100 },
                &Digest::sha256("guessed"),
                11,
                &limits,
                &mut meter,
            )
            .unwrap_err();

        assert_eq!(error, DownloadError::InvalidToken);
    }
}