//! Resource descriptors and utilities to hand resources to tasks and bridges.
//! User view: in `copperconfig.ron`, map the binding names your tasks/bridges
//! expect to the resources exported by your board bundle. Exclusive things
//! (like a serial port) should be bound once; shared things (like a telemetry
//! bus `Arc`) can be bound to multiple consumers.
//!
//! ```ron
//! (
//!     resources: [ ( id: "board", provider: "board_crate::BoardBundle" ) ],
//!     bridges: [
//!         ( id: "crsf", type: "cu_crsf::CrsfBridge<SerialPort, SerialError>",
//!           resources: { serial: "board.uart0" }
//!         ),
//!     ],
//!     tasks: [
//!         ( id: "telemetry", type: "app::TelemetryTask",
//!           resources: { bus: "board.telemetry_bus" }
//!         ),
//!     ],
//! )
//! ```
//!
//! Writing your own task/bridge? Add a small `Resources` struct and implement
//! `ResourceBindings` to pull the names you declared:
//! ```rust,ignore
//! pub struct TelemetryResources<'r> { pub bus: Borrowed<'r, TelemetryBus> }
//! impl<'r> ResourceBindings<'r> for TelemetryResources<'r> {
//!     type Binding = Binding;
//!     fn from_bindings(mgr: &'r mut ResourceManager, map: Option<&ResourceBindingMap<Self::Binding>>) -> CuResult<Self> {
//!         let key = map.expect("bus binding").get(Binding::Bus).expect("bus").typed();
//!         Ok(Self { bus: mgr.borrow(key)? })
//!     }
//! }
//! pub fn new(_cfg: Option<&ComponentConfig>, res: TelemetryResources<'_>) -> CuResult<Self> {
//!     Ok(Self { bus: res.bus })
//! }
//! ```
//! Otherwise, use config to point to the right board resource and you're done.

use crate::config::ComponentConfig;
use core::any::Any;
use core::fmt;
use core::marker::PhantomData;
use cu29_traits::{CuError, CuResult};

use alloc::boxed::Box;
use alloc::format;
use alloc::sync::Arc;
use alloc::vec::Vec;

/// Lightweight wrapper used when a task needs to take ownership of a resource.
pub struct Owned<T>(pub T);

/// Wrapper used when a task needs to borrow a resource that remains managed by
/// the `ResourceManager`.
pub struct Borrowed<'r, T>(pub &'r T);

/// A resource can be exclusive (most common case) or shared.
enum ResourceEntry {
    Owned(Box<dyn Any + Send + Sync>),
    Shared(Arc<dyn Any + Send + Sync>),
}

impl ResourceEntry {
    fn as_shared<T: 'static + Send + Sync>(&self) -> Option<&T> {
        match self {
            ResourceEntry::Shared(arc) => arc.downcast_ref::<T>(),
            ResourceEntry::Owned(boxed) => boxed.downcast_ref::<T>(),
        }
    }

    #[cfg(feature = "std")]
    fn as_shared_arc<T: 'static + Send + Sync>(&self) -> Option<Arc<T>> {
        match self {
            ResourceEntry::Shared(arc) => Arc::downcast::<T>(arc.clone()).ok(),
            ResourceEntry::Owned(_) => None,
        }
    }

    fn into_owned<T: 'static + Send + Sync>(self) -> Option<T> {
        match self {
            ResourceEntry::Owned(boxed) => boxed.downcast::<T>().map(|b| *b).ok(),
            ResourceEntry::Shared(_) => None,
        }
    }
}

/// Typed identifier for a resource entry.
#[derive(Copy, Clone, Eq, PartialEq)]
pub struct ResourceKey<T = ()> {
    bundle: BundleIndex,
    index: usize,
    _boo: PhantomData<fn() -> T>,
}

impl<T> ResourceKey<T> {
    pub const fn new(bundle: BundleIndex, index: usize) -> Self {
        Self {
            bundle,
            index,
            _boo: PhantomData,
        }
    }

    pub const fn bundle(&self) -> BundleIndex {
        self.bundle
    }

    pub const fn index(&self) -> usize {
        self.index
    }

    /// Reinterpret this key as pointing to a concrete resource type.
    pub fn typed<U>(self) -> ResourceKey<U> {
        ResourceKey {
            bundle: self.bundle,
            index: self.index,
            _boo: PhantomData,
        }
    }
}

impl<T> fmt::Debug for ResourceKey<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("ResourceKey")
            .field("bundle", &self.bundle.index())
            .field("index", &self.index)
            .finish()
    }
}

/// Index identifying a resource bundle in the active mission.
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub struct BundleIndex(usize);

impl BundleIndex {
    pub const fn new(index: usize) -> Self {
        Self(index)
    }

    pub const fn index(self) -> usize {
        self.0
    }

