use crate::backend_rust::RustBackend;
use crate::backend_swift::SwiftBackend;
use crate::config::Config;
use crate::dto::{EmbedRequest, EmbedResponse};
use crate::errors::OrchestratorError;
use crate::tokenizer::{round_up, TokenCounter};

/// Build a sub-request containing only the texts at `indices`, preserving order.
fn sub_request(req: &EmbedRequest, indices: &[usize]) -> EmbedRequest {
    EmbedRequest {
        texts: indices.iter().map(|&i| req.texts[i].clone()).collect(),
        input_type: req.input_type,
        size: req.size.clone(),
    }
}

/// Partition indices into (short, long) by the threshold.
pub fn partition(counts: &[usize], threshold: usize) -> (Vec<usize>, Vec<usize>) {
    let mut short_idx = Vec::new();
    let mut long_idx = Vec::new();
    for (i, &c) in counts.iter().enumerate() {
        if c <= threshold {
            short_idx.push(i);
        } else {
            long_idx.push(i);
        }
    }
    (short_idx, long_idx)
}

pub async fn dispatch(
    req: EmbedRequest,
    tokenizer: &TokenCounter,
    swift: &SwiftBackend,
    rust: &RustBackend,
    cfg: &Config,
) -> Result<EmbedResponse, OrchestratorError> {
    if req.texts.is_empty() {
        return Err(OrchestratorError::EmptyTexts);
    }
    if let Some(s) = req.size.as_deref() {
        if s != "0.6B" {
            return Err(OrchestratorError::UnsupportedSize);
        }
    }

    // 1. Wrap + count each text.
    let counts: Vec<usize> = req
        .texts
        .iter()
        .map(|t| tokenizer.count_wrapped(t, req.input_type))
        .collect::<Result<_, _>>()?;

    // 2. Validate model max.
    for (i, &c) in counts.iter().enumerate() {
        if c > cfg.model_max_tokens {
            return Err(OrchestratorError::TextTooLong {
                index: i,
                tokens: c,
                limit: cfg.model_max_tokens,
            });
        }
    }

    // 3. Partition.
    let (short_idx, long_idx) = partition(&counts, cfg.swift_threshold);

    // 4. Dispatch in parallel.
    let short_fut = async {
        if short_idx.is_empty() {
            return Ok::<Vec<Vec<f32>>, OrchestratorError>(vec![]);
        }
        let sub = sub_request(&req, &short_idx);
        Ok(swift.embed(&sub).await?.vectors)
    };
    let long_fut = async {
        if long_idx.is_empty() {
            return Ok::<Vec<Vec<f32>>, OrchestratorError>(vec![]);
        }
        let max_tok = long_idx.iter().map(|i| counts[*i]).max().unwrap();
        let bucket =
            round_up(max_tok, cfg.rust_bucket_granularity).clamp(1024, cfg.model_max_tokens);
        let sub = sub_request(&req, &long_idx);
        Ok(rust.run_one(bucket, &sub).await?.vectors)
    };
    let (shorts, longs) = tokio::try_join!(short_fut, long_fut)?;

    // 5. Merge back in original order.
    let mut out: Vec<Vec<f32>> = vec![Vec::new(); req.texts.len()];
    let mut dim = 0usize;
    for (slot, v) in short_idx.iter().zip(shorts) {
        if dim == 0 {
            dim = v.len();
        }
        out[*slot] = v;
    }
    for (slot, v) in long_idx.iter().zip(longs) {
        if dim == 0 {
            dim = v.len();
        }
        out[*slot] = v;
    }

    Ok(EmbedResponse {
        vectors: out,
        dim,
        model: "qwen3-embedding-0.6B".into(),
    })
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn partition_basic() {
        let counts = vec![25, 30, 8200, 28, 15040];
        let (s, l) = partition(&counts, 8192);
        assert_eq!(s, vec![0, 1, 3]);
        assert_eq!(l, vec![2, 4]);
    }

    #[test]
    fn partition_all_short() {
        let (s, l) = partition(&[1, 2, 3], 8192);
        assert_eq!(s, vec![0, 1, 2]);
        assert!(l.is_empty());
    }
}