use std::cmp::Ordering;

use tantivy::{TantivyDocument, schema::Type};

use crate::utils::{schema::SortField, tantivy_doc::DocFieldAccess};

/// Compare two optional f64 values with a total order:
/// - None (NULL) is smallest
/// - NaN is largest
/// - Otherwise use IEEE total order over f64 values
pub fn cmp_opt_f64(a: Option<f64>, b: Option<f64>) -> Ordering {
    let to_key = |v: Option<f64>| -> (u8, f64) {
        match v {
            None => (0, 0.0),                  // NULL is smallest
            Some(x) if x.is_nan() => (2, 0.0), // NaN is largest
            Some(x) => (1, x),
        }
    };

    let (k1_tag, k1_val) = to_key(a);
    let (k2_tag, k2_val) = to_key(b);
    let ord = k1_tag.cmp(&k2_tag);
    if ord == Ordering::Equal && k1_tag == 1 {
        k1_val.total_cmp(&k2_val)
    } else {
        ord
    }
}

pub fn compare_docs(a: &TantivyDocument, b: &TantivyDocument, columns: &[SortField]) -> Ordering {
    for sort_field in columns {
        let field = &sort_field.field;
        let cmp = match field.value_type {
            Type::Str => {
                let v1 = a.get_str(field).unwrap();
                let v2 = b.get_str(field).unwrap();
                v1.cmp(&v2)
            }
            Type::U64 => {
                let v1 = a.get_u64(field).unwrap();
                let v2 = b.get_u64(field).unwrap();
                v1.cmp(&v2)
            }
            Type::I64 => {
                let v1 = a.get_i64(field).unwrap();
                let v2 = b.get_i64(field).unwrap();
                v1.cmp(&v2)
            }
            Type::F64 => {
                let v1 = a.get_f64(field).unwrap();
                let v2 = b.get_f64(field).unwrap();
                cmp_opt_f64(v1, v2)
            }
            Type::Bytes => {
                let v1 = a.get_bytes(field).unwrap();
                let v2 = b.get_bytes(field).unwrap();
                v1.cmp(&v2)
            }
            _ => panic!("Unsupported field={:?}", field),
        };

        let ord = if sort_field.is_asc {
            cmp
        } else {
            cmp.reverse()
        };
        if ord != Ordering::Equal {
            return ord;
        }
    }

    Ordering::Equal
}

pub fn sort_docs(mut docs: Vec<TantivyDocument>, columns: &[SortField]) -> Vec<TantivyDocument> {
    if columns.is_empty() {
        return docs;
    }
    docs.sort_by(|a, b| compare_docs(a, b, columns));
    docs
}

pub fn is_sorted(docs: &[TantivyDocument], columns: &[SortField]) -> bool {
    if columns.is_empty() || docs.len() <= 1 {
        return true;
    }

    for i in 0..docs.len() - 1 {
        if compare_docs(&docs[i], &docs[i + 1], columns) == Ordering::Greater {
            return false;
        }
    }

    true
}

#[cfg(test)]
mod tests {
    use rand::{rng, seq::SliceRandom};
    use tantivy::{
        TantivyDocument,
        schema::{Field, Type},
    };

    use super::*;
    use crate::{
        reader::query_info::{DocField, TantivySchemaExt},
        test_utils::{TEST_INDEX_META_OVERLAP, test_gen_raw_docs},
        utils::tantivy_doc::TantivyDoc,
    };

    #[test]
    fn test_sort() {
        let index_meta = TEST_INDEX_META_OVERLAP.clone();
        let schema = index_meta.get_tantivy_schema();

        let (raw_docs, _expected_raw_docs) = test_gen_raw_docs(10000, 0, 0);
        let docs = raw_docs
            .iter()
            .map(|raw_doc| TantivyDoc::new(raw_doc, &index_meta).unwrap().doc)
            .collect::<Vec<_>>();

        let doc_field = schema.doc_field("column_1").unwrap();
        let columns = vec![SortField {
            field: doc_field.clone(),
            is_asc: true,
        }];
        let docs = sort_docs(docs, &columns);
        assert!(is_sorted(&docs, &columns));
    }