    pub fn key<T, I: ResourceId>(self, id: I) -> ResourceKey<T> {
        ResourceKey::new(self, id.index())
    }
}

/// Trait implemented by resource id enums generated by `bundle_resources!`.
pub trait ResourceId: Copy + Eq {
    const COUNT: usize;
    fn index(self) -> usize;
}

/// Trait implemented by bundle providers to declare their resource id enum.
pub trait ResourceBundleDecl {
    type Id: ResourceId;
}

/// Static mapping between user-defined binding ids and resource keys.
#[derive(Clone, Copy)]
pub struct ResourceBindingMap<B: Copy + Eq + 'static> {
    entries: &'static [(B, ResourceKey)],
}

impl<B: Copy + Eq + 'static> ResourceBindingMap<B> {
    pub const fn new(entries: &'static [(B, ResourceKey)]) -> Self {
        Self { entries }
    }

    pub fn get(&self, binding: B) -> Option<ResourceKey> {
        self.entries
            .iter()
            .find(|(entry_id, _)| *entry_id == binding)
            .map(|(_, key)| *key)
    }
}

/// Manages the concrete resources available to tasks and bridges.
pub struct ResourceManager {
    bundles: Box<[BundleEntries]>,
}

struct BundleEntries {
    entries: Box<[Option<ResourceEntry>]>,
}

impl ResourceManager {
    /// Creates a new manager sized for the number of resources generated for
    /// each bundle in the current mission.
    pub fn new(bundle_sizes: &[usize]) -> Self {
        let bundles = bundle_sizes
            .iter()
            .map(|size| {
                let mut entries = Vec::with_capacity(*size);
                entries.resize_with(*size, || None);
                BundleEntries {
                    entries: entries.into_boxed_slice(),
                }
            })
            .collect::<Vec<_>>();
        Self {
            bundles: bundles.into_boxed_slice(),
        }
    }

    fn entry_mut<T>(&mut self, key: ResourceKey<T>) -> CuResult<&mut Option<ResourceEntry>> {
        let bundle = self
            .bundles
            .get_mut(key.bundle.index())
            .ok_or_else(|| CuError::from("Resource bundle index out of range"))?;
        bundle
            .entries
            .get_mut(key.index)
            .ok_or_else(|| CuError::from("Resource index out of range"))
    }

    fn entry<T>(&self, key: ResourceKey<T>) -> CuResult<&ResourceEntry> {
        let bundle = self
            .bundles
            .get(key.bundle.index())
            .ok_or_else(|| CuError::from("Resource bundle index out of range"))?;
        bundle
            .entries
            .get(key.index)
            .and_then(|opt| opt.as_ref())
            .ok_or_else(|| CuError::from("Resource not found"))
    }

    fn take_entry<T>(&mut self, key: ResourceKey<T>) -> CuResult<ResourceEntry> {
        let bundle = self
            .bundles
            .get_mut(key.bundle.index())
            .ok_or_else(|| CuError::from("Resource bundle index out of range"))?;
        let entry = bundle
            .entries
            .get_mut(key.index)
            .and_then(|opt| opt.take())
            .ok_or_else(|| CuError::from("Resource not found"))?;
        Ok(entry)
    }

    /// Register an owned resource in the slot identified by `key`.
    pub fn add_owned<T: 'static + Send + Sync>(
        &mut self,
        key: ResourceKey<T>,
        value: T,
    ) -> CuResult<()> {
        let entry = self.entry_mut(key)?;
        if entry.is_some() {
            return Err(CuError::from("Resource already registered"));
        }
        *entry = Some(ResourceEntry::Owned(Box::new(value)));
        Ok(())
    }

    /// Register a shared (borrowed) resource. Callers keep an `Arc` while tasks
    /// receive references.
    pub fn add_shared<T: 'static + Send + Sync>(
        &mut self,
        key: ResourceKey<T>,
        value: Arc<T>,
    ) -> CuResult<()> {
        let entry = self.entry_mut(key)?;
        if entry.is_some() {
            return Err(CuError::from("Resource already registered"));
        }
        *entry = Some(ResourceEntry::Shared(value as Arc<dyn Any + Send + Sync>));
        Ok(())
    }

    /// Borrow a shared resource by key.
    pub fn borrow<'r, T: 'static + Send + Sync>(
        &'r self,
        key: ResourceKey<T>,
    ) -> CuResult<Borrowed<'r, T>> {
        let entry = self.entry(key)?;
        entry.as_shared::<T>().map(Borrowed).ok_or_else(|| {
            CuError::from(format!(
                "Borrowing Resource has unexpected type, expected '{}'",
                core::any::type_name::<T>()
            ))
        })
    }

    /// Borrow a shared `Arc`-backed resource by key, cloning the `Arc` for the caller.
    #[cfg(feature = "std")]
    pub fn borrow_shared_arc<T: 'static + Send + Sync>(
        &self,
        key: ResourceKey<T>,
    ) -> CuResult<Arc<T>> {
        let entry = self.entry(key)?;
        entry.as_shared_arc::<T>().ok_or_else(|| {
            CuError::from(format!(
                "Borrow Shared Resource '{}' has unexpected type",
                core::any::type_name::<T>()
            ))
        })
    }

    /// Take ownership of a resource by key.
    pub fn take<T: 'static + Send + Sync>(&mut self, key: ResourceKey<T>) -> CuResult<Owned<T>> {
        let entry = self.take_entry(key)?;
        entry.into_owned::<T>().map(Owned).ok_or_else(|| {
            CuError::from(format!(
                "Resource {} is not owned or has unexpected type",
                core::any::type_name::<T>()
            ))
        })
    }

    /// Insert a prebuilt bundle by running a caller-supplied function. This is
    /// the escape hatch for resources that must be constructed in application
    /// code (for example, owning handles to embedded peripherals).
    pub fn add_bundle_prebuilt(
        &mut self,
        builder: impl FnOnce(&mut ResourceManager) -> CuResult<()>,
    ) -> CuResult<()> {
        builder(self)
    }
}

