use anyhow::{Context, Result};
use rig_core::client::{CompletionClient, EmbeddingsClient, Nothing};
use rig_core::completion::CompletionModel;
use rig_core::embeddings::EmbeddingModel as _;
use rig_core::providers::ollama;
use serde::Deserialize;

/// Thin wrapper around a local Ollama instance via rig. Two models are used
/// deliberately: a small/cheap embedding model for the first-pass filter
/// over every article, and a larger chat model reserved for the shortlist
/// that clears the embedding threshold (see `feedsignal_core::scoring`).
pub struct Llm {
    client: ollama::Client,
    embedding_model: String,
    embedding_dims: usize,
    chat_model: String,
}

#[derive(Debug, Deserialize)]
pub struct RelevanceJudgment {
    pub score: f32,
    pub rationale: String,
}

impl Llm {
    /// `base_url` e.g. "http://localhost:11434". Model names must already
    /// be pulled locally, e.g. `ollama pull nomic-embed-text` and
    /// `ollama pull llama3.1`. `embedding_dims` must match the pulled
    /// embedding model (768 for nomic-embed-text, 384 for all-minilm).
    pub fn new(base_url: &str, embedding_model: impl Into<String>, embedding_dims: usize, chat_model: impl Into<String>) -> Result<Self> {
        let client = ollama::Client::builder()
            .api_key(Nothing)
            .base_url(base_url)
            .build()
            .context("failed to build ollama client")?;
        Ok(Self {
            client,
            embedding_model: embedding_model.into(),
            embedding_dims,
            chat_model: chat_model.into(),
        })
    }

    pub async fn embed(&self, text: &str) -> Result<Vec<f32>> {
        let model = self.client.embedding_model_with_ndims(&self.embedding_model, self.embedding_dims);
        let embedding = model.embed_text(text).await.context("ollama embedding request failed")?;
        Ok(embedding.vec.into_iter().map(|v| v as f32).collect())
    }

    /// Runs the second-stage relevance judgment for one article. Called
    /// only for the embedding-similarity shortlist, not every article —
    /// see `feedsignal-db::shortlist_for_llm_scoring`.
    pub async fn judge_relevance(
        &self,
        preferences_description: &str,
        topic_affinities_summary: &str,
        article_title: &str,
        article_summary: &str,
    ) -> Result<RelevanceJudgment> {
        let model = self.client.completion_model(&self.chat_model);

        let prompt = format!(
            "Preferences: {preferences_description}\n\
             Historical topic affinities: {topic_affinities_summary}\n\n\
             Article title: {article_title}\n\
             Article summary: {article_summary}"
        );

        let request = model
            .completion_request(prompt)
            .preamble(
                "You screen articles for a personal RSS reader. Given the \
                 reader's stated preferences, their historical topic \
                 affinities, and one article, respond with ONLY a JSON \
                 object: {\"score\": <0.0-1.0>, \"rationale\": \"<one \
                 sentence>\"}. score is how relevant this specific article \
                 is to this specific reader right now."
                    .to_string(),
            )
            .build();

        let response = model.completion(request).await.context("ollama chat request failed")?;
        let text: String = response
            .choice
            .into_iter()
            .filter_map(|content| match content {
                rig_core::completion::AssistantContent::Text(t) => Some(t.text),
                _ => None,
            })
            .collect();
        parse_judgment(&text)
    }
}

fn parse_judgment(raw: &str) -> Result<RelevanceJudgment> {
    // Models sometimes wrap JSON in prose or code fences despite instructions;
    // take the outermost {...} span rather than trusting the whole response.
    let start = raw.find('{').context("no JSON object in LLM response")?;
    let end = raw.rfind('}').context("no JSON object in LLM response")?;
    let json = &raw[start..=end];
    Ok(serde_json::from_str(json)?)
}