use std::{
    collections::BTreeMap,
    fmt::Debug,
    sync::{Arc, OnceLock},
};

use crossbeam_skiplist::SkipMap;
use rayon::{iter::IndexedParallelIterator, prelude::ParallelIterator};
use sled::InlineArray;
use tracing::debug;
use xxhash_rust::xxh3::Xxh3Default;

use super::{KeyValue, raw_tree::RawTree, replication::ReplicationEvent, tree::EventHandlers};
use crate::{
    Cid, GetType,
    database::{
        db::Database,
        error::{DbError, DbResult},
        raw_tree::InnerFields,
        tree::{CollectionTree, IndexesMap},
    },
    indexes::IndexUpdateCollection,
    record::{MigrationFn, MigrationRes, MigrationType, Record},
    utils::{
        Change, CollectionInfo, GroupVal, IdChangeEventFn, MaybeParIter, RemoveEventFn,
        UpsertEventFn,
    },
};

pub struct CollectionBuilder<Item: Record> {
    name: String,
    db: Database,
    tree: super::Tree,
    inner: InnerFields,
    index_maps: IndexesMap<Item>,
    handlers: EventHandlers<Item>,
    lazy: bool,
}

impl<Item> CollectionBuilder<Item>
where
    Item: Record<Value = InlineArray>,
{
    #[tracing::instrument(name = "CollectionBuilder::with_name")]
    pub(crate) fn with_name(db: Database, name: &str) -> Self {
        let tree = db.open_tree(name);

        tracing::debug!("tree opened");

        let info = db
            .db_info
            .cloned()
            .collections
            .get(name)
            .cloned()
            .unwrap_or_default();

        if info.seq != 0 {
            tracing::debug!(seq = info.seq);
        }

        let hashpoints = SkipMap::new();

        let inner = InnerFields {
            info,
            seq: Default::default(),
            hashpoints,
            replication_handler: None,
        };

        Self {
            name: name.to_owned(),
            db: db.clone(),
            tree,
            inner,
            index_maps: Default::default(),
            handlers: Default::default(),
            lazy: false,
        }
    }

    pub(crate) fn new(db: Database) -> Self {
        Self::with_name(db, <Item as Record>::NAME)
    }

    pub fn get_seq(&self) -> usize {
        self.inner.info.seq as usize
    }

    pub fn add_index<T>(mut self, index: T) -> Self
    where
        T: IndexUpdateCollection<Item> + 'static,
    {
        self.index_maps.push(Box::new(index));
        self
    }

    /// Fn(&Item, old_id: &Id, new_id: &Id)
    pub fn add_id_change_fns(mut self, f: IdChangeEventFn<Item>) -> Self {
        self.handlers.on_id_change = Some(f);
        self
    }

    /// Fn(&Id, Change {old: &Item, new: &Item})
    pub fn add_update_fn(mut self, f: UpsertEventFn<Item>) -> Self {
        self.handlers.on_upsert = Some(f);
        self
    }

    /// Fn(&Id, &Item)
    pub fn add_removed_fn(mut self, f: RemoveEventFn<Item>) -> Self {
        self.handlers.on_remove = Some(f);
        self
    }

    pub fn set_lazy(mut self, lazy: bool) -> Self {
        self.lazy = lazy;
        self
    }

    // iterate tree and compute hash for each N items with bromberg_sl2 crate
    #[tracing::instrument(name = "CollectionBuilder::build", skip(self), fields(name = self.name))]
    pub fn build(self) -> CollectionTree<Item>
    where
        Item: GetType + Debug,
    {
        let CollectionBuilder {
            db,
            tree,
            mut inner,
            name,
            index_maps,
            handlers,
            lazy: _,
        } = self;

        let index_maps = Arc::new(index_maps);
        let handlers = Arc::new(handlers);

        let has_handlers = handlers.has_handlers() || !index_maps.is_empty();

        if has_handlers {
            debug!("has handlers");
            let cloned_indexes = index_maps.clone();
            let cloned_handlers = handlers.clone();
            inner.replication_handler = Some(Box::new(move |event| {
                let id = Item::deser_key(event.key());
                let item = Item::deser(event.val());
                match event {
                    ReplicationEvent::Upsert { old_val, .. } => {
                        let old_item = old_val.map(Item::deser);
                        for map in cloned_indexes.iter() {
                            map.update(&id, old_item.as_ref(), Some(&item));
                        }
                        if let Some(f) = cloned_handlers.on_upsert.as_ref() {
                            f(
                                &id,
                                Change {
                                    old: old_item.as_ref(),
                                    new: &item,
                                },
                            );
                        }
                    }
                    ReplicationEvent::Delete { .. } => {
                        for map in cloned_indexes.iter() {
                            map.update(&id, Some(&item), None);
                        }
                        if let Some(f) = cloned_handlers.on_remove.as_ref() {
                            f(&id, &item);
                        }
                    }
                    ReplicationEvent::IdChange { old_key, .. } => {
                        for map in cloned_indexes.iter() {
                            map.update(&id, None, Some(&item));
                        }
                        let old_id = Item::deser_key(old_key);
                        if let Some(f) = cloned_handlers.on_id_change.as_ref() {
                            f(&item, &old_id, &id);
                        }
                    }
                }
            }));
        }

        let removed = db.open_tree(&format!("__{name}_removed"));

        let raw = RawTree {
            db,
            name,
            attributes: &Item::ATTRIBUTES,
            tree,
            removed,
            inner: Arc::new(inner),
            subscribers: OnceLock::new(),
        };

        let collection = CollectionTree {
            raw,
            index_maps,
            handlers,
        };

        let current_type_hash = <Item as GetType>::TYPE.h();

        let mut current_seq = 0;

        if !collection.raw.is_empty() {
            let instant = std::time::Instant::now();
            let info = collection.raw.inner.info;

            if info.typ_hash != current_type_hash {
                tracing::warn!(
                    "collection type hash changed from {:#X} to {:#X}",
                    info.typ_hash,
                    current_type_hash
                );
            }

            let iter = collection
                .par_iter(Some(info.seq))
                .expect("cannot get par iter");

            let migration = Item::MIGRATIONS.iter().copied().find_map(|(version, f)| {
                if version == info.version {
                    tracing::info!(?version, "found migration");
                    Some(f)
                } else {
                    None
                }
            });

            let span = tracing::info_span!("migration collection").entered();

            let hash: u64 = match iter {
                MaybeParIter::Par(iter) => {
                    let mut vec = Vec::with_capacity(64);
                    iter.map(|(iter, table)| {
                        let iter = iter.map(|res| res.map_err(DbError::from));
                        let mut map = BTreeMap::<u32, Xxh3Default>::new();
                        handle_items(iter, migration, &table, &mut map)
                            .expect("cannot compute hash");
                        map
                    })
                    .collect_into_vec(&mut vec);
                    let iter = vec.into_iter().flat_map(|map| map.into_iter());
                    let mut hasher = Xxh3Default::new();
                    for (k, v) in iter {
                        hasher.update(&k.to_le_bytes());
                        collection.raw.inner.hashpoints.insert(
                            k,
                            GroupVal {
                                hash: v.digest(),
                                changed: false,
                            },
                        );
                    }
                    hasher.digest()
                }
                MaybeParIter::Seq(iter) => {
                    let mut map = BTreeMap::<u32, Xxh3Default>::new();
                    let iter = iter.map(|res| res.map_err(DbError::from));
                    handle_items(iter, migration, &collection, &mut map)
                        .expect("cannot compute hash");
                    let mut hasher = Xxh3Default::new();
                    for (k, v) in map {
                        hasher.update(&k.to_le_bytes());
                        collection.raw.inner.hashpoints.insert(
                            k,
                            GroupVal {
                                hash: v.digest(),
                                changed: false,
                            },
                        );
                    }
                    hasher.digest()
                }
                MaybeParIter::Empty => {
                    tracing::error!("empty collection, but tree is not");
                    0
                }
            };

            span.exit();

            let time = instant.elapsed().as_millis();

            tracing::debug!(?hash, ?time);

            let count = collection.count();
            current_seq = count;

            if info.seq != current_seq {
                tracing::warn!(saved = info.seq, current_seq, "saved seq != current seq");
            }
        }

        let version = <Item as Record>::VERSION;

        let collection_info = CollectionInfo {
            typ_hash: current_type_hash,
            seq: current_seq,
            version,
        };
        if version > 0 {
            tracing::debug!(version);
        }

        collection
            .raw
            .db
            .db_info
            .data
            .lock()
            .collections
            .insert(collection.raw.name.clone(), collection_info);

        collection
    }
}

