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

/// Rerank score result
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct RerankScore {
    pub index: usize,
    pub score: f32,
}

/// Rerank response from gateway
#[derive(Deserialize)]
struct RerankResponse {
    results: Vec<RerankScore>,
}

/// Rerank client using api.riotpiao.com gateway (BAAI/bge-reranker-base model)
pub struct RerankClient {
    base_url: String,
    model: String,
    timeout_secs: u64,
}

impl RerankClient {
    /// Create rerank client pointing to gateway
    pub fn new(base_url: &str, _api_key: &str, model: &str) -> Result<Self> {
        Ok(Self {
            base_url: base_url.to_string(),
            model: model.to_string(),
            timeout_secs: 300,
        })
    }

    /// Create from environment (uses api.riotpiao.com)
    pub fn from_env() -> Result<Self> {
        let base_url = std::env::var("LLM_API_BASE")
            .unwrap_or_else(|_| "https://api.riotpiao.com".to_string());
        let model = "BAAI/bge-reranker-base".to_string();
        
        Ok(Self {
            base_url,
            model,
            timeout_secs: 300,
        })
    }

    /// Rerank query against texts, return scored items in score order.
    /// Returns Vec<(index, score)> mapping back to input positions.
    pub async fn rerank(&self, query: &str, texts: &[&str]) -> Result<Vec<(usize, f32)>> {
        // Empty input returns empty without request
        if texts.is_empty() {
            return Ok(vec![]);
        }

        let url = format!("{}/v1/rerank", self.base_url);

        let client = Client::builder()
            .timeout(Duration::from_secs(self.timeout_secs))
            .build()?;

        let payload = serde_json::json!({
            "model": self.model,
            "query": query,
            "texts": texts,
            "top_k": texts.len(),
        });

        let response = client
            .post(&url)
            .header("Content-Type", "application/json")
            .json(&payload)
            .send()
            .await?;

        if !response.status().is_success() {
            let error_text = response.text().await.unwrap_or_default();
            return Err(anyhow::anyhow!("Rerank failed: {}", error_text));
        }

        // Parse gateway response (OpenAI format with results field)
        let resp: RerankResponse = response.json().await?;

        // Map to (index, score) and sort by score descending
        let mut results: Vec<(usize, f32)> = resp
            .results
            .into_iter()
            .map(|s| (s.index, s.score))
            .collect();

        results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());

        Ok(results)
    }
}