//! Generic `all` + per-`AddrType` container, mirrors the `WithSth` pattern
//! along the address-type axis. Used by every metric that tracks one
//! aggregate value alongside a per-address-type breakdown.

use brk_cohort::ByAddrType;
use brk_error::Result;
use brk_indexer::Lengths;
use brk_traversable::Traversable;
use brk_types::{Cents, Height, Sats, Version};
use rayon::prelude::*;
use schemars::JsonSchema;
use vecdb::{AnyStoredVec, AnyVec, CachedBoxedVec, Database, EagerVec, Exit, PcoVec, WritableVec};

use crate::indexes;

use super::{
    CachedWindowStartVec, NumericValue, PerBlock, PerBlockCumulativeRolling, SpotValuePerBlock,
    Windows,
};

/// `all` aggregate plus per-`AddrType` breakdown.
#[derive(Clone, Traversable)]
pub struct WithAddrTypes<T> {
    pub all: T,
    #[traversable(flatten)]
    pub by_addr_type: ByAddrType<T>,
}

impl<T> WithAddrTypes<PerBlock<T>>
where
    T: NumericValue + JsonSchema,
{
    pub(crate) fn forced_import(
        db: &Database,
        name: &str,
        version: Version,
        indexes: &indexes::Vecs,
    ) -> Result<Self> {
        let all = PerBlock::forced_import(db, name, version, indexes)?;
        let by_addr_type = ByAddrType::new_with_name(|type_name| {
            PerBlock::forced_import(db, &format!("{type_name}_{name}"), version, indexes)
        })?;
        Ok(Self { all, by_addr_type })
    }

    pub(crate) fn min_stateful_len(&self) -> usize {
        self.by_addr_type
            .values()
            .map(|v| v.height.len())
            .min()
            .unwrap()
            .min(self.all.height.len())
    }

    pub(crate) fn par_iter_height_mut(
        &mut self,
    ) -> impl ParallelIterator<Item = &mut dyn AnyStoredVec> {
        rayon::iter::once(&mut self.all.height as &mut dyn AnyStoredVec).chain(
            self.by_addr_type
                .par_values_mut()
                .map(|v| &mut v.height as &mut dyn AnyStoredVec),
        )
    }

    pub(crate) fn reset_height(&mut self) -> Result<()> {
        self.all.height.reset()?;
        for v in self.by_addr_type.values_mut() {
            v.height.reset()?;
        }
        Ok(())
    }

    #[inline(always)]
    pub(crate) fn push_height<U>(&mut self, total: U, per_type: impl IntoIterator<Item = U>)
    where
        U: Into<T>,
    {
        self.all.height.push(total.into());
        for (v, value) in self.by_addr_type.values_mut().zip(per_type) {
            v.height.push(value.into());
        }
    }

    /// Compute `all.height` as the per-block sum of the per-type vecs.
    pub(crate) fn compute_rest(&mut self, starting_lengths: &Lengths, exit: &Exit) -> Result<()> {
        let sources: Vec<&EagerVec<PcoVec<Height, T>>> =
            self.by_addr_type.values().map(|v| &v.height).collect();
        self.all
            .height
            .compute_sum_of_others(starting_lengths.height, &sources, exit)?;
        Ok(())
    }
}

impl<T> WithAddrTypes<PerBlockCumulativeRolling<T>>
where
    T: NumericValue + JsonSchema,
{
    pub(crate) fn forced_import(
        db: &Database,
        name: &str,
        version: Version,
        indexes: &indexes::Vecs,
        cached_starts: &Windows<&CachedWindowStartVec>,
    ) -> Result<Self> {
        let all =
            PerBlockCumulativeRolling::forced_import(db, name, version, indexes, cached_starts)?;
        let by_addr_type = ByAddrType::new_with_name(|type_name| {
            PerBlockCumulativeRolling::forced_import(
                db,
                &format!("{type_name}_{name}"),
                version,
                indexes,
                cached_starts,
            )
        })?;
        Ok(Self { all, by_addr_type })
    }

    pub(crate) fn min_stateful_len(&self) -> usize {
        self.by_addr_type
            .values()
            .map(|v| v.cumulative.height.len())
            .min()
            .unwrap()
            .min(self.all.cumulative.height.len())
    }

    pub(crate) fn par_iter_height_mut(
        &mut self,
    ) -> impl ParallelIterator<Item = &mut dyn AnyStoredVec> {
        rayon::iter::once(&mut self.all.cumulative.height as &mut dyn AnyStoredVec).chain(
            self.by_addr_type
                .par_values_mut()
                .map(|v| &mut v.cumulative.height as &mut dyn AnyStoredVec),
        )
    }

    pub(crate) fn reset_height(&mut self) -> Result<()> {
        self.all.cumulative.height.reset()?;
        for v in self.by_addr_type.values_mut() {
            v.cumulative.height.reset()?;
        }
        Ok(())
    }

    #[inline(always)]
    pub(crate) fn push_height<U>(&mut self, total: U, per_type: impl IntoIterator<Item = U>)
    where
        U: Into<T>,
    {
        self.all.push_block(total.into());
        for (v, value) in self.by_addr_type.values_mut().zip(per_type) {
            v.push_block(value.into());
        }
    }
}

impl WithAddrTypes<SpotValuePerBlock> {
    pub(crate) fn forced_import(
        db: &Database,
        name: &str,
        version: Version,
        indexes: &indexes::Vecs,
        spot_price: &CachedBoxedVec<Height, Cents>,
    ) -> Result<Self> {
        let all = SpotValuePerBlock::forced_import(db, name, version, indexes, spot_price)?;
        let by_addr_type = ByAddrType::new_with_name(|type_name| {
            SpotValuePerBlock::forced_import(
                db,
                &format!("{type_name}_{name}"),
                version,
                indexes,
                spot_price,
            )
        })?;
        Ok(Self { all, by_addr_type })
    }

    pub(crate) fn min_stateful_len(&self) -> usize {
        self.by_addr_type
            .values()
            .map(|v| v.sats.height.len())
            .min()
            .unwrap()
            .min(self.all.sats.height.len())
    }

    pub(crate) fn par_iter_height_mut(
        &mut self,
    ) -> impl ParallelIterator<Item = &mut dyn AnyStoredVec> {
        rayon::iter::once(&mut self.all.sats.height as &mut dyn AnyStoredVec).chain(
            self.by_addr_type
                .par_values_mut()
                .map(|v| &mut v.sats.height as &mut dyn AnyStoredVec),
        )
    }

    pub(crate) fn reset_height(&mut self) -> Result<()> {
        self.all.sats.height.reset()?;
        for v in self.by_addr_type.values_mut() {
            v.sats.height.reset()?;
        }
        Ok(())
    }

    /// Push the stateful sats value for `all` and each per-type.
    #[inline(always)]
    pub(crate) fn push_height<U>(&mut self, total: U, per_type: impl IntoIterator<Item = U>)
    where
        U: Into<Sats>,
    {
        self.all.sats.height.push(total.into());
        for (v, value) in self.by_addr_type.values_mut().zip(per_type) {
            v.sats.height.push(value.into());
        }
    }
}