    fn run_sort_test<T, FAdd>(values: Vec<Option<T>>, value_type: Type, add_fn: FAdd)
    where
        T: Clone,
        FAdd: Fn(&mut TantivyDocument, Field, &T),
    {
        let mut shuffle_v = values.clone();
        shuffle_v.shuffle(&mut rng());

        let field = Field::from_field_id(1);
        let mut shuffled_docs = Vec::new();
        for v in shuffle_v {
            let mut doc = TantivyDocument::new();
            if let Some(val) = v {
                add_fn(&mut doc, field, &val);
            }
            shuffled_docs.push(doc);
        }

        let doc_field = DocField {
            name: String::from("column_1"),
            field,
            value_type,
        };

        let asc_columns = vec![SortField {
            field: doc_field.clone(),
            is_asc: true,
        }];

        let desc_columns = vec![SortField {
            field: doc_field.clone(),
            is_asc: false,
        }];

        let asc_docs = sort_docs(shuffled_docs.clone(), &asc_columns);
        assert!(is_sorted(&asc_docs, &asc_columns));
        assert!(!is_sorted(&asc_docs, &desc_columns));

        let desc_docs = sort_docs(shuffled_docs, &desc_columns);
        assert!(is_sorted(&desc_docs, &desc_columns));
        assert!(!is_sorted(&desc_docs, &asc_columns));
    }

    #[test]
    fn test_sort_i64() {
        let values = vec![
            None,
            Some(-5),
            Some(-4),
            Some(-3),
            Some(-2),
            Some(-1),
            Some(0),
            Some(1),
            Some(2),
            Some(3),
            Some(4),
            Some(5),
        ];
        run_sort_test(values, Type::I64, |doc, field, val| {
            doc.add_i64(field, *val)
        });
    }

    #[test]
    fn test_sort_u64() {
        let values = vec![
            None,
            Some(0),
            Some(1),
            Some(2),
            Some(3),
            Some(4),
            Some(5),
            Some(6),
            Some(7),
            Some(8),
            Some(9),
            Some(10),
        ];
        run_sort_test(values, Type::U64, |doc, field, val| {
            doc.add_u64(field, *val)
        });
    }

    #[test]
    fn test_sort_f64() {
        let values = vec![
            None,
            Some(f64::NEG_INFINITY),
            Some(-1.5),
            Some(0.0),
            Some(std::f64::consts::PI),
            Some(f64::INFINITY),
            Some(f64::NAN),
        ];
        run_sort_test(values, Type::F64, |doc, field, val| {
            doc.add_f64(field, *val)
        });
    }

    #[test]
    fn test_sort_str() {
        let values = vec![
            None,
            Some("apple".to_string()),
            Some("banana".to_string()),
            Some("cherry".to_string()),
            Some("date".to_string()),
            Some("zebra".to_string()),
        ];
        run_sort_test(values, Type::Str, |doc, field, val| {
            doc.add_text(field, val)
        });
    }

    #[test]
    fn test_sort_bytes() {
        let values = vec![
            None,
            Some(vec![0u8, 1]),
            Some(vec![1, 2, 3]),
            Some(vec![2]),
            Some(vec![10, 20, 30]),
            Some(vec![255]),
        ];
        run_sort_test(values, Type::Bytes, |doc, field, val| {
            doc.add_bytes(field, val)
        });
    }

