use meili_snap::snapshot;
use meilisearch_auth::AuthFilter;
use meilisearch_types::index_uid_pattern::IndexUidPattern;
use meilisearch_types::milli::update::MissingDocumentPolicy;
use meilisearch_types::tasks::{IndexSwap, KindWithContent, Status};
use time::{Duration, OffsetDateTime};

use crate::insta_snapshot::{snapshot_bitmap, snapshot_index_scheduler};
use crate::test_utils::Breakpoint::*;
use crate::test_utils::{
    index_creation_task, read_json, replace_document_import_task,
    replace_document_import_task_with_opts, sample_documents, FailureLocation,
};
use crate::{IndexScheduler, Query};

#[test]
fn query_batches_from_and_limit() {
    let (index_scheduler, mut handle) = IndexScheduler::test(true, vec![]);

    let kind = index_creation_task("doggo", "bone");
    let _task = index_scheduler.register(kind, None, false).unwrap();
    snapshot!(snapshot_index_scheduler(&index_scheduler), name: "registered_the_first_task");
    let kind = index_creation_task("whalo", "plankton");
    let _task = index_scheduler.register(kind, None, false).unwrap();
    snapshot!(snapshot_index_scheduler(&index_scheduler), name: "registered_the_second_task");
    let kind = index_creation_task("catto", "his_own_vomit");
    let _task = index_scheduler.register(kind, None, false).unwrap();
    snapshot!(snapshot_index_scheduler(&index_scheduler), name: "registered_the_third_task");

    handle.advance_n_successful_batches(3);
    snapshot!(snapshot_index_scheduler(&index_scheduler), name: "processed_all_tasks");

    let proc = index_scheduler.processing_tasks.read().unwrap().clone();
    let rtxn = index_scheduler.env.read_txn().unwrap();
    let query = Query { limit: Some(0), ..Default::default() };
    let (batches, _) = index_scheduler
        .queue
        .get_batch_ids_from_authorized_indexes(&rtxn, &query, &AuthFilter::default(), &proc)
        .unwrap();
    snapshot!(snapshot_bitmap(&batches), @"[]");

    let query = Query { limit: Some(1), ..Default::default() };
    let (batches, _) = index_scheduler
        .queue
        .get_batch_ids_from_authorized_indexes(&rtxn, &query, &AuthFilter::default(), &proc)
        .unwrap();
    snapshot!(snapshot_bitmap(&batches), @"[2,]");

    let query = Query { limit: Some(2), ..Default::default() };
    let (batches, _) = index_scheduler
        .queue
        .get_batch_ids_from_authorized_indexes(&rtxn, &query, &AuthFilter::default(), &proc)
        .unwrap();
    snapshot!(snapshot_bitmap(&batches), @"[1,2,]");

    let query = Query { from: Some(1), ..Default::default() };
    let (batches, _) = index_scheduler
        .queue
        .get_batch_ids_from_authorized_indexes(&rtxn, &query, &AuthFilter::default(), &proc)
        .unwrap();
    snapshot!(snapshot_bitmap(&batches), @"[0,1,]");

    let query = Query { from: Some(2), ..Default::default() };
    let (batches, _) = index_scheduler
        .queue
        .get_batch_ids_from_authorized_indexes(&rtxn, &query, &AuthFilter::default(), &proc)
        .unwrap();
    snapshot!(snapshot_bitmap(&batches), @"[0,1,2,]");

    let query = Query { from: Some(1), limit: Some(1), ..Default::default() };
    let (batches, _) = index_scheduler
        .queue
        .get_batch_ids_from_authorized_indexes(&rtxn, &query, &AuthFilter::default(), &proc)
        .unwrap();
    snapshot!(snapshot_bitmap(&batches), @"[1,]");

    let query = Query { from: Some(1), limit: Some(2), ..Default::default() };
    let (batches, _) = index_scheduler
        .queue
        .get_batch_ids_from_authorized_indexes(&rtxn, &query, &AuthFilter::default(), &proc)
        .unwrap();
    snapshot!(snapshot_bitmap(&batches), @"[0,1,]");
}

