khora_agents/asset_agent/
agent.rs

1// Copyright 2025 eraflo
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! The AssetAgent is responsible for managing asset loading and retrieval.
16
17use anyhow::{anyhow, Context, Result};
18use khora_core::{
19    asset::{Asset, AssetHandle, AssetUUID},
20    vfs::VirtualFileSystem,
21};
22use khora_data::assets::Assets;
23use khora_lanes::asset_lane::{AssetLoaderLane, PackLoadingLane};
24use khora_telemetry::MetricsRegistry;
25use std::{
26    any::{Any, TypeId},
27    collections::HashMap,
28    fs::File,
29    sync::Arc,
30};
31
32use crate::asset_agent::loader::AssetLoaderLaneRegistry;
33
34/// The AssetAgent is responsible for managing asset loading and retrieval.
35pub struct AssetAgent {
36    /// The virtual file system for asset management.
37    vfs: VirtualFileSystem,
38    /// The loading lane for asset I/O operations.
39    loading_lane: PackLoadingLane,
40    /// The registry that manages asset loaders.
41    loaders: AssetLoaderLaneRegistry,
42    /// Type-erased storage for different types of loaded assets.
43    /// Maps a TypeId to a Box<Any> that holds an `Assets<A>`.
44    storages: HashMap<TypeId, Box<dyn Any + Send + Sync>>,
45}
46
47impl AssetAgent {
48    /// Creates a new `AssetAgent` with the given VFS and loading lane.
49    pub fn new(
50        index_bytes: &[u8],
51        data_file: File,
52        metrics_registry: Arc<MetricsRegistry>,
53    ) -> Result<Self> {
54        let vfs = VirtualFileSystem::new(index_bytes)
55            .context("Failed to initialize VirtualFileSystem from index bytes")?;
56
57        let loading_lane = PackLoadingLane::new(data_file);
58
59        Ok(Self {
60            vfs,
61            loading_lane,
62            loaders: AssetLoaderLaneRegistry::new(metrics_registry),
63            storages: HashMap::new(),
64        })
65    }
66
67    /// Registers an `AssetLoaderLane` for a specific asset type name.
68    ///
69    /// The `type_name` should match the one generated by the `xtask` packager
70    /// based on the file extension.
71    pub fn register_loader<A: Asset>(
72        &mut self,
73        type_name: &str,
74        loader: impl AssetLoaderLane<A> + Send + Sync + 'static,
75    ) {
76        self.loaders.register::<A>(type_name, loader);
77    }
78
79    /// Loads, decodes, and returns a typed handle to an asset.
80    ///
81    /// This is the main, fully-featured loading method.
82    pub fn load<A: Asset>(&mut self, uuid: &AssetUUID) -> Result<AssetHandle<A>> {
83        let type_id = TypeId::of::<A>();
84
85        // --- 1. Get or create the specific storage for the asset type `A` ---
86        let storage = self
87            .storages
88            .entry(type_id)
89            .or_insert_with(|| Box::new(Assets::<A>::new()));
90
91        // Downcast the `Box<dyn Any>` to `&mut Assets<A>`
92        let assets = storage
93            .downcast_mut::<Assets<A>>()
94            .ok_or_else(|| anyhow!("Mismatched asset storage type"))?;
95
96        // --- 2. Check if the asset is already cached ---
97        if let Some(handle) = assets.get(uuid) {
98            return Ok(handle.clone()); // Cache hit!
99        }
100
101        // --- 3. If not cached, load the asset (Cache miss) ---
102        let metadata = self
103            .vfs
104            .get_metadata(uuid)
105            .ok_or_else(|| anyhow!("Asset with UUID {:?} not found in VFS", uuid))?;
106
107        let source = metadata
108            .variants
109            .get("default")
110            .ok_or_else(|| anyhow!("Asset {:?} has no 'default' variant", uuid))?;
111
112        let bytes = self.loading_lane.load_asset_bytes(source)?;
113
114        let asset: A = self.loaders.load::<A>(&metadata.asset_type_name, &bytes)?;
115
116        // --- 4. Create the handle and store it in the cache before returning ---
117        let handle = AssetHandle::new(asset);
118        assets.insert(*uuid, handle.clone());
119
120        Ok(handle)
121    }
122}