pub(crate) mod collection;

use super::jwt::*;
use super::pager::Pager;
use anyhow::Result;
use fts_index::protocol::Document;
use fts_index::protocol::DocumentSet;
use log::*;
use model::collection::CollectionRecord;
use model::meta::tenant::TenantCollectionMeta;
use model::meta::CollectionMeta;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use sqlx::mysql::MySqlPool;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::mpsc;
use tokio::task::JoinSet;
use tokio::time::Instant;

#[derive(Debug, Default, Serialize)]
pub struct SourceReport {
    pub successes: i64,
    pub faileds: i64,
    pub tenant_faileds: Vec<String>,
    pub cost: std::time::Duration,
}

#[derive(Debug, Deserialize)]
#[serde(rename_all = "snake_case")]
pub struct PayloadIndexInfo {
    pub data_type: PayloadSchemaType,
}

#[derive(Debug, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum PayloadSchemaType {
    Keyword,
    Integer,
    Float,
    Geo,
    Text,
    Bool,
    Datetime,
    Uuid,
}

#[allow(async_fn_in_trait)]
pub trait Source {
    type Item;
    async fn send(
        &mut self,
        sender: mpsc::Sender<Self::Item>,
        start_version: u64,
        end_versin: u64,
    ) -> Result<SourceReport>;
}

pub struct CollectionDocumentSource {
    db: Arc<MySqlPool>,
    receiver: mpsc::Receiver<CollectionMeta>,
    whole_deleted: bool,
}

impl CollectionDocumentSource {
    pub fn new(
        db: Arc<MySqlPool>,
        receiver: mpsc::Receiver<CollectionMeta>,
        whole_deleted: bool,
    ) -> Self {
        Self {
            db,
            receiver,
            whole_deleted,
        }
    }

    // todo get db by tenant token
    pub async fn db_from_tenant(&self, tenant_token: &str) -> Result<Arc<MySqlPool>> {
        get_db_connection_by_tenant_token(self.db.clone(), tenant_token).await
    }

    async fn fetch_documents(
        db: Arc<MySqlPool>,
        collection: CollectionMeta,
        sender: Arc<mpsc::Sender<DocumentSet>>,
        start_version: u64,
        end_version: u64,
    ) -> Result<i64> {
        let ctx = context::new();
        let pagers = Pager::documents_by_version(
            &db,
            &collection.collection_name,
            5000,
            start_version,
            end_version,
        )
        .await?;
        let totals = pagers.iter().map(|pager| pager.page_size).sum::<i64>();
        info!(
            "Reading {} docs from [{}], version: {} -> {}",
            totals, collection.collection_name, start_version, end_version
        );
        if totals == 0 {
            return Ok(0);
        }

        let payload_schema = TenantCollectionMeta::get_by_collection_name(
            ctx.clone(),
            &db,
            &collection.collection_name,
            false,
        )
        .await?;
        let payload_schema = match payload_schema {
            Some(p) => p,
            None => {
                warn!(
                    "Collection [{}] not found with [{}]",
                    &collection.collection_name, &collection.tenant_token
                );
                return Ok(0);
            }
        };
        let payload_schema = payload_schema.payload_schema.unwrap_or_else(|| {
            panic!(
                "Payload schema not found in collection {}",
                &collection.collection_name
            )
        });
        let payload_schema: HashMap<String, PayloadIndexInfo> =
            serde_json::from_str(&payload_schema).unwrap_or_else(|e| {
                panic!(
                    "Failed to parse payload schema in collection {}: {:?}",
                    &collection.collection_name, e
                )
            });
        let mut fields = Vec::new();
        for (k, v) in payload_schema {
            if let PayloadSchemaType::Text = v.data_type {
                fields.push(k);
            }
        }
        if fields.is_empty() {
            warn!(
                "No text fields found in payload schema of collection {}",
                &collection.collection_name
            );
            return Ok(0);
        }

        for pager in pagers {
            info!(
                "Reading page {} docs from [{}]",
                pager.page_number, collection.collection_name
            );
            let documents = CollectionRecord::updated_between_all(
                &db,
                &collection.collection_name,
                pager.start_version,
                pager.end_version,
            )
            .await?;
            if documents.is_empty() {
                continue;
            }

            for field_name in fields.iter() {
                let documents: Vec<Document> = documents
                    .clone()
                    .into_iter()
                    .map(|r| Document {
                        doc_id: r.id,
                        content: get_content_from_payload(r.payload, field_name),
                        version: r.version.unwrap_or(0),
                    })
                    .collect();

                match sender
                    .send(DocumentSet {
                        tenant_token: collection.tenant_token.clone(),
                        documents: documents.clone(),
                        collection_id: collection.collection_id,
                        field_name: field_name.clone(),
                    })
                    .await
                {
                    Ok(_) => {}
                    Err(e) => {
                        warn!("Failed to send documents error: {:?}", e);
                    }
                };
            }
        }

        Ok(totals)
    }

