1 """

2 Asset management system for pyPhotoAlbum with automatic reference counting

3 """

4

5 import hashlib

6 import os

7 import shutil

8 from typing import Dict, List, Optional, Tuple

9 from pathlib import Path

10

11

12 def compute_file_md5(file_path: str) -> Optional[str]:

13 """

14 Compute MD5 hash of a file.

15

16 Args:

17 file_path: Path to the file

18

19 Returns:

20 MD5 hash as hex string, or None if file doesn't exist

21 """

22 if not os.path.exists(file_path):

23 return None

24

25 hash_md5 = hashlib.md5()

26 try:

27 with open(file_path, "rb") as f:

28 for chunk in iter(lambda: f.read(8192), b""):

29 hash_md5.update(chunk)

30 return hash_md5.hexdigest()

31 except Exception as e:

32 print(f"AssetManager: Error computing MD5 for {file_path}: {e}")

33 return None

34

35

36 class AssetManager:

37 """Manages project assets with automatic reference counting and cleanup"""

38

39 def __init__(self, project_folder: str):

40 """

41 Initialize AssetManager.

42

43 Args:

44 project_folder: Root folder for the project

45 """

46 self.project_folder = project_folder

47 self.assets_folder = os.path.join(project_folder, "assets")

48 self.reference_counts: Dict[str, int] = {} # {relative_path: count}

49 self.asset_hashes: Dict[str, str] = {} # {relative_path: md5_hash}

50

51 # Create assets folder if it doesn't exist

52 os.makedirs(self.assets_folder, exist_ok=True)

53

54 def import_asset(self, source_path: str) -> str:

55 """

56 Import an asset into the project by copying it to the assets folder.

57

58 Args:

59 source_path: Path to the source file

60

61 Returns:

62 Relative path to the imported asset (e.g., "assets/photo_001.jpg")

63 """

64 if not os.path.exists(source_path):

65 raise FileNotFoundError(f"Source file not found: {source_path}")

66

67 # Get filename and extension

68 filename = os.path.basename(source_path)

69 name, ext = os.path.splitext(filename)

70

71 # Find a unique filename if there's a collision

72 counter = 1

73 dest_filename = filename

74 dest_path = os.path.join(self.assets_folder, dest_filename)

75

76 while os.path.exists(dest_path):

77 dest_filename = f"{name}_{counter:03d}{ext}"

78 dest_path = os.path.join(self.assets_folder, dest_filename)

79 counter += 1

80

81 # Copy the file

82 shutil.copy2(source_path, dest_path)

83

84 # Get relative path from project folder (for storage/serialization)

85 relative_path = os.path.relpath(dest_path, self.project_folder)

86

87 # Initialize reference count

88 self.reference_counts[relative_path] = 1

89

90 print(f"AssetManager: Imported {source_path} → {dest_path} (stored as {relative_path}, refs=1)")

91

92 # Return relative path for storage in elements

93 return relative_path

94

95 def acquire_reference(self, asset_path: str):

96 """

97 Increment the reference count for an asset.

98

99 Args:

100 asset_path: Relative path to the asset

101 """

102 if not asset_path:

103 return

104

105 if asset_path in self.reference_counts:

106 self.reference_counts[asset_path] += 1

107 print(f"AssetManager: Acquired reference to {asset_path} (refs={self.reference_counts[asset_path]})")

108 else:

109 # Asset might exist from a loaded project

110 full_path = os.path.join(self.project_folder, asset_path)

111 if os.path.exists(full_path):

112 self.reference_counts[asset_path] = 1

113 print(f"AssetManager: Acquired reference to existing asset {asset_path} (refs=1)")

114 else:

115 print(f"AssetManager: Warning - asset not found: {asset_path}")

116

117 def release_reference(self, asset_path: str):

118 """

119 Decrement the reference count for an asset.

120 If count reaches zero, delete the asset file.

121

122 Args:

123 asset_path: Relative path to the asset

124 """

125 if not asset_path:

126 return

127

128 if asset_path not in self.reference_counts:

129 print(f"AssetManager: Warning - attempting to release unknown asset: {asset_path}")

130 return

131

132 self.reference_counts[asset_path] -= 1

133 print(f"AssetManager: Released reference to {asset_path} (refs={self.reference_counts[asset_path]})")

134

135 if self.reference_counts[asset_path] <= 0:

136 # No more references - safe to delete

137 full_path = os.path.join(self.project_folder, asset_path)

138 try:

139 if os.path.exists(full_path):

140 os.remove(full_path)

141 print(f"AssetManager: Deleted unused asset {asset_path}")

142 del self.reference_counts[asset_path]

143 except Exception as e:

