// BGE-M3 编码管线 — 通过 HTTP 调用本地 ONNX BGE 服务 (localhost:8000)
// 使用 OpenAI-compatible /v1/embeddings 端点
// 替代方案: 等 ort 2.0 stable 后用 Rust 原生 ONNX

use reqwest::blocking::Client;
use serde::{Deserialize, Serialize};
use std::time::Duration;

#[derive(Debug, Serialize)]
struct EmbedRequest {
    input: EmbedInput,
    model: String,
}

#[derive(Debug, Serialize)]
#[serde(untagged)]
enum EmbedInput {
    Single(String),
    Batch(Vec<String>),
}

#[derive(Debug, Deserialize)]
struct EmbedResponse {
    data: Vec<EmbedData>,
}

#[derive(Debug, Deserialize)]
struct EmbedData {
    embedding: Vec<f32>,
}

pub struct BGEEncoder {
    pub model_dir: String,
    pub dim: usize,
    endpoint: String,
    client: Client,
}

impl BGEEncoder {
    pub fn new(model_dir: &str) -> Result<Self, Box<dyn std::error::Error>> {
        Ok(Self {
            model_dir: model_dir.to_string(),
            dim: 1024,
            endpoint: "http://localhost:8000/v1/embeddings".to_string(),
            client: Client::builder()
                .timeout(Duration::from_secs(30))
                .build()?,
        })
    }

    pub fn encode(&self, text: &str) -> Result<Vec<f32>, Box<dyn std::error::Error>> {
        let results = self.encode_batch(&[text.to_string()])?;
        results
            .into_iter()
            .next()
            .ok_or_else(|| "BGE encode: empty response".into())
    }

    pub fn encode_batch(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, Box<dyn std::error::Error>> {
        if texts.is_empty() {
            return Ok(Vec::new());
        }

        let req = EmbedRequest {
            input: EmbedInput::Batch(texts.to_vec()),
            model: "bge-m3".to_string(),
        };

        let resp = self
            .client
            .post(&self.endpoint)
            .header("Content-Type", "application/json")
            .timeout(Duration::from_secs(10))
            .body(serde_json::to_string(&req)?)
            .send()?;

        if !resp.status().is_success() {
            return Err(format!("BGE HTTP error: {} — {}", resp.status(), resp.text().unwrap_or_default()).into());
        }

        let body = resp.text()?;
        let parsed: EmbedResponse = serde_json::from_str(&body)?;

        let vectors: Vec<Vec<f32>> = parsed.data.into_iter().map(|d| d.embedding).collect();

        if vectors.len() != texts.len() {
            return Err(format!(
                "BGE encode: expected {} vectors, got {}",
                texts.len(),
                vectors.len()
            ).into());
        }

        Ok(vectors)
    }
}