#[tracing::instrument(name = "handle_items", skip(iter, map))]
fn handle_items<Item>(
    iter: impl Iterator<Item = Result<KeyValue, DbError>>,
    migration: Option<MigrationFn<Item, InlineArray>>,
    collection: &CollectionTree<Item>,
    map: &mut BTreeMap<u32, Xxh3Default>,
) -> DbResult<()>
where
    Item: Record<Value = InlineArray> + Debug,
{
    debug!("handle_items start");
    for item in iter {
        let (k, bytes) = item?;

        let entry = if let Some(migration) = migration {
            let key = k.as_ref();
            match migration(key, &bytes) {
                MigrationRes::Unchanged(k, v) => Some((k, v)),
                MigrationRes::Deleted => {
                    collection.raw.tree.remove(key)?;
                    None
                }
                MigrationRes::Changed {
                    migration_type,
                    new_key,
                    new_value,
                } => {
                    match migration_type {
                        MigrationType::Key => {
                            collection.raw.tree.remove(key)?;
                            let key = Cid::encode(&new_key);
                            collection.raw.tree.insert(key.as_ref(), bytes.clone())?;
                        }
                        MigrationType::Value => {
                            let val = Record::ser(&new_value);
                            collection.raw.tree.insert(key, val)?;
                        }
                        MigrationType::Entry => {
                            collection.raw.tree.remove(key)?;
                            let key = Cid::encode(&new_key);
                            let val = Record::ser(&new_value);
                            collection.raw.tree.insert(key.as_ref(), val)?;
                        }
                    }
                    Some((new_key, new_value))
                }
            }
        } else {
            let key = <Item as Record>::deser_key(&k);
            let val = <Item as Record>::deser(&bytes);

            Some((key, val))
        };

        let Some((key, val)) = entry else {
            continue;
        };

        collection.update_indexes(&key, None, Some(&val));
        collection.upsert_event(&key, None, &val);

        let group_id = key.group_id();
        let hasher = map.entry(group_id).or_default();
        hasher.update(&k);
        hasher.update(&bytes);
    }

    Ok(())
}