144 print(f"AssetManager: Error deleting asset {asset_path}: {e}")

145

146 def get_absolute_path(self, relative_path: str) -> str:

147 """

148 Convert a relative asset path to an absolute path.

149

150 Args:

151 relative_path: Relative path from project folder

152

153 Returns:

154 Absolute path to the asset

155 """

156 return os.path.join(self.project_folder, relative_path)

157

158 def get_reference_count(self, asset_path: str) -> int:

159 """

160 Get the current reference count for an asset.

161

162 Args:

163 asset_path: Relative path to the asset

164

165 Returns:

166 Reference count (0 if not tracked)

167 """

168 return self.reference_counts.get(asset_path, 0)

169

170 def serialize(self) -> Dict:

171 """Serialize asset manager state"""

172 return {

173 "reference_counts": self.reference_counts,

174 "asset_hashes": self.asset_hashes,

175 }

176

177 def deserialize(self, data: Dict):

178 """Deserialize asset manager state"""

179 self.reference_counts = data.get("reference_counts", {})

180 self.asset_hashes = data.get("asset_hashes", {})

181 print(f"AssetManager: Loaded {len(self.reference_counts)} asset references")

182

183 def compute_asset_hash(self, asset_path: str) -> Optional[str]:

184 """

185 Compute and cache the MD5 hash for an asset.

186

187 Args:

188 asset_path: Relative path to the asset

189

190 Returns:

191 MD5 hash as hex string, or None if computation fails

192 """

193 full_path = self.get_absolute_path(asset_path)

194 md5_hash = compute_file_md5(full_path)

195 if md5_hash:

196 self.asset_hashes[asset_path] = md5_hash

197 return md5_hash

198

199 def compute_all_hashes(self) -> Dict[str, str]:

200 """

201 Compute MD5 hashes for all assets in the assets folder.

202

203 Returns:

204 Dictionary mapping relative paths to MD5 hashes

205 """

206 self.asset_hashes.clear()

207

208 if not os.path.exists(self.assets_folder):

209 return self.asset_hashes

210

211 for root, dirs, files in os.walk(self.assets_folder):

212 for filename in files:

213 file_path = os.path.join(root, filename)

214 relative_path = os.path.relpath(file_path, self.project_folder)

215 md5_hash = compute_file_md5(file_path)

216 if md5_hash:

217 self.asset_hashes[relative_path] = md5_hash

218

219 print(f"AssetManager: Computed hashes for {len(self.asset_hashes)} assets")

220 return self.asset_hashes

221

222 def find_duplicates(self) -> Dict[str, List[str]]:

223 """

224 Find duplicate assets based on MD5 hash.

225

226 Returns:

227 Dictionary mapping MD5 hash to list of asset paths with that hash.

228 Only includes hashes with more than one file.

229 """

230 # Compute hashes if not already done

231 if not self.asset_hashes:

232 self.compute_all_hashes()

233

234 # Group assets by hash

235 hash_to_paths: Dict[str, List[str]] = {}

236 for path, md5_hash in self.asset_hashes.items():

237 if md5_hash not in hash_to_paths:

238 hash_to_paths[md5_hash] = []

239 hash_to_paths[md5_hash].append(path)

240

241 # Filter to only duplicates (more than one file with same hash)

242 duplicates = {h: paths for h, paths in hash_to_paths.items() if len(paths) > 1}

243

244 if duplicates:

245 total_dups = sum(len(paths) - 1 for paths in duplicates.values())

246 print(f"AssetManager: Found {total_dups} duplicate files in {len(duplicates)} groups")

247

248 return duplicates

249

250 def deduplicate_assets(self, update_references_callback=None) -> Tuple[int, int]:

251 """

252 Remove duplicate assets, keeping one canonical copy and updating references.

253

254 Args:

255 update_references_callback: Optional callback function that takes

256 (old_path, new_path) to update external references (e.g., ImageData elements)

257

258 Returns:

259 Tuple of (files_removed, bytes_saved)

260 """

261 duplicates = self.find_duplicates()

262 if not duplicates:

263 print("AssetManager: No duplicates found")

264 return (0, 0)

265

266 files_removed = 0

267 bytes_saved = 0

268

269 for md5_hash, paths in duplicates.items():

270 # Sort paths to get consistent canonical path (first alphabetically)

271 paths.sort()

272 canonical_path = paths[0]

273

274 # Remove duplicates and update references

275 for dup_path in paths[1:]:

276 full_dup_path = self.get_absolute_path(dup_path)

277

278 # Get file size before deletion

279 try:

280 file_size = os.path.getsize(full_dup_path)

281 except OSError:

282 file_size = 0

283

284 # Update references if callback provided

285 if update_references_callback:

