use std::collections::HashSet;
use std::sync::Arc;

use crate::pipeline::jwt::get_db_connection_by_tenant_token;
use crate::pipeline::source::collection::CollectionSource;
use crate::pipeline::source::Source;
use anyhow::Result;
use lazy_static::lazy_static;
use log::*;
use model::collection::CollectionRecord;
use model::meta::tenant::{IndexType, PayloadIndex};
use model::meta::{tenant::TenantCollectionMeta, CollectionMeta};
use sqlx::MySqlPool;
use tokio::{join, spawn};

lazy_static! {
    static ref ALTERS_SLEEP_COUNTER_GAP: u64 = std::env::var("ALTERS_SLEEP_COUNTER_GAP")
        .unwrap_or("100".to_string())
        .parse::<u64>()
        .unwrap_or(100);
    static ref ALTERS_SLEEP_SECONDS: u64 = std::env::var("ALTERS_SLEEP_SECONDS")
        .unwrap_or("30".to_string())
        .parse::<u64>()
        .unwrap_or(30);
    static ref ALTERS_THREADS: usize = std::env::var("ALTERS_THREADS")
        .unwrap_or("10".to_string())
        .parse::<usize>()
        .unwrap_or(10);
}

pub(crate) async fn alter_payload_schema(
    dsn: String,
    tenant: Option<String>,
    collection: Option<String>,
) -> Result<()> {
    let ctx = context::new();
    info!(
        "Alter payload schema: tenant: {:?}, collection: {:?}",
        tenant, collection
    );
    let meta_db = Arc::new(MySqlPool::connect(&dsn).await?);

    // single tenant
    if let Some(tenant) = tenant {
        let tenant_db = get_db_connection_by_tenant_token(meta_db.clone(), &tenant).await?;
        let collections = match collection {
            Some(collection) => vec![collection],
            None => {
                let collections = CollectionMeta::all(ctx, &tenant, &meta_db).await?;
                collections.into_iter().map(|c| c.collection_name).collect()
            }
        };
        let mut tx = tenant_db.begin().await?;

        info!("Alter payload schema for tenant: {}", tenant);
        TenantCollectionMeta::add_field_payload_schema(&mut tx).await;

        for collection in collections {
            info!("Drop columns[meta,content] for collection: {}", collection);
            CollectionRecord::drop_columns(&mut tx, &collection).await;
        }
        tx.commit().await?;
        return Ok(());
    }

    let mut collection_source = CollectionSource::new(meta_db.clone());
    let (sender, mut receiver) = tokio::sync::mpsc::channel(100);
    let handler = collection_source.send(sender, 0, u64::MAX);
    let alters_handler = spawn(async move {
        let mut handled_tenant = HashSet::new();
        let mut counter = 1;
        while let Some(collection) = receiver.recv().await {
            let tenant_db =
                match get_db_connection_by_tenant_token(meta_db.clone(), &collection.tenant_token)
                    .await
                {
                    Ok(db) => db,
                    Err(e) => {
                        error!("Failed to get db connection: {:?}", e);
                        continue;
                    }
                };
            let mut tx = match tenant_db.begin().await {
                Ok(tx) => tx,
                Err(e) => {
                    error!("Failed to start transaction: {:?}", e);
                    continue;
                }
            };

            if !handled_tenant.contains(&collection.tenant_token) {
                info!(
                    "Alter payload schema for tenant: {}",
                    collection.tenant_token
                );
                TenantCollectionMeta::add_field_payload_schema(&mut tx).await;
                handled_tenant.insert(collection.tenant_token.clone());
            }

            info!(
                "Drop columns[meta,content] for collection: {}",
                collection.collection_name
            );
            CollectionRecord::drop_columns(&mut tx, &collection.collection_name).await;
            if let Err(e) = tx.commit().await {
                error!("Failed to commit transaction: {:?}", e);
            }
            counter += 1;

            if counter % *ALTERS_SLEEP_COUNTER_GAP == 0 {
                info!(
                    "Handled {} collections, sleeping {} seconds",
                    counter, *ALTERS_SLEEP_SECONDS
                );
                tokio::time::sleep(tokio::time::Duration::from_secs(*ALTERS_SLEEP_SECONDS)).await;
            }
        }
    });

    let (sr, _) = join!(handler, alters_handler);
    if let Err(e) = sr {
        error!("{:?}", e);
    }

    Ok(())
}

