//! [`AssetServer`]: the central registry that loads, deduplicates, and hands out
//! [`Handle`]s, plus the [`AssetLoader`] trait that makes it extensible.

use std::any::{Any, TypeId};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, RwLock, Weak};

use super::handle::{AssetCell, AssetId, Handle};

/// Errors produced while loading assets.
#[derive(Debug, thiserror::Error)]
pub enum AssetError {
    /// The path had no file extension to pick a loader by.
    #[error("path has no file extension: {0}")]
    NoExtension(PathBuf),

    /// No loader was registered for the file's extension.
    #[error("no loader registered for extension '.{0}'")]
    NoLoader(String),

    /// A loader exists for the extension, but it produces a different asset
    /// type than the one requested at the call site.
    #[error("loader for '.{ext}' produces a different asset type than requested")]
    TypeMismatch {
        /// The extension whose loader was selected.
        ext: String,
    },

    /// The loader itself failed (I/O, parse, etc.).
    #[error("failed to load {path}: {message}")]
    Load {
        /// The asset path.
        path: PathBuf,
        /// The loader's error message.
        message: String,
    },
}

/// A pluggable importer that turns a file into an asset of one concrete type.
///
/// Implement this for each asset format and register it with
/// [`AssetServer::register_loader`]. The server dispatches by file extension and
/// checks that the loader's [`Asset`](Self::Asset) type matches what the caller
/// asked to load.
pub trait AssetLoader: Send + Sync + 'static {
    /// The type this loader produces.
    type Asset: Send + Sync + 'static;

    /// The lower-or-mixed-case extensions (without the dot) this loader handles,
    /// e.g. `&["gltf", "glb"]`.
    fn extensions(&self) -> &'static [&'static str];

    /// Loads and parses the asset at `path`.
    fn load(&self, path: &Path) -> Result<Self::Asset, AssetError>;
}

/// Type-erased view of an [`AssetLoader`] so loaders of different output types
/// can share one registry.
trait ErasedLoader: Send + Sync {
    fn output_type(&self) -> TypeId;
    fn load(&self, path: &Path) -> Result<Box<dyn Any + Send + Sync>, AssetError>;
}

impl<L: AssetLoader> ErasedLoader for L {
    fn output_type(&self) -> TypeId {
        TypeId::of::<L::Asset>()
    }

    fn load(&self, path: &Path) -> Result<Box<dyn Any + Send + Sync>, AssetError> {
        Ok(Box::new(<L as AssetLoader>::load(self, path)?))
    }
}

type CacheKey = (TypeId, PathBuf);

/// One entry in the dedup cache. Carries a weak reference to the asset cell so
/// dropped assets are pruned, plus a function pointer that knows how to rerun
/// the loader for the cell's concrete type. Storing the reload-by-type as a
/// per-entry `fn` is what lets [`AssetServer::reload_path`] reload an asset
/// without knowing its `T` at the call site — the original `insert_cache::<T>`
/// captures `T` into the function pointer.
#[derive(Clone)]
struct CacheEntry {
    weak: Weak<dyn Any + Send + Sync>,
    reload_in_place: fn(&AssetServer, &Path),
}

struct Inner {
    loaders: RwLock<HashMap<String, Arc<dyn ErasedLoader>>>,
    /// Dedup cache: weak references, so a cached asset with no live handles is
    /// collected and reloaded fresh next time.
    cache: Mutex<HashMap<CacheKey, CacheEntry>>,
    next_id: AtomicU64,
}

/// The central asset registry.
///
/// Cloning an `AssetServer` is cheap (it shares one inner state via `Arc`) so it
/// can be handed to background load threads and stored across systems. Loading
/// the same path+type twice returns handles to **one** shared asset; when the
/// last handle is dropped the asset is freed.
///
/// ```no_run
/// use oxide_engine::asset::AssetServer;
/// use oxide_engine::asset::GltfModel;
///
/// let assets = AssetServer::new(); // glTF loader registered by default
/// let model = assets.load::<GltfModel>("assets/models/cube.gltf");
/// if let Some(model) = model.get() {
///     println!("{} meshes", model.meshes.len());
/// }
/// ```
#[derive(Clone)]
pub struct AssetServer {
    inner: Arc<Inner>,
}

impl AssetServer {
    /// A server with the engine's built-in loaders registered (currently glTF).
    pub fn new() -> Self {
        let server = Self::empty();
        super::register_default_loaders(&server);
        server
    }

