//! Generic LLM client abstraction for the Rust Unmute conversational stack.
//!
//! Provides a uniform [`LlmClient`] trait with `generate_stream` returning
//! a stream of text tokens. Implementations:
//!
//!   - [`OpenAiCompatibleClient`] — works with OpenAI's Chat Completions
//!     API and any compatible endpoint (Z.AI, vLLM, llama.cpp's HTTP server,
//!     LiteLLM, etc.). Streams via Server-Sent Events.
//!
//! This is the bridge layer in the STT → LLM → TTS conversational pipeline.
//! Token-level streaming is critical: TTS can start speaking the assistant
//! response as soon as the first token arrives, instead of waiting for the
//! full LLM completion.
//!
//! ## Example
//!
//! ```no_run
//! use rtx_csm::llm_client::{ChatMessage, GenConfig, LlmClient, OpenAiCompatibleClient, Role};
//! use futures_util::StreamExt;
//!
//! # async fn run() -> anyhow::Result<()> {
//! let client = OpenAiCompatibleClient::new(
//!     "https://api.openai.com/v1",
//!     std::env::var("OPENAI_API_KEY")?,
//!     "gpt-4o-mini",
//! );
//! let messages = vec![
//!     ChatMessage::system("You are a concise assistant."),
//!     ChatMessage::user("Say hello in one word."),
//! ];
//! let mut stream = client.generate_stream(messages, GenConfig::default()).await?;
//! while let Some(tok) = stream.next().await {
//!     print!("{}", tok?);
//! }
//! # Ok(()) }
//! ```

use crate::error::{CsmError, Result};
use async_trait::async_trait;
use eventsource_stream::Eventsource;
use futures_util::stream::{Stream, StreamExt};
use serde::{Deserialize, Serialize};
use std::pin::Pin;

/// Boxed stream of text tokens. Each token is one chunk of the assistant
/// response — typically a sub-word from the LLM tokenizer. Concatenate to
/// rebuild the full message; pipe to TTS as they arrive for streaming UX.
pub type TokenStream = Pin<Box<dyn Stream<Item = Result<String>> + Send>>;

#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum Role {
    System,
    User,
    Assistant,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatMessage {
    pub role: Role,
    pub content: String,
}

impl ChatMessage {
    pub fn system(content: impl Into<String>) -> Self {
        Self {
            role: Role::System,
            content: content.into(),
        }
    }
    pub fn user(content: impl Into<String>) -> Self {
        Self {
            role: Role::User,
            content: content.into(),
        }
    }
    pub fn assistant(content: impl Into<String>) -> Self {
        Self {
            role: Role::Assistant,
            content: content.into(),
        }
    }
}

#[derive(Debug, Clone)]
pub struct GenConfig {
    /// Maximum tokens in the response.
    pub max_tokens: Option<u32>,
    /// Sampling temperature; 0.0 = greedy.
    pub temperature: f32,
    /// Top-p nucleus cutoff. 1.0 = disabled.
    pub top_p: f32,
    /// Optional stop sequences.
    pub stop: Vec<String>,
    /// Provider-specific extra fields merged into the JSON request body
    /// at top level. Used for things the OpenAI schema doesn't expose:
    ///
    /// - **Z.AI thinking-disabled** (skips reasoning_content for
    ///   voice-AI use): `{"thinking": {"type": "disabled"}}`. Without
    ///   this, glm-4.5/4.6/4.7 burn most of `max_tokens` on
    ///   reasoning_content and emit empty `content`.
    /// - **Anthropic thinking budget** (when proxied through an
    ///   OpenAI-compatible Anthropic shim): `{"thinking": {"type":
    ///   "enabled", "budget_tokens": 1024}}`.
    /// - **vLLM guided decoding**: `{"guided_json": {...}}`.
    pub extra_body: serde_json::Map<String, serde_json::Value>,
}

impl Default for GenConfig {
    fn default() -> Self {
        Self {
            max_tokens: Some(512),
            temperature: 0.7,
            top_p: 1.0,
            stop: Vec::new(),
            extra_body: serde_json::Map::new(),
        }
    }
}

#[async_trait]
pub trait LlmClient: Send + Sync {
    /// Generate a streaming response. The returned stream yields text
    /// chunks as the model produces them. Each chunk is typically one or
    /// a few tokens; concatenate for the full response.
    async fn generate_stream(
        &self,
        messages: Vec<ChatMessage>,
        config: GenConfig,
    ) -> Result<TokenStream>;

    /// Convenience: collect the full response. Default impl folds the
    /// stream into a single String. Implementations are free to override
    /// for non-streaming endpoints.
    async fn generate(&self, messages: Vec<ChatMessage>, config: GenConfig) -> Result<String> {
        let mut stream = self.generate_stream(messages, config).await?;
        let mut out = String::new();
        while let Some(chunk) = stream.next().await {
            out.push_str(&chunk?);
        }
        Ok(out)
    }
}

// -- OpenAI-compatible Chat Completions impl ------------------------------