pub(crate) async fn alter_payload_index(
    dsn: String,
    tenant: Option<String>,
    collection: Option<String>,
) -> Result<()> {
    let ctx = context::new();
    info!(
        "Alter payload index: tenant: {:?}, collection: {:?}",
        tenant, collection
    );
    let meta_db = Arc::new(MySqlPool::connect(&dsn).await?);

    // single tenant
    if let Some(tenant) = tenant {
        match alter_payload_index_for_collection(ctx.clone(), &tenant, meta_db.clone()).await {
            Ok(_) => {
                info!("Altered payload index for tenant: {}", tenant,);
                return Ok(());
            }
            Err(e) => {
                error!(
                    "Failed to alter payload index for tenant {}: {:?}",
                    tenant, e
                );
                return Err(e);
            }
        };
    }

    let mut collection_source = CollectionSource::new(meta_db.clone());
    let (sender, mut receiver) = tokio::sync::mpsc::channel(100);
    let handler = collection_source.send(sender, 0, u64::MAX);
    let alters_handler = spawn(async move {
        let mut counter = 1;
        let mut handled_tenant = HashSet::new();
        while let Some(collection) = receiver.recv().await {
            if handled_tenant.contains(&collection.tenant_token) {
                continue;
            }

            match alter_payload_index_for_collection(
                ctx.clone(),
                &collection.tenant_token,
                meta_db.clone(),
            )
            .await
            {
                Ok(_) => {
                    info!(
                        "[{}] Altered payload index for tenant: {}",
                        counter, collection.tenant_token
                    );
                }
                Err(e) => {
                    error!(
                        "Failed to alter payload index for tenant {}: {:?}",
                        collection.tenant_token, e
                    );
                }
            };
            handled_tenant.insert(collection.tenant_token.clone());

            counter += 1;
            if counter % *ALTERS_SLEEP_COUNTER_GAP == 0 {
                info!(
                    "Handled {} collections, sleeping {} seconds",
                    counter, *ALTERS_SLEEP_SECONDS
                );
                tokio::time::sleep(tokio::time::Duration::from_secs(*ALTERS_SLEEP_SECONDS)).await;
            }
        }
    });

    let (sr, _) = join!(handler, alters_handler);
    if let Err(e) = sr {
        error!("{:?}", e);
    }

    Ok(())
}

const PAYLOAD_INDEX: &str = r#"
{"group_id": {"data_type": "keyword", "points": 0}, "metadata.doc_id": {"data_type": "keyword", "points": 0}, "page_content": {"data_type": "text", "params": {"lowercase": true, "max_token_len": 20, "min_token_len": 2, "tokenizer": "multilingual", "type": "text"}, "points": 0}}
"#;

async fn alter_payload_index_for_collection(
    ctx: Arc<context::Context>,
    tenant: &str,
    meta_db: Arc<MySqlPool>,
) -> Result<()> {
    info!("Alter payload index for tenant: {}", tenant);
    let tenant_db = get_db_connection_by_tenant_token(meta_db.clone(), tenant).await?;
    let collections = CollectionMeta::all(context::new(), tenant, &meta_db).await?;

    for collection in collections {
        info!(
            "Alter page_content for collection: {}",
            collection.collection_name
        );
        if let Err(e) = CollectionRecord::add_tokens_column(
            ctx.clone(),
            &tenant_db,
            &collection.collection_name,
            "page_content",
        )
        .await
        {
            error!(
                "Failed to alter page_content for collection {}: {:?}",
                collection.collection_name, e
            );
        };

        info!(
            "Alter group_id for collection: {}",
            collection.collection_name
        );
        if let Err(e) = TenantCollectionMeta::add_payload_idx(
            ctx.clone(),
            &tenant_db,
            &collection.collection_name,
            PayloadIndex {
                path: "group_id".to_string(),
                length: Some(255),
                index_type: IndexType::Varchar,
            },
        )
        .await
        {
            error!(
                "Failed to alter group_id for collection {}: {:?}",
                collection.collection_name, e
            );
        };

        info!(
            "Alter metadata.doc_id for collection: {}",
            collection.collection_name
        );
        if let Err(e) = TenantCollectionMeta::add_payload_idx(
            ctx.clone(),
            &tenant_db,
            &collection.collection_name,
            PayloadIndex {
                path: "metadata.doc_id".to_string(),
                length: Some(255),
                index_type: IndexType::Varchar,
            },
        )
        .await
        {
            error!(
                "Failed to alter metadata.doc_id for collection {}: {:?}",
                collection.collection_name, e
            );
        };

        info!(
            "Update payload_index for collection: {}",
            collection.collection_name
        );
        if let Err(e) = TenantCollectionMeta::update_payload_index(
            ctx.clone(),
            &tenant_db,
            &collection.collection_name,
            PAYLOAD_INDEX.to_string(),
        )
        .await
        {
            error!(
                "Failed to update payload_index for collection {}: {:?}",
                collection.collection_name, e
            );
        };
    }

    Ok(())
}