    /// A server with **no** loaders registered. Use [`register_loader`] to add
    /// them; handy for tests or fully custom asset pipelines.
    ///
    /// [`register_loader`]: Self::register_loader
    pub fn empty() -> Self {
        Self {
            inner: Arc::new(Inner {
                loaders: RwLock::new(HashMap::new()),
                cache: Mutex::new(HashMap::new()),
                next_id: AtomicU64::new(1),
            }),
        }
    }

    /// Registers `loader`, mapping each of its extensions to it.
    pub fn register_loader<L: AssetLoader>(&self, loader: L) {
        let exts: Vec<String> = loader
            .extensions()
            .iter()
            .map(|e| e.to_lowercase())
            .collect();
        let erased: Arc<dyn ErasedLoader> = Arc::new(loader);
        let mut loaders = self.inner.loaders.write().unwrap();
        for ext in exts {
            loaders.insert(ext, erased.clone());
        }
    }

    /// Removes the loader registered for `extension` (without the dot). Returns
    /// whether one was present. Used when a module that added a loader is removed.
    pub fn unregister_loader(&self, extension: &str) -> bool {
        self.inner
            .loaders
            .write()
            .unwrap()
            .remove(&extension.to_lowercase())
            .is_some()
    }

    /// Loads the asset at `path` as type `T`, blocking until it is ready.
    ///
    /// Returns a handle to a cached asset if one of the same path+type is
    /// already live. On failure the returned handle is in the
    /// [`Failed`](super::LoadState::Failed) state (inspect [`Handle::error`]).
    pub fn load<T: Send + Sync + 'static>(&self, path: impl AsRef<Path>) -> Handle<T> {
        let path = path.as_ref().to_path_buf();
        let key = (TypeId::of::<T>(), path.clone());
        if let Some(handle) = self.cached::<T>(&key) {
            return handle;
        }
        match self.run_loader::<T>(&path) {
            Ok(value) => {
                let cell = AssetCell::new_loaded(self.next_id(), Some(path), value);
                self.insert_cache(key, &cell);
                Handle::from_cell(cell)
            }
            // Failures are not cached, so a later load retries from scratch.
            Err(err) => Handle::from_cell(AssetCell::new_failed(
                self.next_id(),
                Some(path),
                err.to_string(),
            )),
        }
    }

    /// Loads the asset at `path` as type `T` on a background thread, returning a
    /// handle immediately in the [`Loading`](super::LoadState::Loading) state.
    ///
    /// Poll [`Handle::state`]/[`Handle::get`], or block with [`Handle::wait`].
    pub fn load_async<T: Send + Sync + 'static>(&self, path: impl AsRef<Path>) -> Handle<T> {
        let path = path.as_ref().to_path_buf();
        let key = (TypeId::of::<T>(), path.clone());
        if let Some(handle) = self.cached::<T>(&key) {
            return handle;
        }
        // Insert the loading cell up front so concurrent requests dedup onto it.
        let cell = AssetCell::<T>::new_loading(self.next_id(), Some(path.clone()));
        self.insert_cache(key.clone(), &cell);

        let server = self.clone();
        let worker_cell = cell.clone();
        std::thread::spawn(move || match server.run_loader::<T>(&path) {
            Ok(value) => worker_cell.set_loaded(value),
            Err(err) => {
                worker_cell.set_failed(err.to_string());
                // Don't leave a failed slot cached.
                server.inner.cache.lock().unwrap().remove(&key);
            }
        });
        Handle::from_cell(cell)
    }

    /// Adds an already-constructed, in-memory asset and returns a handle to it.
    /// In-memory assets have no source path and are not cached for dedup.
    pub fn add<T: Send + Sync + 'static>(&self, value: T) -> Handle<T> {
        Handle::from_cell(AssetCell::new_loaded(self.next_id(), None, value))
    }

    /// Returns a handle to an already-loaded asset of this path+type, if one is
    /// still live, without triggering a load.
    pub fn get<T: Send + Sync + 'static>(&self, path: impl AsRef<Path>) -> Option<Handle<T>> {
        let key = (TypeId::of::<T>(), path.as_ref().to_path_buf());
        self.cached::<T>(&key)
    }

    /// Re-runs the loader for `path` and updates the existing asset in place, so
    /// every live handle observes the new contents. If no handle is currently
    /// live, behaves like [`load`](Self::load). This is the foundation the
    /// live-reload stage builds on.
    pub fn reload<T: Send + Sync + 'static>(&self, path: impl AsRef<Path>) -> Handle<T> {
        let path = path.as_ref().to_path_buf();
        let key = (TypeId::of::<T>(), path.clone());
        let existing = self.cached::<T>(&key);
        match self.run_loader::<T>(&path) {
            Ok(value) => match existing {
                Some(handle) => {
                    handle.set_loaded(value);
                    handle
                }
                None => {
                    let cell = AssetCell::new_loaded(self.next_id(), Some(path), value);
                    self.insert_cache(key, &cell);
                    Handle::from_cell(cell)
                }
            },
            Err(err) => match existing {
                Some(handle) => {
                    handle.set_failed(err.to_string());
                    handle
                }
                None => Handle::from_cell(AssetCell::new_failed(
                    self.next_id(),
                    Some(path),
                    err.to_string(),
                )),
            },
        }
    }

    /// The number of distinct assets still alive (have at least one live
    /// handle). Prunes collected entries as a side effect.
    pub fn live_asset_count(&self) -> usize {
        let mut cache = self.inner.cache.lock().unwrap();
        cache.retain(|_, entry| entry.weak.strong_count() > 0);
        cache.len()
    }

    /// Reruns the loader for every cached asset whose source path is `path`,
    /// updating each existing handle in place. Returns the number of assets
    /// reloaded.
    ///
    /// Unlike [`reload`](Self::reload) this does **not** need `T` at the call
    /// site — it dispatches on what types are actually cached for `path`. The
    /// file-watcher uses this to react to disk changes without knowing every
    /// asset type at compile time. Paths that are not currently cached return
    /// `0`; they will be loaded fresh by the next [`load`](Self::load) call.
    pub fn reload_path(&self, path: &Path) -> usize {
        // Snapshot the set of typed reload fns to call so we don't hold the
        // cache lock while re-running loaders (which would deadlock — `reload`
        // takes the lock too).
        let reloaders: Vec<fn(&AssetServer, &Path)> = {
            let cache = self.inner.cache.lock().unwrap();
            cache
                .iter()
                .filter_map(|(key, entry)| {
                    if key.1 == path && entry.weak.strong_count() > 0 {
                        Some(entry.reload_in_place)
                    } else {
                        None
                    }
                })
                .collect()
        };
        let n = reloaders.len();
        for f in reloaders {
            f(self, path);
        }
        n
    }

    // --- internals ---------------------------------------------------------

    fn next_id(&self) -> AssetId {
        AssetId(self.inner.next_id.fetch_add(1, Ordering::Relaxed))
    }

    fn cached<T: Send + Sync + 'static>(&self, key: &CacheKey) -> Option<Handle<T>> {
        let cache = self.inner.cache.lock().unwrap();
        let arc = cache.get(key)?.weak.upgrade()?;
        let cell = arc.downcast::<AssetCell<T>>().ok()?;
        Some(Handle::from_cell(cell))
    }

    fn insert_cache<T: Send + Sync + 'static>(&self, key: CacheKey, cell: &Arc<AssetCell<T>>) {
        let erased: Arc<dyn Any + Send + Sync> = cell.clone();
        // `reload_in_place` keeps the concrete `T` in its signature, so the
        // path-keyed `reload_path` can rebuild the typed handle without
        // knowing `T` at the call site.
        let entry = CacheEntry {
            weak: Arc::downgrade(&erased),
            reload_in_place: |server, path| {
                server.reload::<T>(path);
            },
        };
        self.inner.cache.lock().unwrap().insert(key, entry);
    }

    fn run_loader<T: Send + Sync + 'static>(&self, path: &Path) -> Result<T, AssetError> {
        let ext = path
            .extension()
            .and_then(|e| e.to_str())
            .ok_or_else(|| AssetError::NoExtension(path.to_path_buf()))?
            .to_lowercase();
        let loader = self
            .inner
            .loaders
            .read()
            .unwrap()
            .get(&ext)
            .cloned()
            .ok_or_else(|| AssetError::NoLoader(ext.clone()))?;
        if loader.output_type() != TypeId::of::<T>() {
            return Err(AssetError::TypeMismatch { ext });
        }
        let boxed = loader.load(path)?;
        Ok(*boxed
            .downcast::<T>()
            .expect("loader output_type matched the request but downcast failed"))
    }
}

