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

use crossbeam_skiplist::SkipMap;
use fjall::{Keyspace, KeyspaceCreateOptions, Slice};
use xxhash_rust::xxh3::Xxh3Default;

use super::{
    db::Db,
    events::change::*,
    raw::{InnerFields, RawTree},
    tree::TypedTree,
};
use crate::{
    Cid, GetType, Record,
    error::{DbError, DbResult},
    logdb::ReadTree,
    record::{MigrationFn, MigrationRes, MigrationType},
    utils::{CollectionInfo, GroupVal},
};

pub struct TreeBuilder<Item: Record> {
    name: String,
    db: Db,
    tree: Keyspace,
    inner: InnerFields,
    handlers: LogDbEventHandlers<Item>,
}

impl<Item: Record> TreeBuilder<Item> {
    #[instrument(name = "TreeBuilder::with_name", skip(db, options))]
    pub(crate) fn with_name(db: &Db, name: &str, options: Option<KeyspaceCreateOptions>) -> Self {
        let start = std::time::Instant::now();
        let tree = db.open_tree(name, options);

        debug!("tree opened");

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

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

        let hashpoints = SkipMap::new();

        let inner = InnerFields { info, hashpoints };

        histogram!("armour_logdb_builder_with_name_duration", "name" => name.to_owned())
            .record(start.elapsed().as_secs_f64());

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

    pub(crate) fn new(db: &Db, options: Option<KeyspaceCreateOptions>) -> Self {
        Self::with_name(db, <Item as Record>::NAME, options)
    }

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

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

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

    pub fn build(self) -> TypedTree<Item>
    where
        Item: Record<Value = Slice> + GetType + Debug,
    {
        self.inner_build(None::<fn(&Item::SelfId, &Item)>)
    }

    pub fn build_with_handler<F>(self, f: F) -> TypedTree<Item>
    where
        Item: Record<Value = Slice> + GetType + Debug,
        F: for<'a> FnMut(&'a Item::SelfId, &'a Item),
    {
        self.inner_build(Some(f))
    }

    /// Fn(&Id, &Item)
    // iterate tree and compute hash for each N items with bromberg_sl2 crate
    #[instrument(name = "TreeBuilder::build", skip(self, f), fields(name = self.name))]
    fn inner_build<F>(self, f: Option<F>) -> TypedTree<Item>
    where
        Item: Record<Value = Slice> + GetType + Debug,
        F: for<'a> FnMut(&'a Item::SelfId, &'a Item),
    {
        let start = std::time::Instant::now();
        let TreeBuilder {
            db,
            tree,
            inner,
            name,
            handlers,
        } = self;

        let is_empty = tree.is_empty().expect("cannot check if tree is empty");

        let info = inner.info;

        let removed = db
            .db
            .keyspace(&format!("__{name}_removed"), || {
                KeyspaceCreateOptions::default()
            })
            .expect("Failed to open partition");

        let raw = RawTree {
            db: db.clone(),
            name: name.clone(),
            attributes: &Item::ATTRIBUTES,
            tree,
            removed,
            inner: Arc::new(inner),
            meta_saved: Default::default(),
            seq_lock: Arc::default(),
        };

        let handlers = Arc::new(handlers);

        let collection = TypedTree {
            inner: ReadTree {
                raw,
                phantom: std::marker::PhantomData,
            },
            handlers,
        };

        let current_type_hash = <Item as GetType>::TYPE.h();
        let mut seq = collection.tree().approximate_len() as u64;

        if !is_empty {
            let instant = std::time::Instant::now();

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

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

            debug!("starting tree handle_items");
            let res = handle_items(&collection, migration, f);

            let (hash, count) = res.expect("cannot compute hash");
            let time = instant.elapsed().as_millis();

            debug!(?hash, ?count, ?time, "tree handle_items end");
            seq = count;
        }

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

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

        db.db_info.update(|info| {
            info.collections.insert(name.clone(), collection_info);
        });

        histogram!("armour_logdb_builder_build_duration", "name" => name.clone())
            .record(start.elapsed().as_secs_f64());

        collection
    }
}

#[instrument(name = "handle_items", skip(handle_fn))]
fn handle_items<Item, F>(
    collection: &TypedTree<Item>,
    migration: Option<MigrationFn<Item, Slice>>,
    mut handle_fn: Option<F>,
) -> DbResult<(u64, u64)>
where
    Item: Record<Value = Slice> + Debug,
    F: for<'a> FnMut(&'a Item::SelfId, &'a Item),
{
    let iter = collection.tree().iter();

    let mut map = BTreeMap::<u32, Xxh3Default>::new();

    let mut count = 0u64;

    for item in iter {
        let (k, bytes) = item.into_inner()?;

        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.tree().remove(key)?;
                    None
                }
                MigrationRes::Changed {
                    migration_type,
                    new_key,
                    new_value,
                } => {
                    match migration_type {
                        MigrationType::Key => {
                            collection.tree().remove(key)?;
                            let key = Cid::encode(&new_key);
                            collection.tree().insert(key.as_ref(), bytes.clone())?;
                        }
                        MigrationType::Value => {
                            let val = Record::ser(&new_value);
                            collection.tree().insert(key.as_ref(), val)?;
                        }
                        MigrationType::Entry => {
                            collection.tree().remove(key)?;
                            let key = Cid::encode(&new_key);
                            let val = Record::ser(&new_value);
                            collection.tree().insert(key.as_ref(), val)?;
                        }
                    }
                    Some((new_key, new_value))
                }
            }
        } else {
            let key = <Item as Record>::SelfId::from_bytes(&k).map_err(DbError::from)?;
            let val = <Item as Record>::deser(&bytes);

            Some((key, val))
        };

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

        if let Some(f) = handle_fn.as_mut() {
            f(&key, &val);
        }
        collection.upsert_event(&key, &val);

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

    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,
            },
        );
    }

    Ok((hasher.digest(), count))
}