/// Trait implemented by resource binding structs passed to task/bridge
/// constructors. Implementors pull the concrete resources they need from the
/// `ResourceManager`, using the symbolic mapping provided in the Copper config
/// (`resources: { name: "bundle.resource" }`).
pub trait ResourceBindings<'r>: Sized {
    type Binding: Copy + Eq + 'static;

    fn from_bindings(
        manager: &'r mut ResourceManager,
        mapping: Option<&ResourceBindingMap<Self::Binding>>,
    ) -> CuResult<Self>;
}

impl<'r> ResourceBindings<'r> for () {
    type Binding = ();

    fn from_bindings(
        _manager: &'r mut ResourceManager,
        _mapping: Option<&ResourceBindingMap<Self::Binding>>,
    ) -> CuResult<Self> {
        Ok(())
    }
}

/// Bundle providers implement this trait to populate the `ResourceManager` with
/// concrete resources for a given bundle id.
pub trait ResourceBundle: ResourceBundleDecl + Sized {
    fn build(
        bundle: BundleContext<Self>,
        config: Option<&ComponentConfig>,
        manager: &mut ResourceManager,
    ) -> CuResult<()>;
}

/// Context passed to bundle providers when building resources.
pub struct BundleContext<B: ResourceBundleDecl> {
    bundle_index: BundleIndex,
    bundle_id: &'static str,
    _boo: PhantomData<B>,
}

impl<B: ResourceBundleDecl> BundleContext<B> {
    pub const fn new(bundle_index: BundleIndex, bundle_id: &'static str) -> Self {
        Self {
            bundle_index,
            bundle_id,
            _boo: PhantomData,
        }
    }

    pub const fn bundle_id(&self) -> &'static str {
        self.bundle_id
    }

    pub const fn bundle_index(&self) -> BundleIndex {
        self.bundle_index
    }

    pub fn key<T>(&self, id: B::Id) -> ResourceKey<T> {
        ResourceKey::new(self.bundle_index, id.index())
    }
}

#[cfg(feature = "std")]
pub struct ThreadPoolBundle;

#[cfg(feature = "std")]
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
#[repr(usize)]
pub enum ThreadPoolId {
    BgThreads,
}

#[cfg(feature = "std")]
impl ResourceId for ThreadPoolId {
    const COUNT: usize = 1;

    fn index(self) -> usize {
        self as usize
    }
}

#[cfg(feature = "std")]
impl ResourceBundleDecl for ThreadPoolBundle {
    type Id = ThreadPoolId;
}

#[cfg(feature = "std")]
impl ResourceBundle for ThreadPoolBundle {
    fn build(
        bundle: BundleContext<Self>,
        config: Option<&ComponentConfig>,
        manager: &mut ResourceManager,
    ) -> CuResult<()> {
        use rayon::ThreadPoolBuilder;

        const DEFAULT_THREADS: usize = 2;
        let threads: usize = match config {
            Some(cfg) => cfg
                .get::<u64>("threads")?
                .map(|v| v as usize)
                .unwrap_or(DEFAULT_THREADS),
            None => DEFAULT_THREADS,
        };

        let pool = ThreadPoolBuilder::new()
            .num_threads(threads)
            .build()
            .map_err(|e| CuError::from(format!("Failed to build threadpool: {e}")))?;

        let key = bundle.key::<rayon::ThreadPool>(ThreadPoolId::BgThreads);
        manager.add_shared(key, Arc::new(pool))?;
        Ok(())
    }
}