#[test]
fn query_batches_simple() {
    let start_time = OffsetDateTime::now_utc();

    let (index_scheduler, mut handle) =
        IndexScheduler::test(true, vec![(3, FailureLocation::InsideProcessBatch)]);

    let kind = index_creation_task("catto", "mouse");
    let _task = index_scheduler.register(kind, None, false).unwrap();
    let kind = index_creation_task("doggo", "sheep");
    let _task = index_scheduler.register(kind, None, false).unwrap();
    let kind = index_creation_task("whalo", "fish");
    let _task = index_scheduler.register(kind, None, false).unwrap();

    snapshot!(snapshot_index_scheduler(&index_scheduler), name: "start");

    handle.advance_till([Start, BatchCreated]);

    let query = Query { statuses: Some(vec![Status::Processing]), ..Default::default() };
    let (mut batches, _) = index_scheduler
        .get_batches_from_authorized_indexes(&query, &AuthFilter::default())
        .unwrap();
    assert_eq!(batches.len(), 1);
    batches[0].started_at = OffsetDateTime::UNIX_EPOCH;
    assert!(batches[0].enqueued_at.is_some());
    batches[0].enqueued_at = None;

    if !batches[0].stats.progress_trace.is_empty() {
        batches[0].stats.progress_trace.clear();
        batches[0]
            .stats
            .progress_trace
            .insert("processing tasks".to_string(), "deterministic_duration".into());
    }

    // Insta cannot snapshot our batches because the batch stats contains an enum as key: https://github.com/mitsuhiko/insta/issues/689
    let batch = serde_json::to_string_pretty(&batches[0]).unwrap();
    snapshot!(batch, @r###"
    {
      "uid": 0,
      "details": {
        "primaryKey": "mouse"
      },
      "stats": {
        "totalNbTasks": 1,
        "status": {
          "processing": 1
        },
        "types": {
          "indexCreation": 1
        },
        "indexUids": {
          "catto": 1
        },
        "progressTrace": {
          "processing tasks": "deterministic_duration"
        }
      },
      "startedAt": "1970-01-01T00:00:00Z",
      "finishedAt": null,
      "enqueuedAt": null,
      "stopReason": "created batch containing only task with id 0 of type `indexCreation` that cannot be batched with any other task."
    }
    "###);

    let query = Query { statuses: Some(vec![Status::Enqueued]), ..Default::default() };
    let (batches, _) = index_scheduler
        .get_batch_ids_from_authorized_indexes(&query, &AuthFilter::default())
        .unwrap();
    snapshot!(snapshot_bitmap(&batches), @"[]"); // The batches don't contains any enqueued tasks

    let query =
        Query { statuses: Some(vec![Status::Enqueued, Status::Processing]), ..Default::default() };
    let (batches, _) = index_scheduler
        .get_batch_ids_from_authorized_indexes(&query, &AuthFilter::default())
        .unwrap();
    snapshot!(snapshot_bitmap(&batches), @"[0,]"); // both enqueued and processing tasks in the first tick

    let query = Query {
        statuses: Some(vec![Status::Enqueued, Status::Processing]),
        after_started_at: Some(start_time),
        ..Default::default()
    };
    let (batches, _) = index_scheduler
        .get_batch_ids_from_authorized_indexes(&query, &AuthFilter::default())
        .unwrap();
    // both enqueued and processing tasks in the first tick, but limited to those with a started_at
    // that comes after the start of the test, which should excludes the enqueued tasks
    snapshot!(snapshot_bitmap(&batches), @"[0,]");

    let query = Query {
        statuses: Some(vec![Status::Enqueued, Status::Processing]),
        before_started_at: Some(start_time),
        ..Default::default()
    };
    let (batches, _) = index_scheduler
        .get_batch_ids_from_authorized_indexes(&query, &AuthFilter::default())
        .unwrap();
    // both enqueued and processing tasks in the first tick, but limited to those with a started_at
    // that comes before the start of the test, which should excludes all of them
    snapshot!(snapshot_bitmap(&batches), @"[]");

    let query = Query {
        statuses: Some(vec![Status::Enqueued, Status::Processing]),
        after_started_at: Some(start_time),
        before_started_at: Some(start_time + Duration::minutes(1)),
        ..Default::default()
    };
    let (batches, _) = index_scheduler
        .get_batch_ids_from_authorized_indexes(&query, &AuthFilter::default())
        .unwrap();
    // both enqueued and processing tasks in the first tick, but limited to those with a started_at
    // that comes after the start of the test and before one minute after the start of the test,
    // which should exclude the enqueued tasks and include the only processing task
    snapshot!(snapshot_bitmap(&batches), @"[0,]");

    handle.advance_till([
        InsideProcessBatch,
        InsideProcessBatch,
        ProcessBatchSucceeded,
        AfterProcessing,
        Start,
        BatchCreated,
    ]);
    snapshot!(snapshot_index_scheduler(&index_scheduler), name: "after-advancing-a-bit");

    let second_start_time = OffsetDateTime::now_utc();

    let query = Query {
        statuses: Some(vec![Status::Succeeded, Status::Processing]),
        after_started_at: Some(start_time),
        before_started_at: Some(start_time + Duration::minutes(1)),
        ..Default::default()
    };
    let (batches, _) = index_scheduler
        .get_batch_ids_from_authorized_indexes(&query, &AuthFilter::default())
        .unwrap();
    // both succeeded and processing tasks in the first tick, but limited to those with a started_at
    // that comes after the start of the test and before one minute after the start of the test,
    // which should include all tasks
    snapshot!(snapshot_bitmap(&batches), @"[0,1,]");

    let query = Query {
        statuses: Some(vec![Status::Succeeded, Status::Processing]),
        before_started_at: Some(start_time),
        ..Default::default()
    };
    let (batches, _) = index_scheduler
        .get_batch_ids_from_authorized_indexes(&query, &AuthFilter::default())
        .unwrap();
    // both succeeded and processing tasks in the first tick, but limited to those with a started_at
    // that comes before the start of the test, which should exclude all tasks
    snapshot!(snapshot_bitmap(&batches), @"[]");

    let query = Query {
        statuses: Some(vec![Status::Enqueued, Status::Succeeded, Status::Processing]),
        after_started_at: Some(second_start_time),
        before_started_at: Some(second_start_time + Duration::minutes(1)),
        ..Default::default()
    };
    let (batches, _) = index_scheduler
        .get_batch_ids_from_authorized_indexes(&query, &AuthFilter::default())
        .unwrap();
    // both succeeded and processing tasks in the first tick, but limited to those with a started_at
    // that comes after the start of the second part of the test and before one minute after the
    // second start of the test, which should exclude all tasks
    snapshot!(snapshot_bitmap(&batches), @"[]");

    // now we make one more batch, the started_at field of the new tasks will be past `second_start_time`
    handle.advance_till([
        InsideProcessBatch,
        InsideProcessBatch,
        ProcessBatchSucceeded,
        AfterProcessing,
        Start,
        BatchCreated,
    ]);

    let (batches, _) = index_scheduler
        .get_batch_ids_from_authorized_indexes(&query, &AuthFilter::default())
        .unwrap();
    // we run the same query to verify that, and indeed find that the last task is matched
    snapshot!(snapshot_bitmap(&batches), @"[2,]");

    let query = Query {
        statuses: Some(vec![Status::Enqueued, Status::Succeeded, Status::Processing]),
        after_started_at: Some(second_start_time),
        before_started_at: Some(second_start_time + Duration::minutes(1)),
        ..Default::default()
    };
    let (batches, _) = index_scheduler
        .get_batch_ids_from_authorized_indexes(&query, &AuthFilter::default())
        .unwrap();
    // enqueued, succeeded, or processing tasks started after the second part of the test, should
    // again only return the last task
    snapshot!(snapshot_bitmap(&batches), @"[2,]");

    handle.advance_till([ProcessBatchFailed, AfterProcessing]);

    // now the last task should have failed
    snapshot!(snapshot_index_scheduler(&index_scheduler), name: "end");
    let (batches, _) = index_scheduler
        .get_batch_ids_from_authorized_indexes(&query, &AuthFilter::default())
        .unwrap();
    // so running the last query should return nothing
    snapshot!(snapshot_bitmap(&batches), @"[]");

    let query = Query {
        statuses: Some(vec![Status::Failed]),
        after_started_at: Some(second_start_time),
        before_started_at: Some(second_start_time + Duration::minutes(1)),
        ..Default::default()
    };
    let (batches, _) = index_scheduler
        .get_batch_ids_from_authorized_indexes(&query, &AuthFilter::default())
        .unwrap();
    // but the same query on failed tasks should return the last task
    snapshot!(snapshot_bitmap(&batches), @"[2,]");

    let query = Query {
        statuses: Some(vec![Status::Failed]),
        after_started_at: Some(second_start_time),
        before_started_at: Some(second_start_time + Duration::minutes(1)),
        ..Default::default()
    };
    let (batches, _) = index_scheduler
        .get_batch_ids_from_authorized_indexes(&query, &AuthFilter::default())
        .unwrap();
    // but the same query on failed tasks should return the last task
    snapshot!(snapshot_bitmap(&batches), @"[2,]");

    let query = Query {
        statuses: Some(vec![Status::Failed]),
        uids: Some(vec![1]),
        after_started_at: Some(second_start_time),
        before_started_at: Some(second_start_time + Duration::minutes(1)),
        ..Default::default()
    };
    let (batches, _) = index_scheduler
        .get_batch_ids_from_authorized_indexes(&query, &AuthFilter::default())
        .unwrap();
    // same query but with an invalid uid
    snapshot!(snapshot_bitmap(&batches), @"[]");

    let query = Query {
        statuses: Some(vec![Status::Failed]),
        uids: Some(vec![2]),
        after_started_at: Some(second_start_time),
        before_started_at: Some(second_start_time + Duration::minutes(1)),
        ..Default::default()
    };
    let (batches, _) = index_scheduler
        .get_batch_ids_from_authorized_indexes(&query, &AuthFilter::default())
        .unwrap();
    // same query but with a valid uid
    snapshot!(snapshot_bitmap(&batches), @"[2,]");
}

#[test]
fn query_batches_special_rules() {
    let (index_scheduler, mut handle) =
        IndexScheduler::test(true, vec![(3, FailureLocation::InsideProcessBatch)]);

    let kind = index_creation_task("catto", "mouse");
    let _task = index_scheduler.register(kind, None, false).unwrap();
    let kind = index_creation_task("doggo", "sheep");
    let _task = index_scheduler.register(kind, None, false).unwrap();
    let kind = KindWithContent::IndexSwap {
        swaps: vec![IndexSwap { indexes: ("catto".to_owned(), "doggo".to_owned()), rename: false }],
    };
    let _task = index_scheduler.register(kind, None, false).unwrap();
    let kind = KindWithContent::IndexSwap {
        swaps: vec![IndexSwap { indexes: ("catto".to_owned(), "whalo".to_owned()), rename: false }],
    };
    let _task = index_scheduler.register(kind, None, false).unwrap();

    snapshot!(snapshot_index_scheduler(&index_scheduler), name: "start");

    handle.advance_till([Start, BatchCreated]);

    let rtxn = index_scheduler.env.read_txn().unwrap();
    let proc = index_scheduler.processing_tasks.read().unwrap().clone();

    let query = Query { index_uids: Some(vec!["catto".to_owned()]), ..Default::default() };
    let (batches, _) = index_scheduler
        .queue
        .get_batch_ids_from_authorized_indexes(&rtxn, &query, &AuthFilter::default(), &proc)
        .unwrap();
    // only the first task associated with catto is returned, the indexSwap tasks are excluded!
    snapshot!(snapshot_bitmap(&batches), @"[0,]");

    let query = Query { index_uids: Some(vec!["catto".to_owned()]), ..Default::default() };
    let (batches, _) = index_scheduler
        .queue
        .get_batch_ids_from_authorized_indexes(
            &rtxn,
            &query,
            &AuthFilter::with_allowed_indexes(
                vec![IndexUidPattern::new_unchecked("doggo")].into_iter().collect(),
            ),
            &proc,
        )
        .unwrap();
    // we have asked for only the tasks associated with catto, but are only authorized to retrieve the tasks
    // associated with doggo -> empty result
    snapshot!(snapshot_bitmap(&batches), @"[]");

    drop(rtxn);
    // We're going to advance and process all the batches for the next query to actually hit the db
    handle.advance_till([
        InsideProcessBatch,
        InsideProcessBatch,
        ProcessBatchSucceeded,
        AfterProcessing,
    ]);
    handle.advance_one_successful_batch();
    handle.advance_n_failed_batches(2);
    snapshot!(snapshot_index_scheduler(&index_scheduler), name: "after-processing-everything");
    let rtxn = index_scheduler.env.read_txn().unwrap();

    let query = Query::default();
    let (batches, _) = index_scheduler
        .queue
        .get_batch_ids_from_authorized_indexes(
            &rtxn,
            &query,
            &AuthFilter::with_allowed_indexes(
                vec![IndexUidPattern::new_unchecked("doggo")].into_iter().collect(),
            ),
            &proc,
        )
        .unwrap();
    // we asked for all the tasks, but we are only authorized to retrieve the doggo tasks
    // -> only the index creation of doggo should be returned
    snapshot!(snapshot_bitmap(&batches), @"[1,]");

    let query = Query::default();
    let (batches, _) = index_scheduler
        .queue
        .get_batch_ids_from_authorized_indexes(
            &rtxn,
            &query,
            &AuthFilter::with_allowed_indexes(
                vec![
                    IndexUidPattern::new_unchecked("catto"),
                    IndexUidPattern::new_unchecked("doggo"),
                ]
                .into_iter()
                .collect(),
            ),
            &proc,
        )
        .unwrap();
    // we asked for all the tasks, but we are only authorized to retrieve the doggo and catto tasks
    // -> all tasks except the swap of catto with whalo are returned
    snapshot!(snapshot_bitmap(&batches), @"[0,1,]");

    let query = Query::default();
    let (batches, _) = index_scheduler
        .queue
        .get_batch_ids_from_authorized_indexes(&rtxn, &query, &AuthFilter::default(), &proc)
        .unwrap();
    // we asked for all the tasks with all index authorized -> all tasks returned
    snapshot!(snapshot_bitmap(&batches), @"[0,1,2,3,]");
}

#[test]
fn query_batches_canceled_by() {
    let (index_scheduler, mut handle) =
        IndexScheduler::test(true, vec![(3, FailureLocation::InsideProcessBatch)]);

    let kind = index_creation_task("catto", "mouse");
    let _ = index_scheduler.register(kind, None, false).unwrap();
    let kind = index_creation_task("doggo", "sheep");
    let _ = index_scheduler.register(kind, None, false).unwrap();
    let kind = KindWithContent::IndexSwap {
        swaps: vec![IndexSwap { indexes: ("catto".to_owned(), "doggo".to_owned()), rename: false }],
    };
    let _task = index_scheduler.register(kind, None, false).unwrap();

    handle.advance_n_successful_batches(1);
    let kind = KindWithContent::TaskCancelation {
        query: "test_query".to_string(),
        tasks: [0, 1, 2, 3].into_iter().collect(),
    };
    let task_cancelation = index_scheduler.register(kind, None, false).unwrap();
    handle.advance_n_successful_batches(1);

    snapshot!(snapshot_index_scheduler(&index_scheduler), name: "start");

    let query = Query { canceled_by: Some(vec![task_cancelation.uid]), ..Query::default() };
    let (batches, _) = index_scheduler
        .get_batch_ids_from_authorized_indexes(&query, &AuthFilter::default())
        .unwrap();
    // The batch zero was the index creation task, the 1 is the task cancellation
    snapshot!(snapshot_bitmap(&batches), @"[1,]");

    let query = Query { canceled_by: Some(vec![task_cancelation.uid]), ..Query::default() };
    let (batches, _) = index_scheduler
        .get_batch_ids_from_authorized_indexes(
            &query,
            &AuthFilter::with_allowed_indexes(
                vec![IndexUidPattern::new_unchecked("doggo")].into_iter().collect(),
            ),
        )
        .unwrap();
    // Return only 1 because the user is not authorized to see task 2
    snapshot!(snapshot_bitmap(&batches), @"[1,]");
}

#[test]
fn batch_skip_creation_with_deletion() {
    let (index_scheduler, mut handle) = IndexScheduler::test(true, vec![]);
    let kind = index_creation_task("docs", "id");
    let _task = index_scheduler.register(kind, None, false).unwrap();

    handle.advance_one_successful_batch();

    let (file0, documents_count0) = sample_documents(&index_scheduler, 1, 1);
    let (file1, documents_count1) = sample_documents(&index_scheduler, 2, 1);
    file0.persist().unwrap();
    file1.persist().unwrap();
    let kind = replace_document_import_task("docs", Some("id"), 1, documents_count0);
    index_scheduler.register(kind, None, false).unwrap();
    index_scheduler
        .register(
            KindWithContent::DocumentDeletion {
                index_uid: "docs".to_string(),
                documents_ids: vec!["1".to_string()],
            },
            None,
            false,
        )
        .unwrap();
    let kind = replace_document_import_task_with_opts(
        "docs",
        Some("id"),
        2,
        documents_count1,
        MissingDocumentPolicy::Skip,
    );
    index_scheduler.register(kind, None, false).unwrap();

    handle.advance_one_successful_batch();

    snapshot!(snapshot_index_scheduler(&index_scheduler));

    let index = index_scheduler.index("docs").unwrap();
    let rtxn = index.read_txn().unwrap();
    snapshot!(snapshot_bitmap(&index.documents_ids(&rtxn).unwrap()));
}

#[test]
fn batch_deletion_nothing_and_add_documents_no_guess_pk() {
    let (index_scheduler, mut handle) = IndexScheduler::test(true, vec![]);

    let _task = index_scheduler
        .register(
            KindWithContent::IndexCreation { index_uid: "docs".to_string(), primary_key: None },
            None,
            false,
        )
        .unwrap();

    handle.advance_one_successful_batch();

    // First register a deletion task
    index_scheduler
        .register(
            KindWithContent::DocumentDeletion {
                index_uid: "docs".to_string(),
                documents_ids: vec!["1".to_string()],
            },
            None,
            false,
        )
        .unwrap();

    // conflicting pk
    let content = r#"{ "id": "Hello!", "noid": "World!" }"#.to_string();

    let (_uuid, mut file0) = index_scheduler.queue.create_update_file_with_uuid(1).unwrap();
    let documents_count0 = read_json(content.as_bytes(), &mut file0).unwrap();
    file0.persist().unwrap();
    let kind = replace_document_import_task_with_opts(
        "docs",
        None,
        1,
        documents_count0,
        MissingDocumentPolicy::Create,
    );
    index_scheduler.register(kind, None, false).unwrap();

    // no primary
    let content = r#"{ "toto": "Hello!", "titi": "World!" }"#.to_string();

    let (_uuid, mut file1) = index_scheduler.queue.create_update_file_with_uuid(2).unwrap();
    let documents_count1 = read_json(content.as_bytes(), &mut file1).unwrap();
    file1.persist().unwrap();

    let kind = replace_document_import_task_with_opts(
        "docs",
        None,
        2,
        documents_count1,
        MissingDocumentPolicy::Create,
    );
    index_scheduler.register(kind, None, false).unwrap();

    // actually guessing the pk
    let content = r#"{ "id": "1234", "content": "Hello World!" }"#.to_string();

    let (_uuid, mut file2) = index_scheduler.queue.create_update_file_with_uuid(3).unwrap();
    let documents_count2 = read_json(content.as_bytes(), &mut file2).unwrap();
    file2.persist().unwrap();

    let kind = replace_document_import_task_with_opts(
        "docs",
        None,
        3,
        documents_count2,
        MissingDocumentPolicy::Create,
    );
    index_scheduler.register(kind, None, false).unwrap();

    // We now know the pk, GG!
    let content = r#"{ "id": "12345", "noid": "World!" }"#.to_string();

    let (_uuid, mut file3) = index_scheduler.queue.create_update_file_with_uuid(4).unwrap();
    let documents_count0 = read_json(content.as_bytes(), &mut file3).unwrap();
    file3.persist().unwrap();
    let kind = replace_document_import_task_with_opts(
        "docs",
        None,
        4,
        documents_count0,
        MissingDocumentPolicy::Create,
    );
    index_scheduler.register(kind, None, false).unwrap();

    handle.advance_one_successful_batch();

    snapshot!(snapshot_index_scheduler(&index_scheduler));

    let index = index_scheduler.index("docs").unwrap();
    let rtxn = index.read_txn().unwrap();
    snapshot!(snapshot_bitmap(&index.documents_ids(&rtxn).unwrap()));
}
