mod alters;
mod fts;
mod index;
mod pipeline;

use crate::alters::*;
use crate::fts::*;
use anyhow::Result;
use clap::{Parser, Subcommand};

#[derive(Parser)]
#[command(
    name = "ctl",
    version,
    arg_required_else_help(true),
    about = "Command-line tool for Qdrant on TiDB services"
)]
struct Cli {
    #[command(subcommand)]
    command: Commands,
}

#[derive(Subcommand)]
enum Commands {
    CreateIndex {
        index_name: String,
    },
    Indexes {
        base_url: Option<String>,
    },
    Versions {
        #[arg(
            id = "base-url",
            long = "base-url",
            env = "BASE_URL",
            help = "base url of the service",
            num_args(1)
        )]
        base_url: Option<String>,
        index_name: String,
    },
    Search {
        #[arg(
            id = "base-url",
            long = "base-url",
            env = "BASE_URL",
            help = "base url of the service",
            num_args(1)
        )]
        base_url: Option<String>,
        tenant_token: String,
        collection_id: u64,
        index_name: String,
        query: String,
    },
    ReloadIndex {
        #[arg(
            id = "base-url",
            long = "base-url",
            env = "BASE_URL",
            help = "base url of the service",
            num_args(1)
        )]
        base_url: Option<String>,
    },
    UpdateIndex {
        #[arg(
            id = "base-url",
            long = "base-url",
            env = "BASE_URL",
            help = "base url of the service",
            num_args(1)
        )]
        base_url: Option<String>,
        version: u64,
    },
    IndexUpdatedCollections {
        #[arg(
            id = "start-version",
            long = "start-version",
            help = "Load documents newer than or equals to the specified end version(inclusive)",
            num_args(1)
        )]
        start_version: Option<u64>,
        #[arg(
            id = "end-version",
            long = "end-version",
            help = "Load documents newer than or equals to the specified end version(inclusive)",
            num_args(1)
        )]
        end_version: Option<u64>,
        #[arg(
            id = "memory-budget",
            long = "memory-budget",
            env = "MEMORY_BUDGET",
            help = "Memory budget when building index",
            num_args(1)
        )]
        memory_budget: Option<bytesize::ByteSize>,
        #[arg(
            id = "c",
            long = "concurrency",
            default_value = "5",
            help = "concurrency thread number, default 5",
            num_args(1)
        )]
        concurrency: usize,
        #[arg(id = "dsn", long = "dsn", env = "DSN", help = "TiDB DSN", num_args(1))]
        dsn: String,
        storage_path: String,
        index_name: String,
        #[arg(
            id = "whole-deleted",
            long = "whole-deleted",
            default_value = "false",
            help = "with all deleted documents",
            num_args(1)
        )]
        whole_deleted: String,
        #[arg(
            id = "compact",
            long = "compact",
            default_value = "false",
            help = "compact index segments and upload",
            num_args(1)
        )]
        compact: String,
    },
    CompactIndex {
        storage_path: String,
        index_name: String,
    },
    AlterPayloadSchema {
        #[arg(id = "dsn", long = "dsn", env = "DSN", help = "TiDB DSN", num_args(1))]
        dsn: String,
        #[arg(
            id = "tenant",
            long = "tenant",
            help = "specific tenant, if none, alter all tenants",
            num_args(1)
        )]
        tenant: Option<String>,
        #[arg(
            id = "collection",
            long = "collection",
            help = "specific collection, if none, alter all collections",
            num_args(1)
        )]
        collection: Option<String>,
    },
    AlterPayloadIndex {
        #[arg(id = "dsn", long = "dsn", env = "DSN", help = "TiDB DSN", num_args(1))]
        dsn: String,
        #[arg(
            id = "tenant",
            long = "tenant",
            help = "specific tenant, if none, alter all tenants",
            num_args(1)
        )]
        tenant: Option<String>,
        #[arg(
            id = "collection",
            long = "collection",
            help = "specific collection, if none, alter all collections",
            num_args(1)
        )]
        collection: Option<String>,
    },
    MoveTokens {
        #[arg(id = "dsn", long = "dsn", env = "DSN", help = "TiDB DSN", num_args(1))]
        dsn: String,
        #[arg(
            id = "tenant",
            long = "tenant",
            help = "specific tenant, if none, alter all tenants",
            num_args(1)
        )]
        tenant: Option<String>,
        #[arg(
            id = "collection",
            long = "collection",
            help = "specific collection, if none, alter all collections",
            num_args(1)
        )]
        collection: Option<String>,
    },
}

#[tokio::main]
#[allow(clippy::needless_return)]
async fn main() -> Result<()> {
    env_logger::init();
    let args = Cli::parse();
    match args.command {
        Commands::CreateIndex { index_name } => {
            index::create_index(&index_name).await?;
        }
        Commands::Indexes { base_url } => {
            indexes(base_url).await?;
        }
        Commands::Versions {
            base_url,
            index_name,
        } => {
            versions(base_url, index_name).await?;
        }
        Commands::Search {
            tenant_token,
            collection_id,
            index_name,
            base_url,
            query,
        } => {
            search(base_url, tenant_token, collection_id, index_name, query).await?;
        }
        Commands::ReloadIndex { base_url } => {
            reload_index(base_url).await?;
        }
        Commands::UpdateIndex { base_url, version } => {
            update_index(base_url, version).await?;
        }
        Commands::IndexUpdatedCollections {
            start_version,
            end_version,
            memory_budget,
            dsn,
            storage_path,
            index_name,
            concurrency,
            whole_deleted,
            compact,
        } => {
            match pipeline::index_updated_collections(
                dsn,
                (start_version, end_version),
                memory_budget.map(|b| b.0 as usize),
                storage_path.clone(),
                index_name.clone(),
                concurrency,
                whole_deleted == "true",
            )
            .await
            {
                Err(e) => {
                    let msg = format!("Failed to update index to FTS service: {:?}", e);
                    notice::feishu::critical(&msg, Some("ctl:IndexUpdatedCollections")).await;
                    log::error!("{}", msg);
                }
                Ok(version) => {
                    if version == 0 {
                        return Ok(());
                    }

                    if compact == "true" || version % REMAINS_META_VERSION_COUNT == 0 {
                        if let Err(e) =
                            compact_index(storage_path + "/compacting", index_name).await
                        {
                            let msg =
                                format!("Failed to compact version({}) index: {:?}", version, e);
                            notice::feishu::critical(&msg, Some("ctl:CompactIndex")).await;
                            log::error!("{}", msg);
                        };
                    }
                }
            }
        }
        Commands::CompactIndex {
            storage_path,
            index_name,
        } => {
            compact_index(storage_path, index_name).await?;
        }
        Commands::AlterPayloadSchema {
            dsn,
            tenant,
            collection,
        } => {
            alter_payload_schema(dsn, tenant, collection).await?;
        }
        Commands::AlterPayloadIndex {
            dsn,
            tenant,
            collection,
        } => {
            alter_payload_index(dsn, tenant, collection).await?;
        }
        Commands::MoveTokens {
            dsn,
            tenant,
            collection,
        } => {
            move_tokens(dsn, tenant, collection).await?;
        }
    };

    Ok(())
}