286 update_references_callback(dup_path, canonical_path)

287

288 # Transfer reference count to canonical path

289 if dup_path in self.reference_counts:

290 dup_refs = self.reference_counts[dup_path]

291 if canonical_path in self.reference_counts:

292 self.reference_counts[canonical_path] += dup_refs

293 else:

294 self.reference_counts[canonical_path] = dup_refs

295 del self.reference_counts[dup_path]

296

297 # Delete the duplicate file

298 try:

299 if os.path.exists(full_dup_path):

300 os.remove(full_dup_path)

301 files_removed += 1

302 bytes_saved += file_size

303 print(f"AssetManager: Removed duplicate {dup_path} (kept {canonical_path})")

304 except Exception as e:

305 print(f"AssetManager: Error removing duplicate {dup_path}: {e}")

306

307 # Remove from hash tracking

308 if dup_path in self.asset_hashes:

309 del self.asset_hashes[dup_path]

310

311 print(f"AssetManager: Deduplication complete - removed {files_removed} files, saved {bytes_saved} bytes")

312 return (files_removed, bytes_saved)

313

314 def get_duplicate_stats(self) -> Tuple[int, int, int]:

315 """

316 Get statistics about duplicate assets without modifying anything.

317

318 Returns:

319 Tuple of (duplicate_groups, total_duplicate_files, estimated_bytes_to_save)

320 """

321 duplicates = self.find_duplicates()

322 if not duplicates:

323 return (0, 0, 0)

324

325 duplicate_groups = len(duplicates)

326 total_duplicate_files = sum(len(paths) - 1 for paths in duplicates.values())

327

328 # Calculate bytes that would be saved

329 bytes_to_save = 0

330 for paths in duplicates.values():

331 # Skip the first (canonical) file, count size of the rest

332 for dup_path in paths[1:]:

333 full_path = self.get_absolute_path(dup_path)

334 try:

335 bytes_to_save += os.path.getsize(full_path)

336 except OSError:

337 pass

338

339 return (duplicate_groups, total_duplicate_files, bytes_to_save)

340

341 def find_unused_assets(self) -> List[str]:

342 """

343 Find assets that exist in the assets folder but have no references.

344

345 Returns:

346 List of relative paths to unused assets

347 """

348 unused: list[str] = []

349

350 if not os.path.exists(self.assets_folder):

351 return unused

352

353 # Get all files in assets folder

354 for root, dirs, files in os.walk(self.assets_folder):

355 for filename in files:

356 file_path = os.path.join(root, filename)

357 relative_path = os.path.relpath(file_path, self.project_folder)

358

359 # Check if this asset has any references

360 ref_count = self.reference_counts.get(relative_path, 0)

361 if ref_count <= 0:

362 unused.append(relative_path)

363

364 if unused:

365 print(f"AssetManager: Found {len(unused)} unused assets")

366

367 return unused

368

369 def get_unused_stats(self) -> Tuple[int, int]:

370 """

371 Get statistics about unused assets without modifying anything.

372

373 Returns:

374 Tuple of (unused_file_count, total_bytes)

375 """

376 unused = self.find_unused_assets()

377 if not unused:

378 return (0, 0)

379

380 total_bytes = 0

381 for asset_path in unused:

382 full_path = self.get_absolute_path(asset_path)

383 try:

384 total_bytes += os.path.getsize(full_path)

385 except OSError:

386 pass

387

388 return (len(unused), total_bytes)

389

390 def remove_unused_assets(self) -> Tuple[int, int]:

391 """

392 Remove all unused assets from the assets folder.

393

394 Returns:

395 Tuple of (files_removed, bytes_freed)

396 """

397 unused = self.find_unused_assets()

398 if not unused:

399 print("AssetManager: No unused assets to remove")

400 return (0, 0)

401

402 files_removed = 0

403 bytes_freed = 0

404

405 for asset_path in unused:

406 full_path = self.get_absolute_path(asset_path)

407

408 try:

409 file_size = os.path.getsize(full_path)

410 except OSError:

411 file_size = 0

412

413 try:

414 if os.path.exists(full_path):

415 os.remove(full_path)

416 files_removed += 1

417 bytes_freed += file_size

418 print(f"AssetManager: Removed unused asset {asset_path}")

419

420 # Clean up tracking

421 if asset_path in self.reference_counts:

422 del self.reference_counts[asset_path]

423 if asset_path in self.asset_hashes:

424 del self.asset_hashes[asset_path]

425

426 except Exception as e:

427 print(f"AssetManager: Error removing unused asset {asset_path}: {e}")

428

429 print(f"AssetManager: Removed {files_removed} unused assets, freed {bytes_freed} bytes")

430 return (files_removed, bytes_freed)