/// HTTP client for OpenAI-compatible Chat Completions endpoints.
/// Tested patterns:
///   - OpenAI: base="https://api.openai.com/v1", model="gpt-4o-mini"
///   - Z.AI:   base="https://api.z.ai/api/coding/paas/v4", model="glm-4.6"
///   - vLLM:   base="http://localhost:8000/v1", model="<served-model-name>"
///   - llama.cpp: base="http://localhost:8080/v1", model="<arbitrary>"
#[derive(Debug, Clone)]
pub struct OpenAiCompatibleClient {
    base_url: String,
    api_key: String,
    model: String,
    http: reqwest::Client,
}

impl OpenAiCompatibleClient {
    pub fn new(
        base_url: impl Into<String>,
        api_key: impl Into<String>,
        model: impl Into<String>,
    ) -> Self {
        Self {
            base_url: base_url.into(),
            api_key: api_key.into(),
            model: model.into(),
            http: reqwest::Client::builder()
                .timeout(std::time::Duration::from_secs(120))
                .build()
                .expect("reqwest client init"),
        }
    }
}

// -- Wire types: OpenAI Chat Completions request/response -----------------

#[derive(Serialize)]
struct ChatRequest<'a> {
    model: &'a str,
    messages: &'a [ChatMessage],
    stream: bool,
    #[serde(skip_serializing_if = "Option::is_none")]
    max_tokens: Option<u32>,
    temperature: f32,
    top_p: f32,
    #[serde(skip_serializing_if = "Vec::is_empty")]
    stop: Vec<String>,
}

#[derive(Deserialize)]
struct ChatStreamEvent {
    choices: Vec<ChatStreamChoice>,
}

#[derive(Deserialize)]
struct ChatStreamChoice {
    delta: ChatStreamDelta,
    /// Captured for completeness; OpenAI sets to `"stop"`/`"length"` on
    /// the last event. We don't currently surface it to callers.
    #[serde(default)]
    #[allow(dead_code)]
    finish_reason: Option<String>,
}

#[derive(Deserialize, Default)]
struct ChatStreamDelta {
    #[serde(default)]
    content: Option<String>,
}

#[async_trait]
impl LlmClient for OpenAiCompatibleClient {
    async fn generate_stream(
        &self,
        messages: Vec<ChatMessage>,
        config: GenConfig,
    ) -> Result<TokenStream> {
        let url = format!("{}/chat/completions", self.base_url.trim_end_matches('/'));
        let req = ChatRequest {
            model: &self.model,
            messages: &messages,
            stream: true,
            max_tokens: config.max_tokens,
            temperature: config.temperature,
            top_p: config.top_p,
            stop: config.stop.clone(),
        };
        // Serialize the typed request, then merge in extra_body at the
        // top level so provider-specific fields (e.g., Z.AI's `thinking`)
        // ride alongside the OpenAI schema. extra_body is empty in the
        // default case so this is a no-op.
        let mut body = serde_json::to_value(&req)
            .map_err(|e| CsmError::Config(format!("LLM serialize: {e}")))?;
        if !config.extra_body.is_empty()
            && let Some(obj) = body.as_object_mut()
        {
            for (k, v) in config.extra_body.iter() {
                obj.insert(k.clone(), v.clone());
            }
        }
        let response = self
            .http
            .post(&url)
            .bearer_auth(&self.api_key)
            .header("accept", "text/event-stream")
            .json(&body)
            .send()
            .await
            .map_err(|e| CsmError::Config(format!("LLM request: {e}")))?;
        if !response.status().is_success() {
            let status = response.status();
            let body = response.text().await.unwrap_or_default();
            return Err(CsmError::Config(format!("LLM HTTP {status}: {body}")));
        }

        // SSE stream: each event has data: <json>\n\n. Final event is data: [DONE].
        let bytes_stream = response.bytes_stream();
        let event_stream = bytes_stream.eventsource();
        let token_stream = event_stream.filter_map(|ev| async move {
            match ev {
                Err(e) => Some(Err(CsmError::Config(format!("SSE: {e}")))),
                Ok(event) => {
                    let data = event.data;
                    if data.trim() == "[DONE]" {
                        None
                    } else {
                        match serde_json::from_str::<ChatStreamEvent>(&data) {
                            Ok(parsed) => parsed
                                .choices
                                .into_iter()
                                .next()
                                .and_then(|c| c.delta.content)
                                .filter(|s| !s.is_empty())
                                .map(Ok),
                            Err(e) => {
                                Some(Err(CsmError::Config(format!("SSE parse: {e} body={data}"))))
                            }
                        }
                    }
                }
            }
        });
        Ok(Box::pin(token_stream))
    }
}

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

    #[test]
    fn chat_message_constructors() {
        assert_eq!(ChatMessage::system("s").role, Role::System);
        assert_eq!(ChatMessage::user("u").role, Role::User);
        assert_eq!(ChatMessage::assistant("a").role, Role::Assistant);
    }

    #[test]
    fn role_serializes_lowercase() {
        let m = ChatMessage::user("hi");
        let json = serde_json::to_string(&m).unwrap();
        assert!(json.contains("\"role\":\"user\""), "got {json}");
    }

    #[test]
    fn gen_config_defaults_sane() {
        let c = GenConfig::default();
        assert_eq!(c.max_tokens, Some(512));
        assert!((c.temperature - 0.7).abs() < 1e-6);
        assert_eq!(c.top_p, 1.0);
        assert!(c.stop.is_empty());
    }
}