
use std::{fmt, sync::{Arc, MutexGuard, RwLock, RwLockReadGuard, RwLockWriteGuard, Weak}, vec};
use croaring::Bitmap;
use pyo3::prelude::*;
use rustc_hash::FxHashMap;
use smallvec::SmallVec;
use smol_str::SmolStr;

use crate::index::{HybridHashmap, Indexable, PyQueryExpr, core::{query::{BulkQueryMapAdder, query_ops::{QueryExpr, evaluate_and_queries_vec}}, structures::{hybrid_set::{HybridSet, HybridSetOps}, m2m::M2MU32, string_interner::INTERNER}}, interfaces::filtered_index::FilteredIndex, types::{DEFAULT_INDEXABLE_ARC, IndexTree, StrId}};
use crate::index::core::query::{QueryMap, attr_parts, evaluate_query};

use crate::index::core::stored_item::StoredItem;
use crate::index::value::PyValue;

const QUERY_DEPTH_LEN: usize = 12;

#[derive(Clone, Default)]
pub struct IndexAPI{
    pub index: IndexTree,
    pub items: Arc<RwLock<Vec<StoredItem>>>,
    pub allowed_items: Arc<RwLock<Bitmap>>,
    pub parent_child_map: Arc<RwLock<M2MU32>>,
    pub parent_index: Option<Weak<IndexAPI>>,
}

impl IndexAPI{

    pub fn new(parent_index: Option<Weak<IndexAPI>>) -> Self {
        Self {
            index: Arc::new(RwLock::new(vec![])),
            items: Arc::new(RwLock::new(vec![])),
            allowed_items: Arc::new(RwLock::new(Bitmap::new())),
            parent_child_map: Arc::new(RwLock::new(M2MU32::new())),
            parent_index: parent_index,
        }
    }

    pub fn collect(&self, py:Python) -> PyResult<Vec<Py<Indexable>>> {
        let mut result = vec![];
        let allowed_items = self.get_allowed_items_reader();
        let items_reader = self.get_items_reader();

        for idx in allowed_items.iter(){
            result.push(items_reader[idx as usize].get_py_ref(py));
        }
        Ok(result)
    }

    pub fn get_from_parent_ids(&self, parent_ids: &Bitmap) -> Bitmap {
        self.get_parent_child_map_reader().get_for_forward_many(parent_ids)
    }

    pub fn get_parent_from_ids(&self, ids: &Bitmap) -> Bitmap {
        let res = self.get_parent_child_map_reader().get_for_reverse_many(ids);
        res
    }

    pub fn add_object_many(
        &self,
        weak_self: Weak<Self>,
        raw_objs: Vec<(Indexable, Py<Indexable>)>
    ) {
        // 3 pass - wrap in ARC - add meta to index with locks - add to index maps which may call meta locks
        let arc_objs: Vec<(Arc<Indexable>, Arc<Py<Indexable>>)> = raw_objs
            .into_iter()
            .map(|(idx, py)| (Arc::new(idx), Arc::new(py)))
            .collect();

        let mut allowed_writer: RwLockWriteGuard<'_, Bitmap> = self.get_allowed_items_writer();
        let mut items_writer = self.get_items_writer();

        for (rust_handle, py_handle) in &arc_objs {

            rust_handle.add_index(weak_self.clone());
            allowed_writer.add(rust_handle.id);

            let idx = rust_handle.id as usize;
            let stored_item = StoredItem::new(py_handle.clone(), rust_handle.clone());

            if items_writer.len() <= idx{
                items_writer.resize(idx * 2, StoredItem::default());
            }

            items_writer[idx] = stored_item;

        }
        drop(allowed_writer);
        drop(items_writer);

        let mut index_reader = self.get_index_reader();
        let mut delayed_adders: Vec<BulkQueryMapAdder> = index_reader.iter().map(|i| {
            i.get_bulk_writer()
        }).collect();

        for (rust_handle, _) in arc_objs {
            let object_id = rust_handle.id;
            for (key, value) in rust_handle.get_py_values().iter() {
                let attr_id = *key as usize;
                if let Some(qmap) = delayed_adders.get_mut(attr_id) {
                    qmap.insert(value, object_id);
                } else {
                    drop(delayed_adders);
                    drop(index_reader);

                    let mut writer = self.get_index_writer();
                    let qmap = QueryMap::new(weak_self.clone());
                    qmap.insert(value, object_id);
                    if attr_id >= writer.len() {
                        writer.resize_with((attr_id + 1) as usize, Default::default); // or None if Option
                    }
                    writer[attr_id as usize] = qmap;
                    
                    drop(writer);
                    index_reader = self.get_index_reader();
                    delayed_adders = index_reader.iter().map(|i| {
                        i.get_bulk_writer()
                    }).collect();
                }
            }
        }

//        for (rust_handle, _) in arc_objs {
//            for (key, value) in rust_handle.get_py_values().iter() {
//                self.add_index(weak_self.clone(), rust_handle.id, *key, value);
//            }
//        }
    }