pub(crate) async fn move_tokens(
    dsn: String,
    tenant: Option<String>,
    collection: Option<String>,
) -> Result<()> {
    let ctx = context::new();
    info!(
        "Ctl move tokens: tenant: {:?}, collection: {:?}",
        tenant, collection
    );
    let meta_db = Arc::new(MySqlPool::connect(&dsn).await?);

    // single tenant
    if let Some(tenant) = tenant {
        let tenant_db = get_db_connection_by_tenant_token(meta_db.clone(), &tenant).await?;
        let collections = match collection {
            Some(collection) => vec![collection],
            None => {
                let collections = CollectionMeta::all(ctx.clone(), &tenant, &meta_db).await?;
                collections.into_iter().map(|c| c.collection_name).collect()
            }
        };

        info!("Moving tokens for tenant: {}", tenant);
        for collection in collections {
            CollectionRecord::move_tokens(ctx.clone(), &tenant_db, &collection).await?;
        }
        return Ok(());
    }

    let mut collection_source = CollectionSource::new(meta_db.clone());
    let (sender, mut receiver) = tokio::sync::mpsc::channel(100);
    let handler = collection_source.send(sender, 0, u64::MAX);
    let alters_handler = spawn(async move {
        let mut handled_tenant = HashSet::new();
        let mut counter = 1;
        while let Some(collection) = receiver.recv().await {
            if handled_tenant.contains(&collection.tenant_token) {
                continue;
            }

            let _ = move_job(
                ctx.clone(),
                meta_db.clone(),
                collection.tenant_token.clone(),
            )
            .await;
            handled_tenant.insert(collection.tenant_token.clone());

            counter += 1;
            if counter % *ALTERS_SLEEP_COUNTER_GAP == 0 {
                info!(
                    "Handled {} collections, sleeping {} seconds",
                    counter, *ALTERS_SLEEP_SECONDS
                );
                tokio::time::sleep(tokio::time::Duration::from_secs(*ALTERS_SLEEP_SECONDS)).await;
            }
        }
    });

    let (sr, _) = join!(handler, alters_handler);
    if let Err(e) = sr {
        error!("{:?}", e);
    }

    Ok(())
}

async fn move_job(
    ctx: Arc<context::Context>,
    meta_db: Arc<MySqlPool>,
    tenant: String,
) -> Result<()> {
    let tenant_db = get_db_connection_by_tenant_token(meta_db.clone(), &tenant).await?;
    let collections = CollectionMeta::all(ctx.clone(), &tenant, &meta_db).await?;

    for collection in collections {
        match CollectionRecord::move_tokens(ctx.clone(), &tenant_db, &collection.collection_name)
            .await
        {
            Ok(_) => {
                info!(
                    "Moved tokens for tenant: {}, collection: {}",
                    collection.tenant_token, collection.collection_name
                );
            }
            Err(e) => {
                error!(
                    "Failed to move tokens for tenant: {}, collection: {}: {:?}",
                    collection.tenant_token, collection.collection_name, e
                );
            }
        };
    }
    Ok(())
}