use crate::{
    exchange::Client,
    types::{ExchangeResult, Operation, OperationResult},
    DebugInfo, Exchange, ExchangeFactory, GraphQLQuery, OperationMeta, OperationType, QueryError,
    RequestPolicy, Response, ResultSource
};
use std::{
    any::Any,
    collections::{HashMap, HashSet},
    sync::{Arc, Mutex}
};

type ResultCache = Arc<Mutex<HashMap<u64, Box<dyn Any + Send>>>>;
type OperationCache = Arc<Mutex<HashMap<&'static str, HashSet<u64>>>>;

/// The default caching exchange.
///
/// This does eager invalidation, invalidating all queries related to all types that are returned by a mutation.
/// It doesn't do any normalization, operations are stored by their key (query x variables).
/// Mutations are never cached.
pub struct CacheExchange;
impl<TNext: Exchange> ExchangeFactory<TNext> for CacheExchange {
    type Output = CacheExchangeImpl<TNext>;

    fn build(self, next: TNext) -> Self::Output {
        CacheExchangeImpl {
            result_cache: Arc::new(Mutex::new(HashMap::new())),
            operation_cache: Arc::new(Mutex::new(HashMap::new())),

            next
        }
    }
}

pub struct CacheExchangeImpl<TNext: Exchange> {
    result_cache: ResultCache,
    operation_cache: OperationCache,

    next: TNext
}

#[inline]
fn should_skip<Q: GraphQLQuery>(operation: &Operation<Q::Variables>) -> bool {
    let operation_type = &operation.meta.operation_type;
    operation_type != &OperationType::Query && operation_type != &OperationType::Mutation
}

impl<TNext: Exchange> CacheExchangeImpl<TNext> {
    fn is_operation_cached<Q: GraphQLQuery>(&self, operation: &Operation<Q::Variables>) -> bool {
        let OperationMeta { operation_type, .. } = &operation.meta;
        let key = operation.key;
        let request_policy = &operation.options.request_policy;

        operation_type == &OperationType::Query
            && request_policy != &RequestPolicy::NetworkOnly
            && (request_policy == &RequestPolicy::CacheOnly
                || self.result_cache.lock().unwrap().contains_key(&key))
    }

    fn after_query<Q: GraphQLQuery>(
        &self,
        operation_result: OperationResult<Q::ResponseData>
    ) -> Result<OperationResult<Q::ResponseData>, QueryError> {
        if operation_result.response.data.is_none() {
            return Ok(operation_result);
        }

        let OperationMeta { involved_types, .. } = &operation_result.meta;
        let key = operation_result.key;

        {
            let mut result_cache = self.result_cache.lock().unwrap();
            let data = operation_result.response.data.as_ref().unwrap().clone();
            result_cache.insert(key, Box::new(data));
        }
        {
            let mut operation_cache = self.operation_cache.lock().unwrap();
            for involved_type in involved_types {
                operation_cache
                    .entry(*involved_type)
                    .and_modify(|entry| {
                        entry.insert(key);
                    })
                    .or_insert_with(|| {
                        let mut set = HashSet::with_capacity(1);
                        set.insert(key);
                        set
                    });
            }
        }

        Ok(operation_result)
    }

    fn after_mutation<Q: GraphQLQuery, C: Client>(
        &self,
        operation_result: OperationResult<Q::ResponseData>,
        client: C
    ) -> Result<OperationResult<Q::ResponseData>, QueryError> {
        if operation_result.response.data.is_none() {
            return Ok(operation_result);
        }

        let OperationMeta { involved_types, .. } = &operation_result.meta;
        let key = operation_result.key;

        let ops_to_remove: HashSet<u64> = {
            let cache = self.operation_cache.lock().unwrap();
            let mut ops = HashSet::new();
            for involved_type in involved_types {
                let ops_for_type = cache.get(involved_type);
                if let Some(ops_for_type) = ops_for_type {
                    ops.extend(ops_for_type)
                }
            }
            ops.insert(key);
            ops
        };
        {
            let mut cache = self.result_cache.lock().unwrap();
            for op in ops_to_remove.iter() {
                cache.remove(op);
            }
        }
        for op in ops_to_remove {
            client.rerun_query(op);
        }
        Ok(operation_result)
    }
}

#[async_trait]
impl<TNext: Exchange> Exchange for CacheExchangeImpl<TNext> {
    async fn run<Q: GraphQLQuery, C: Client>(
        &self,
        operation: Operation<Q::Variables>,
        client: C
    ) -> ExchangeResult<Q::ResponseData> {
        if should_skip::<Q>(&operation) {
            return self.next.run::<Q, _>(operation, client).await;
        }

        if !self.is_operation_cached::<Q>(&operation) {
            let res = self.next.run::<Q, _>(operation, client.clone()).await?;

            match res.meta.operation_type {
                OperationType::Query => self.after_query::<Q>(res),
                OperationType::Mutation => self.after_mutation::<Q, _>(res, client),
                _ => Ok(res)
            }
        } else {
            let key = &operation.key;

            let cached_result = {
                let cache = self.result_cache.lock().unwrap();
                cache.get(key).map(|res| {
                    let res: &Q::ResponseData = (&**res).downcast_ref::<Q::ResponseData>().unwrap();
                    res.clone()
                })
            };

            if let Some(cached) = cached_result {
                let result = OperationResult {
                    key: operation.key,
                    meta: operation.meta,
                    response: Response {
                        debug_info: Some(DebugInfo {
                            source: ResultSource::Cache,
                            did_dedup: false
                        }),
                        data: Some(cached),
                        errors: None
                    }
                };
                Ok(result)
            } else {
                self.next.run::<Q, _>(operation, client).await
            }
        }
    }
}