impl Default for AssetServer {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::asset::LoadState;
    use std::sync::atomic::{AtomicU32, Ordering};

    // A trivial asset + loader: each "load" reads a file's text and counts how
    // many times the loader actually ran, so dedup can be observed.
    struct Counter(Arc<AtomicU32>);

    #[derive(Debug, PartialEq, Eq)]
    struct TextAsset(String);

    struct TextLoader(Arc<AtomicU32>);
    impl AssetLoader for TextLoader {
        type Asset = TextAsset;
        fn extensions(&self) -> &'static [&'static str] {
            &["txt"]
        }
        fn load(&self, path: &Path) -> Result<TextAsset, AssetError> {
            self.0.fetch_add(1, Ordering::SeqCst);
            let text = std::fs::read_to_string(path).map_err(|e| AssetError::Load {
                path: path.to_path_buf(),
                message: e.to_string(),
            })?;
            Ok(TextAsset(text.trim().to_string()))
        }
    }

    fn temp_file(name: &str, contents: &str) -> PathBuf {
        let mut path = std::env::temp_dir();
        path.push(format!(
            "oxide_asset_test_{}_{name}.txt",
            std::process::id()
        ));
        std::fs::write(&path, contents).unwrap();
        path
    }

    fn server() -> (AssetServer, Counter) {
        let counter = Arc::new(AtomicU32::new(0));
        let server = AssetServer::empty();
        server.register_loader(TextLoader(counter.clone()));
        (server, Counter(counter))
    }

    #[test]
    fn loads_and_reads_an_asset() {
        let (server, _c) = server();
        let path = temp_file("hello", "  hello world  ");
        let handle = server.load::<TextAsset>(&path);
        assert_eq!(handle.state(), LoadState::Loaded);
        assert_eq!(handle.get().unwrap().0, "hello world");
        assert_eq!(handle.source(), Some(path.as_path()));
        std::fs::remove_file(path).ok();
    }

    #[test]
    fn loading_twice_yields_one_resource() {
        let (server, c) = server();
        let path = temp_file("dedup", "data");
        let a = server.load::<TextAsset>(&path);
        let b = server.load::<TextAsset>(&path);
        // Same allocation: loader ran once, ids match, two handles share it.
        assert_eq!(c.0.load(Ordering::SeqCst), 1);
        assert_eq!(a.id(), b.id());
        assert_eq!(a.ref_count(), 2);
        assert_eq!(server.live_asset_count(), 1);
        std::fs::remove_file(path).ok();
    }

    #[test]
    fn dropping_all_handles_frees_the_asset() {
        let (server, _c) = server();
        let path = temp_file("free", "data");
        let handle = server.load::<TextAsset>(&path);
        assert_eq!(server.live_asset_count(), 1);
        drop(handle);
        // With no live handles, the weak cache entry is dead and pruned.
        assert_eq!(server.live_asset_count(), 0);
        assert!(server.get::<TextAsset>(&path).is_none());
        std::fs::remove_file(path).ok();
    }

    #[test]
    fn missing_loader_and_type_mismatch_are_distinct_errors() {
        let (server, _c) = server();
        let path = temp_file("x", "data");

        // No loader for ".dat".
        let bad_ext = path.with_extension("dat");
        std::fs::write(&bad_ext, "data").unwrap();
        let h = server.load::<TextAsset>(&bad_ext);
        assert_eq!(h.state(), LoadState::Failed);
        assert!(h.error().unwrap().contains("no loader"));

        // A ".txt" loader exists but produces TextAsset, not String.
        let renamed = path.with_extension("txt");
        std::fs::write(&renamed, "data").unwrap();
        let h2 = server.load::<String>(&renamed);
        assert!(h2.error().unwrap().contains("different asset type"));

        std::fs::remove_file(path).ok();
        std::fs::remove_file(bad_ext).ok();
        std::fs::remove_file(renamed).ok();
    }

    #[test]
    fn async_load_completes_and_dedups() {
        let (server, c) = server();
        let path = temp_file("async", "background");
        let handle = server.load_async::<TextAsset>(&path);
        let value = handle.wait().expect("async load should succeed");
        assert_eq!(value.0, "background");
        // A second request dedups onto the same now-loaded asset.
        let again = server.load::<TextAsset>(&path);
        assert_eq!(again.id(), handle.id());
        assert_eq!(c.0.load(Ordering::SeqCst), 1);
        std::fs::remove_file(path).ok();
    }

    #[test]
    fn reload_updates_in_place_for_existing_handles() {
        let (server, _c) = server();
        let path = temp_file("reload", "before");
        let handle = server.load::<TextAsset>(&path);
        assert_eq!(handle.get().unwrap().0, "before");

        // Change the file on disk and reload: the SAME handle sees new contents.
        std::fs::write(&path, "after").unwrap();
        let reloaded = server.reload::<TextAsset>(&path);
        assert_eq!(reloaded.id(), handle.id());
        assert_eq!(handle.get().unwrap().0, "after");
        std::fs::remove_file(path).ok();
    }

    #[test]
    fn add_stores_in_memory_assets() {
        let (server, _c) = server();
        let handle = server.add(TextAsset("in-memory".to_string()));
        assert_eq!(handle.get().unwrap().0, "in-memory");
        assert!(handle.source().is_none());
    }
}
