use std::{cmp::Ordering, sync::Arc};

use anyhow::{Result, anyhow};
use binary_heap_plus::{BinaryHeap, FnComparator};
use log::warn;
use tantivy::{TantivyDocument, schema::Type};

use crate::{
    reader::{
        FragReader, query_info::QueryInfo, query_result::TopKValue, shard_reader::ShardReaderStat,
    },
    utils::{cmp_doc::cmp_opt_f64, schema::SortField},
};

pub struct HeapItem {
    pub val: TopKValue,
    pub seg_idx: usize,
}

pub struct FastMerger<F>
where
    F: Fn(&HeapItem, &HeapItem) -> Ordering,
{
    list_iters: Vec<std::vec::IntoIter<TopKValue>>,
    heap: BinaryHeap<HeapItem, FnComparator<F>>,
}

impl<F> FastMerger<F>
where
    F: Fn(&HeapItem, &HeapItem) -> Ordering,
{
    pub fn new(lists: Vec<Vec<TopKValue>>, cmp: F) -> Result<Self> {
        let mut heap = BinaryHeap::new_by(cmp);
        let mut list_iters: Vec<std::vec::IntoIter<TopKValue>> =
            lists.into_iter().map(|v| v.into_iter()).collect();

        for (i, iter) in list_iters.iter_mut().enumerate() {
            if let Some(val) = iter.next() {
                heap.push(HeapItem { val, seg_idx: i });
            }
        }

        Ok(Self { list_iters, heap })
    }

    pub fn next_item(&mut self) -> Result<Option<HeapItem>> {
        if let Some(top) = self.heap.pop() {
            let seg_idx = top.seg_idx;
            let val = top.val;
            if let Some(next_val) = self.list_iters[seg_idx].next() {
                self.heap.push(HeapItem {
                    val: next_val,
                    seg_idx,
                });
            }
            return Ok(Some(HeapItem { val, seg_idx }));
        }
        Ok(None)
    }
}

fn compare_topk(a: &TopKValue, b: &TopKValue, sort_fields: &[SortField]) -> Ordering {
    // Contract: a.cols/b.cols are aligned with sort_fields by index.
    for (i, sf) in sort_fields.iter().enumerate() {
        let ac = &a.cols[i];
        let bc = &b.cols[i];

        let ord = match sf.field.value_type {
            Type::I64 | Type::U64 => ac.int_value.cmp(&bc.int_value),
            Type::F64 => cmp_opt_f64(ac.float_value, bc.float_value),
            Type::Str | Type::Bytes => ac.string_value.cmp(&bc.string_value),
            _ => {
                warn!(
                    "Unsupported sort field type in compare_topk, treat as equal: field={:?}",
                    sf.field
                );
                Ordering::Equal
            }
        };

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

    // fallback deterministic tie-breaker using DocAddress (segment ord then doc id)
    let seg_cmp = a.doc_addr.segment_ord.cmp(&b.doc_addr.segment_ord);
    if seg_cmp != Ordering::Equal {
        return seg_cmp;
    }
    a.doc_addr.doc_id.cmp(&b.doc_addr.doc_id)
}

pub fn search_with_top_collector(
    frag_readers: &[Arc<FragReader>],
    query_info: &QueryInfo,
    search_stat: &mut ShardReaderStat,
) -> Result<Vec<Arc<TantivyDocument>>> {
    let top_k_param = query_info
        .topk_params()
        .ok_or_else(|| anyhow!("topk_params is not set"))?;
    if top_k_param.limit == 0 {
        return Ok(Vec::new());
    }

    // collect partial TopK lists from each frag reader
    let mut lists: Vec<Vec<TopKValue>> = Vec::with_capacity(frag_readers.len());
    for frag_reader in frag_readers {
        let partial_result = frag_reader.search_with_top_collector(query_info, search_stat)?;
        lists.push(partial_result);
    }

    let sort_fields = top_k_param.sort_fields.clone();
    let cmp = move |a: &HeapItem, b: &HeapItem| -> Ordering {
        // keep same direction as previous code: compare b vs a
        compare_topk(&b.val, &a.val, &sort_fields)
    };

    let mut merger = FastMerger::new(lists, cmp)?;
    let mut docs: Vec<Arc<TantivyDocument>> = Vec::with_capacity(top_k_param.limit);

    // Collect DocAddress from merged TopKValue and fetch actual docs from
    // frag_readers
    while docs.len() < top_k_param.limit {
        match merger.next_item()? {
            Some(item) => {
                let addr = item.val.doc_addr;
                let seg_idx = item.seg_idx;
                if seg_idx >= frag_readers.len() {
                    continue;
                }
                let frag_reader = &frag_readers[seg_idx];
                let doc = frag_reader.fetch_doc(addr)?;
                docs.push(doc);
            }
            None => break,
        }
    }

    Ok(docs)
}