    #[test]
    fn test_sort_multi_fields_mixed_types() {
        let int_values = [
            Some(1),
            Some(1),
            Some(2),
            Some(2),
            Some(3),
            Some(3),
            Some(4),
            Some(4),
            Some(5),
            Some(5),
        ];

        let str_values = [
            Some("apple".to_string()),
            Some("banana".to_string()),
            Some("apple".to_string()),
            Some("banana".to_string()),
            Some("apple".to_string()),
            Some("banana".to_string()),
            Some("apple".to_string()),
            Some("banana".to_string()),
            Some("apple".to_string()),
            Some("banana".to_string()),
        ];

        let f64_values = [
            Some(1.1),
            Some(0.5),
            Some(2.0),
            Some(1.5),
            Some(3.3),
            Some(3.2),
            Some(4.0),
            Some(4.5),
            Some(5.0),
            Some(5.1),
        ];

        let bytes_values = [
            Some(vec![0u8, 1]),
            Some(vec![1, 2]),
            Some(vec![0, 5]),
            Some(vec![1, 3]),
            Some(vec![2, 0]),
            Some(vec![2, 1]),
            Some(vec![3, 0]),
            Some(vec![3, 1]),
            Some(vec![4, 0]),
            Some(vec![4, 1]),
        ];

        let len = int_values.len();
        let mut indices: Vec<usize> = (0..len).collect();
        indices.shuffle(&mut rng());

        let int_field = Field::from_field_id(1);
        let str_field = Field::from_field_id(2);
        let f64_field = Field::from_field_id(3);
        let bytes_field = Field::from_field_id(4);

        let mut docs = Vec::new();
        for &idx in &indices {
            let mut doc = TantivyDocument::new();
            if let Some(v) = int_values[idx] {
                doc.add_i64(int_field, v);
            }
            if let Some(ref s) = str_values[idx] {
                doc.add_text(str_field, s);
            }
            if let Some(f) = f64_values[idx] {
                doc.add_f64(f64_field, f);
            }
            if let Some(ref b) = bytes_values[idx] {
                doc.add_bytes(bytes_field, b);
            }
            docs.push(doc);
        }

        let int_doc_field = DocField {
            name: "int_col".to_string(),
            field: int_field,
            value_type: Type::I64,
        };
        let str_doc_field = DocField {
            name: "str_col".to_string(),
            field: str_field,
            value_type: Type::Str,
        };
        let f64_doc_field = DocField {
            name: "f64_col".to_string(),
            field: f64_field,
            value_type: Type::F64,
        };
        let bytes_doc_field = DocField {
            name: "bytes_col".to_string(),
            field: bytes_field,
            value_type: Type::Bytes,
        };

        let columns = vec![
            SortField {
                field: int_doc_field.clone(),
                is_asc: true,
            },
            SortField {
                field: str_doc_field.clone(),
                is_asc: false,
            },
            SortField {
                field: f64_doc_field.clone(),
                is_asc: true,
            },
            SortField {
                field: bytes_doc_field.clone(),
                is_asc: false,
            },
        ];

        let sorted_docs = sort_docs(docs.clone(), &columns);
        assert!(is_sorted(&sorted_docs, &columns));

        let desc_columns = vec![
            SortField {
                field: int_doc_field.clone(),
                is_asc: false,
            },
            SortField {
                field: str_doc_field.clone(),
                is_asc: false,
            },
            SortField {
                field: f64_doc_field.clone(),
                is_asc: false,
            },
            SortField {
                field: bytes_doc_field.clone(),
                is_asc: false,
            },
        ];
        let sorted_desc_docs = sort_docs(docs.clone(), &desc_columns);
        assert!(is_sorted(&sorted_desc_docs, &desc_columns));

        let asc_columns = vec![
            SortField {
                field: int_doc_field,
                is_asc: true,
            },
            SortField {
                field: str_doc_field,
                is_asc: true,
            },
            SortField {
                field: f64_doc_field,
                is_asc: true,
            },
            SortField {
                field: bytes_doc_field,
                is_asc: true,
            },
        ];
        let sorted_asc_docs = sort_docs(docs, &asc_columns);
        assert!(is_sorted(&sorted_asc_docs, &asc_columns));
    }
}