    async fn fetch_documents_deleted(
        db: Arc<MySqlPool>,
        collection: CollectionMeta,
        sender: Arc<mpsc::Sender<DocumentSet>>,
        start_version: u64,
        end_version: u64,
        whole_deleted: bool,
    ) -> Result<i64> {
        let start_version = if whole_deleted { 0 } else { start_version };

        let documents = CollectionRecord::deleted_between_all(
            &db,
            &collection.collection_name,
            start_version,
            end_version,
        )
        .await?;

        let totals = documents.len() as i64;
        if totals == 0 {
            return Ok(0);
        }

        let documents = documents
            .into_iter()
            .map(|r| Document {
                doc_id: r.id,
                content: None,
                version: r.version.unwrap_or(0),
            })
            .collect();

        let doc_set = DocumentSet {
            tenant_token: collection.tenant_token.clone(),
            documents,
            collection_id: collection.collection_id,
            field_name: "".to_string(),
        };
        match sender.send(doc_set).await {
            Ok(_) => {}
            Err(e) => {
                warn!("Failed to send documents error: {:?}", e);
            }
        };

        Ok(totals)
    }
}

impl Source for CollectionDocumentSource {
    type Item = DocumentSet;
    async fn send(
        &mut self,
        sender: mpsc::Sender<Self::Item>,
        start_version: u64,
        end_version: u64,
    ) -> Result<SourceReport> {
        let instant = Instant::now();
        let mut report = SourceReport::default();
        let sender = Arc::new(sender);
        let mut set = JoinSet::<Result<i64>>::new();
        while let Some(collection) = self.receiver.recv().await {
            let db = match self.db_from_tenant(&collection.tenant_token).await {
                Ok(db) => db,
                Err(e) => {
                    // report.tenant_faileds.push(collection.tenant_token.clone());
                    let msg = format!(
                        "Failed to get db from tenant: {}, error: {:?}",
                        collection.tenant_token, e
                    );
                    error!("{}", msg);
                    return Err(anyhow::anyhow!(msg));
                }
            };
            set.spawn(Self::fetch_documents(
                db.clone(),
                collection.clone(),
                sender.clone(),
                start_version,
                end_version,
            ));
            set.spawn(Self::fetch_documents_deleted(
                db,
                collection,
                sender.clone(),
                start_version,
                end_version,
                self.whole_deleted,
            ));
        }

        while let Some(res) = set.join_next().await {
            match res {
                Ok(t) => {
                    if let Ok(t) = t {
                        report.successes += t;
                    } else {
                        warn!("Task failed with error: {:?}", t);
                    }
                }
                Err(e) => {
                    warn!("Task failed with error: {:?}", e);
                }
            }
        }

        debug!("All tasks of document source is done.");

        drop(sender);

        report.cost = instant.elapsed();
        Ok(report)
    }
}

fn get_content_from_payload(payload: Option<String>, field_name: &str) -> Option<String> {
    serde_json::from_str(payload.as_ref()?)
        .ok()
        .and_then(|v: Value| {
            v.get(field_name).and_then(|c| match c {
                Value::String(s) => {
                    if s.is_empty() {
                        None
                    } else {
                        Some(s.clone())
                    }
                }
                _ => {
                    warn!("Content is not string: {:?}", c);
                    None
                }
            })
        })
}

// #[cfg(test)]
// mod tests {
//     use sqlx::mysql::MySqlPoolOptions;

//     use super::*;
//     use std::env;

//     #[tokio::test]
//     async fn test_db_from_tenant_localhost() {
//         env::set_var("LOCALHOST", "true");
//         let db_url = "mysql://abc.root:xxxx@localhost:4000/ai";
//         let db = MySqlPoolOptions::new()
//             .connect(db_url)
//             .await
//             .expect("Failed to connect to database");
//         let db = Arc::new(db);
//         let (_, receiver) = mpsc::channel(1);
//         let source = CollectionDocumentSource::new(db.clone(), receiver, false);

//         let result = source.db_from_tenant("some_tenant_token").await;
//         assert!(result.is_ok());
//     }

//     #[tokio::test]
//     async fn test_db_from_tenant_remote() {
//         env::set_var("LOCALHOST", "false");
//         let (_, receiver) = mpsc::channel(1);
//         let db_url = "mysql://abc.root:xxxx@localhost:4000/ai";
//         let db = MySqlPoolOptions::new()
//             .connect(db_url)
//             .await
//             .expect("Failed to connect to database");
//         let source = CollectionDocumentSource::new(Arc::new(db.clone()), receiver, false);

//         let tenant_token = "xxxx";

//         let _ = source.db_from_tenant(tenant_token).await.unwrap();
//         // assert!(result.is_ok());
//     }
// }