use std::sync::Mutex;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};

use anyhow::{Context, Result};
use fastembed::{EmbeddingModel, InitOptions, TextEmbedding};

pub struct Embedder {
    model: Mutex<Option<TextEmbedding>>,
    last_used_secs: AtomicU64,
}

pub fn query_prompt(text: &str) -> String {
    format!("Represent this sentence for searching relevant passages: {text}")
}

fn now_secs() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs()
}

fn load_model() -> Result<TextEmbedding> {
    TextEmbedding::try_new(
        InitOptions::new(EmbeddingModel::BGESmallENV15).with_show_download_progress(false),
    )
    .context("loading embedding model")
}

impl Embedder {
    pub fn new() -> Result<Self> {
        let model = TextEmbedding::try_new(
            InitOptions::new(EmbeddingModel::BGESmallENV15).with_show_download_progress(true),
        )
        .context("initializing embedding model")?;

        Ok(Self {
            model: Mutex::new(Some(model)),
            last_used_secs: AtomicU64::new(now_secs()),
        })
    }

    pub fn embed(&self, text: &str) -> Result<Vec<f32>> {
        let mut guard = self.model.lock().unwrap();
        if guard.is_none() {
            tracing::info!("reloading embedding model after idle drop");
            *guard = Some(load_model()?);
        }
        self.last_used_secs.store(now_secs(), Ordering::Relaxed);
        let results = guard
            .as_mut()
            .unwrap()
            .embed(vec![text], None)
            .context("generating embedding")?;
        results
            .into_iter()
            .next()
            .ok_or_else(|| anyhow::anyhow!("no embedding returned"))
    }

    // BGE retrieval is asymmetric: the instruction prefix goes on the query only,
    // documents stay bare, so embed() is left unchanged.
    pub fn embed_query(&self, text: &str) -> Result<Vec<f32>> {
        self.embed(&query_prompt(text))
    }

    pub fn embed_batch(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
        let mut guard = self.model.lock().unwrap();
        if guard.is_none() {
            tracing::info!("reloading embedding model after idle drop");
            *guard = Some(load_model()?);
        }
        self.last_used_secs.store(now_secs(), Ordering::Relaxed);
        let owned: Vec<String> = texts.iter().map(|t| t.to_string()).collect();
        let refs: Vec<&str> = owned.iter().map(|s| s.as_str()).collect();
        guard
            .as_mut()
            .unwrap()
            .embed(refs, None)
            .context("generating batch embeddings")
    }

    pub fn drop_if_idle(&self, idle_secs: u64) {
        let last = self.last_used_secs.load(Ordering::Relaxed);
        if now_secs().saturating_sub(last) >= idle_secs {
            let mut guard = self.model.lock().unwrap();
            if guard.is_some() {
                *guard = None;
                tracing::info!("embedding model dropped after {}s idle", idle_secs);
            }
        }
    }

    pub fn dimension(&self) -> usize {
        384
    }
}
