use std::collections::HashMap;

use axum::Json;
use model::collection::CollectionRecord;
use model::meta::CollectionMeta;
use tokio::time::Instant;

use crate::collection::actions::QdrantMeta;
use crate::error::QdrantError;
use crate::error::Result;

use crate::protocol::{ApiStatus, UsageResponse};
use crate::server::ExtractApiKey;
use crate::{metrics, state::ApiState};
use axum::extract::State;
use log::*;

const VECTOR_DATA_SIZE_BYTES_EACH_DIM: u64 = 4;
const TIKV_ROW_OVERHEAD_BYTES: u64 = (2764.0 * 0.8) as u64;
const COLUMNAR_ROW_OVERHEAD_BYTES: u64 = 1228;

pub(crate) async fn cluster(
    State(state): State<ApiState>,
    ExtractApiKey(api_key): ExtractApiKey,
) -> Result<Json<UsageResponse<HashMap<String, HashMap<String, u64>>>>> {
    let ctx = context::with_metrics(metrics::INTERFACE_PERFORMANCE_HISTOGRAM.clone());
    info!("[{}] GET /cluster", ctx.id());
    metrics::ROUTER_COLLECTION_WITH_PATH
        .with_label_values(&["", "/cluster", "get"])
        .inc();
    let (meta_db, tenant_db, tenant_token) = state.get_meta_and_tenant_db(&ctx, &api_key).await?;
    let meta_db = meta_db.ok_or(QdrantError::InternalServerErrorWithMsg(
        "meta-db not configured".to_string(),
    ))?;

    let start = Instant::now();
    // query all collections
    let mut collections =
        match CollectionMeta::get_by_tenant(ctx.clone(), &meta_db, &tenant_token).await? {
            None => {
                let usage = HashMap::from([("row_based".to_string(), 0), ("columnar".to_string(), 0)]);
                let stoarge = HashMap::from([("storage".to_string(), usage)]);

                return Ok(Json(UsageResponse {
                    time: start.elapsed().as_secs_f64() as f32,
                    status: ApiStatus::Ok,
                    usage: Some(stoarge),
                }));
            }
            Some(cols) => cols,
        };

    // if data_count is -1, we need to query from tenant db
    let mut row_based_storage_size: u64 = 0;
    let mut columnar_storage_size: u64 = 0;
    for collection in collections.iter_mut() {
        if collection.data_count.is_none() || collection.data_count <= Some(-1) {
            let count =
                CollectionRecord::count(ctx.clone(), &tenant_db, &collection.collection_name)
                    .await?;
            collection.data_count = Some(count as i64);

            // update data_count in collection with async way
            let tenant_token_clone = tenant_token.to_string();
            let collection_name_clone = collection.collection_name.clone();
            let meta_db_clone = meta_db.clone();
            tokio::spawn(async move {
                if let Err(e) = CollectionMeta::update_data_count(
                    &meta_db_clone,
                    &tenant_token_clone,
                    &collection_name_clone,
                    (count as i64) + 1,
                )
                .await
                {
                    error!(
                        "Failed to update data_count for collection {}: {}",
                        &collection_name_clone, e
                    );
                }
            });
        }
        let collection_meta: QdrantMeta =
            serde_json::from_str(collection.meta.as_ref().unwrap_or(&"{}".to_string()))?;
        let vector_dim_size = if collection_meta.qdrant.vectors.size == 0 {
            3072
        } else {
            collection_meta.qdrant.vectors.size
        };

        let rows = collection.data_count.unwrap_or(0) as u64;
        row_based_storage_size += vector_dim_size * VECTOR_DATA_SIZE_BYTES_EACH_DIM * rows
            + TIKV_ROW_OVERHEAD_BYTES * rows;
        columnar_storage_size += vector_dim_size * VECTOR_DATA_SIZE_BYTES_EACH_DIM * rows
            + COLUMNAR_ROW_OVERHEAD_BYTES * rows;
    }

    let usage = HashMap::from([
        ("row_based".to_string(), row_based_storage_size),
        ("columnar".to_string(), columnar_storage_size),
    ]);

    let stoarge = HashMap::from([("storage".to_string(), usage)]);
    Ok(Json(UsageResponse {
        time: start.elapsed().as_secs_f64() as f32,
        status: ApiStatus::Ok,
        usage: Some(stoarge),
    }))
}