1use anyhow::{Context, Result};
6use serde::{Deserialize, Serialize};
7use sha2::{Digest, Sha256};
8use std::collections::BTreeMap;
9use std::io::Read;
10use std::path::{Path, PathBuf};
11
12use crate::Artifact;
13
14pub const HUB_DIR_ENV: &str = "DSPK_HUB_DIR";
16
17const HASH_CHUNK_BYTES: usize = 64 * 1024;
19
20#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
22pub enum ArtifactStatus {
23 Installed,
25 Available,
27 Corrupted,
29}
30
31impl ArtifactStatus {
32 pub fn badge(self) -> &'static str {
34 match self {
35 Self::Installed => "[INSTALLED]",
36 Self::Available => "[AVAILABLE]",
37 Self::Corrupted => "[CORRUPTED]",
38 }
39 }
40}
41
42impl std::fmt::Display for ArtifactStatus {
43 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
44 f.write_str(self.badge())
45 }
46}
47
48#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
50pub struct CachedArtifactRecord {
51 pub id: String,
52 pub family: String,
53 pub sha256: String,
54 pub size_bytes: u64,
55 pub relative_path: String,
56 pub source_url: String,
57}
58
59#[derive(Debug, Clone, Default, Serialize, Deserialize)]
61pub struct HubCacheIndex {
62 pub entries: BTreeMap<String, CachedArtifactRecord>,
63}
64
65#[derive(Debug, Clone)]
67pub struct HubCache {
68 root_dir: PathBuf,
69}
70
71impl HubCache {
72 pub fn from_env_or_default() -> Self {
74 if let Ok(dir) = std::env::var(HUB_DIR_ENV)
75 && !dir.trim().is_empty()
76 {
77 return Self::new(PathBuf::from(dir.trim()));
78 }
79 let base = dirs::cache_dir()
80 .or_else(|| dirs::home_dir().map(|h| h.join(".cache")))
81 .unwrap_or_else(|| PathBuf::from(".cache"));
82 Self::new(base.join("dspk").join("hub"))
83 }
84
85 pub fn new(root_dir: impl Into<PathBuf>) -> Self {
86 Self { root_dir: root_dir.into() }
87 }
88
89 pub fn root_dir(&self) -> &Path {
90 &self.root_dir
91 }
92
93 pub fn index_path(&self) -> PathBuf {
94 self.root_dir.join("index.json")
95 }
96
97 pub fn artifact_dir(&self, artifact: &Artifact) -> PathBuf {
99 self.root_dir.join("models").join(&artifact.family).join(artifact.slug())
100 }
101
102 pub fn path(&self, artifact: &Artifact) -> PathBuf {
104 self.artifact_dir(artifact).join(&artifact.file_name)
105 }
106
107 pub fn load_index(&self) -> Result<HubCacheIndex> {
109 let path = self.index_path();
110 if !path.exists() {
111 return Ok(HubCacheIndex::default());
112 }
113 let content = std::fs::read_to_string(&path).with_context(|| format!("failed to read hub cache index {}", path.display()))?;
114 serde_json::from_str(&content).with_context(|| format!("failed to parse hub cache index {}", path.display()))
115 }
116
117 pub fn save_index(&self, index: &HubCacheIndex) -> Result<()> {
118 std::fs::create_dir_all(&self.root_dir).with_context(|| format!("failed to create hub cache directory {}", self.root_dir.display()))?;
119 let path = self.index_path();
120 std::fs::write(&path, serde_json::to_string_pretty(index)?).with_context(|| format!("failed to write hub cache index {}", path.display()))
121 }
122
123 pub fn record_installed(&self, artifact: &Artifact, sha256: &str, size_bytes: u64) -> Result<()> {
125 let path = self.path(artifact);
126 let relative_path = path.strip_prefix(&self.root_dir).unwrap_or(&path).to_string_lossy().into_owned();
127 let mut index = self.load_index().unwrap_or_default();
128 index.entries.insert(
129 artifact.id.clone(),
130 CachedArtifactRecord {
131 id: artifact.id.clone(),
132 family: artifact.family.clone(),
133 sha256: sha256.to_ascii_lowercase(),
134 size_bytes,
135 relative_path,
136 source_url: artifact.url.clone(),
137 },
138 );
139 self.save_index(&index)
140 }
141
142 pub fn status(&self, artifact: &Artifact) -> ArtifactStatus {
144 let path = self.path(artifact);
145 let Ok(meta) = path.metadata() else { return ArtifactStatus::Available };
146 let recorded = self.load_index().ok().and_then(|i| i.entries.get(&artifact.id).cloned());
147 match recorded {
148 Some(r) if r.sha256.eq_ignore_ascii_case(&artifact.sha256) && r.size_bytes == artifact.size_bytes && meta.len() == artifact.size_bytes => {
149 ArtifactStatus::Installed
150 }
151 _ => ArtifactStatus::Corrupted,
152 }
153 }
154
155 pub fn verify(&self, artifact: &Artifact) -> Result<(ArtifactStatus, Option<String>)> {
157 let path = self.path(artifact);
158 if !path.exists() {
159 return Ok((ArtifactStatus::Available, None));
160 }
161 let (digest, size) = compute_file_sha256(&path)?;
162 let ok = digest.eq_ignore_ascii_case(&artifact.sha256) && size == artifact.size_bytes;
163 Ok((if ok { ArtifactStatus::Installed } else { ArtifactStatus::Corrupted }, Some(digest)))
164 }
165
166 pub fn remove(&self, artifact: &Artifact) -> Result<bool> {
168 let mut removed = false;
169 let dir = self.artifact_dir(artifact);
170 if dir.exists() {
171 std::fs::remove_dir_all(&dir).with_context(|| format!("failed to remove cached artifact directory {}", dir.display()))?;
172 removed = true;
173 }
174 if let Ok(mut index) = self.load_index()
175 && index.entries.remove(&artifact.id).is_some()
176 {
177 self.save_index(&index)?;
178 removed = true;
179 }
180 Ok(removed)
181 }
182
183 pub fn clean_all(&self) -> Result<u64> {
185 let models_dir = self.root_dir.join("models");
186 let freed = dir_size_bytes(&models_dir);
187 if models_dir.exists() {
188 std::fs::remove_dir_all(&models_dir).with_context(|| format!("failed to clean hub models directory {}", models_dir.display()))?;
189 }
190 let _ = std::fs::remove_file(self.index_path());
191 Ok(freed)
192 }
193}
194
195pub fn compute_file_sha256(path: &Path) -> Result<(String, u64)> {
197 let mut file = std::fs::File::open(path).with_context(|| format!("failed to open file for SHA-256 check: {}", path.display()))?;
198 let mut hasher = Sha256::new();
199 let mut buf = vec![0u8; HASH_CHUNK_BYTES];
200 let mut total = 0u64;
201 loop {
202 let n = file.read(&mut buf)?;
203 if n == 0 {
204 break;
205 }
206 hasher.update(&buf[..n]);
207 total += n as u64;
208 }
209 Ok((hex::encode(hasher.finalize()), total))
210}
211
212fn dir_size_bytes(path: &Path) -> u64 {
213 let Ok(entries) = std::fs::read_dir(path) else { return 0 };
214 entries
215 .flatten()
216 .map(|e| {
217 let p = e.path();
218 if p.is_dir() { dir_size_bytes(&p) } else { e.metadata().map_or(0, |m| m.len()) }
219 })
220 .sum()
221}