    pub fn has_object_id(&self, id: u32) -> bool {
        !Arc::ptr_eq(
            self.get_items_reader().get(id as usize).unwrap_or(&StoredItem::default()).get_owned_handle(),
            &DEFAULT_INDEXABLE_ARC
        )
    }

    pub fn register_path(&self, object_id: u32, parent_id: u32) {
        self.get_parent_child_map_writer().add(object_id, parent_id);
//        let mut writer = self.get_items_writer();
//        let obj = writer.get_mut(object_id as usize).unwrap();
//        obj.add_parent(parent_id);
    }

    pub fn deregister_path(&self, object_id: u32, parent_id: u32) {
        self.get_parent_child_map_writer().remove(parent_id, object_id);
    }

    pub fn add_object(
        &self,
        weak_self: Weak<IndexAPI>,
        idx: u32,
        stored_item: StoredItem,
        py_val_hashmap: MutexGuard<HybridHashmap<StrId, PyValue>>
    ) {

        self.get_allowed_items_writer().add(idx);
        {
            let mut items_writer = self.get_items_writer();
            if items_writer.len() <= idx as usize{
                items_writer.resize(idx as usize + 1, StoredItem::default());
            }

            items_writer[idx as usize] = stored_item;
        }

        for (attr_id, value) in py_val_hashmap.iter() {
            // if key.starts_with("_"){continue;}
            self.add_index(weak_self.clone(), idx, *attr_id, value);
        }
    }

    pub fn get_parents_from_id(&self, id: usize) -> Bitmap {
        let mut res = self.get_parent_child_map_reader().get_for_reverse(id as u32);
        if let Some(p_index) = &self.parent_index {
            if let Some(p_index) = p_index.upgrade() {
                let parent_parents = p_index.get_parent_from_ids(&res);
                res.or_inplace(&parent_parents);
            }
        }
        res
    }

    pub fn remove(&self, item: &Indexable, parent_id: u32) {
        let item_id = item.id;
        let mut parent_child_map_writer = self.get_parent_child_map_writer();
        
        if parent_child_map_writer.get_for_reverse(parent_id).cardinality() == 1 {
            parent_child_map_writer.remove(parent_id, item_id);
            let mut items_writer = self.get_items_writer();
            items_writer[item_id as usize] = StoredItem::default();
            drop(items_writer);

            for (key, value) in (*item.get_py_values()).iter(){
                // if key.starts_with("_"){continue;}
                self.remove_index(item_id, *key as usize, value);
            }

            self.get_allowed_items_writer().remove(item_id);
        }
    }
    

    pub fn keep_only_with_parent_ids(&self, parent_ids: &Bitmap) {
        let to_keep = self.get_from_parent_ids(parent_ids);
        self.keep_only_from_bitmap(&to_keep);
    }

    pub fn reduce(
        &self,
        query: FxHashMap<SmolStr, PyValue>,
    ){
        // this one is not ideal as it needs to trim in place
        // ideally this is a bunch of andnot_inplace calls
        let index = self.get_index_reader();
        let all_valid = self.get_allowed_items_reader();
        let exprs: Vec<QueryExpr> = query.into_iter().map(|(k, v)| {
            QueryExpr::Eq(k, v)
        }).collect();
        
        let keep = evaluate_and_queries_vec(&index, &all_valid, &exprs);
        drop(all_valid);
        self.keep_only_from_bitmap(&keep);
    }

    fn keep_only_from_bitmap(&self, keep: &Bitmap) {
        let index = self.get_index_reader();
        let to_remove = self.get_allowed_items_reader().andnot(&keep);
        
        let mut allowed_items = self.get_allowed_items_writer();
        allowed_items.and_inplace(&keep);
        drop(allowed_items);

        for map in index.iter() {
            map.keep_only(&keep);
        }
        let mut stored_items = self.get_items_writer();
        for idx in to_remove.iter(){
            stored_items[idx as usize] = StoredItem::default();
        }
    }


    pub fn reduced(
        &self,
        query: FxHashMap<SmolStr, PyValue>
    ) -> FilteredIndex {
        let index = self.get_index_reader();
        let all_valid = self.get_allowed_items_reader();
        let exprs: Vec<QueryExpr> = query.into_iter().map(|(k, v)| {
            QueryExpr::Eq(k, v)
        }).collect();
        self.filter_from_bitmap(
            evaluate_and_queries_vec(&index, &all_valid, &exprs)
        )
    }

    pub fn reduced_query(
        &self,
        query: PyQueryExpr,
    ) -> FilteredIndex {
        let index = self.get_index_reader();
        let allowed = self.get_allowed_items_reader();
        self.filter_from_bitmap(
            evaluate_query(&index, &allowed, &query.inner)
        )
    }

    pub fn union_with(&self, other: &IndexAPI) -> PyResult<()>{
        let mut self_index = self.get_index_reader();
        let other_index = other.get_index_reader();

        if self_index.len() < other_index.len() {
            let additional = other_index.len() - self_index.len();
            drop(self_index);
            let mut self_index_writer = self.get_index_writer();
            self_index_writer.reserve(additional);
            drop(self_index_writer);
            self_index = self.get_index_reader();
        }
        
        for (self_qm, other_qm) in self_index.iter().zip(other_index.iter()) {
            self_qm.merge(other_qm);
        }

        let mut items_writer = self.get_items_writer();

        // iterate other bitset to get allowed items
        let other_allowed_items_reader = other.get_allowed_items_reader();
        let other_items_reader = other.get_items_reader();

        for idx in other_allowed_items_reader.iter(){
            let other_item = other_items_reader.get(idx as usize).unwrap().clone();
            if items_writer.len() <= idx as usize{
                items_writer.resize(usize::max(idx as usize * 2, 1), StoredItem::default());
            }
            items_writer[idx as usize] = other_item;
        }
        self.get_allowed_items_writer().or_inplace(&other_allowed_items_reader);

        Ok(())
    }


    pub fn update_index(
        &self,
        weak_self: Weak<IndexAPI>,
        attr: StrId, 
        old_pv: Option<&PyValue>,
        new_pv: &PyValue,
        item_id: u32,
    ) {
//        if attr.starts_with("_") {
//            return;
//        }
        
        if let Some(old_val) = old_pv {
            self.remove_index(item_id, attr as usize, old_val);
        }
        self.add_index(weak_self, item_id, attr, &new_pv);
    }

    pub fn get_from_indexes(&self, py: Python, indexes: Bitmap) -> PyResult<Vec<Py<Indexable>>>{
        let items_read = self.get_items_reader();
        let results: Vec<Py<Indexable>> = indexes.iter()
            .map(|arc| items_read.get(arc as usize).unwrap().get_py_ref(py))
            .collect();

        Ok(results)
    }

    pub fn add_index(
        &self,
        weak_self: Weak<IndexAPI>,
        obj_id: u32,
        attr_id: StrId,
        value: &PyValue
    ){
        if let Some(qmap) = self.get_index_reader().get(attr_id as usize) {
            qmap.insert(value, obj_id);
            return;
        }

        let qmap = QueryMap::new(weak_self);
        qmap.insert(value, obj_id);
        let mut writer = self.get_index_writer();

        if attr_id >= writer.len() as u32 {
            writer.resize_with((attr_id + 1) as usize, Default::default); // or None if Option
        }
        writer[attr_id as usize] = qmap;
    }

    fn remove_index(
        &self,
        idx: u32,
        attr_id: usize,
        py_value: &PyValue
    ){
        let index = self.get_index_reader();
        if index.len() > attr_id {
            if let Some(val) = index.get(attr_id) { 
                val.remove_id(py_value, idx);
                val.check_prune(py_value);
            };

            if index[attr_id].is_empty(){
                drop(index);
                self.get_index_writer()[attr_id] = Default::default();
            }
        }
    }

    pub fn filter_from_bitmap(&self, bm: Bitmap) -> FilteredIndex {
        FilteredIndex {
            index: self.index.clone(),
            items: self.items.clone(),
            allowed_items: bm
        }
    }

    pub fn is_attr_equal(&self, id: usize, str_id: StrId, val: &PyValue) -> bool {
        self.get_items_reader()
            .get(id)
            .and_then(|item| {
                item.with_attr_id(str_id, |py_val| py_val == val)
            })
            .unwrap_or(false)
    }

    fn get_items_writer(&self) -> RwLockWriteGuard<'_, Vec<StoredItem>> {
        self.items.write().unwrap()
        //self.items.try_write().expect("items writer deadlock")
    }

    fn get_items_reader(&self) -> RwLockReadGuard<'_, Vec<StoredItem>> {
        self.items.read().unwrap()
        //self.items.try_read().expect("cannot read from items")
    }

    pub fn get_index_writer(&self) -> RwLockWriteGuard<'_, Vec<QueryMap>> {
        self.index.write().unwrap()
        //self.index.try_write().expect("index writer deadlock")
    }

    pub fn get_index_reader(&self) -> RwLockReadGuard<'_, Vec<QueryMap>> {
        self.index.read().unwrap()
        //self.index.try_read().expect("cannot read from index")
    }

    fn get_allowed_items_writer(&self) -> RwLockWriteGuard<'_, Bitmap> {
        self.allowed_items.write().unwrap()
        //self.allowed_items.try_write().expect("index writer deadlock")
    }

    fn get_allowed_items_reader(&self) -> RwLockReadGuard<'_, Bitmap> {
        self.allowed_items.read().unwrap()
        //self.allowed_items.try_read().expect("cannot read from index")
    }

    fn get_parent_child_map_reader(&self) -> RwLockReadGuard<'_, M2MU32> {
        self.parent_child_map.read().unwrap()
        //self.parent_child_map.try_read().expect("cannot read from index")
    }

    fn get_parent_child_map_writer(&self) -> RwLockWriteGuard<'_, M2MU32> {
        self.parent_child_map.write().unwrap()
        //self.parent_child_map.try_write().expect("cannot read from index")
    }
}

impl fmt::Debug for IndexAPI {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let index = self.get_index_reader();
        let items = self.get_items_reader();
        let allowed_items = self.get_allowed_items_reader();

        f.debug_struct("IndexAPI")
            .field("index_len", &index.len())
            .field("items_len", &items.len())
            .field("allowed_items_cardinality", &allowed_items.cardinality())
            .finish()
    }
}
