// PUBLIC DOMAIN - NO LICENSE, NO WARRANTY
// Copyright 2025 TimeHexOn & foxhop & russell@unturf
// https://www.permacomputer.com

//! UncloseAI - Rust Client Library
//!
//! A Rust client for interacting with vLLM, Ollama, and OpenAI-compatible APIs.
//!
//! # Features
//!
//! - Automatic model discovery from environment variables
//! - Streaming and non-streaming chat completions
//! - Text-to-speech generation
//! - Support for multiple endpoints
//! - Type-safe API with comprehensive error handling
//!
//! # Example
//!
//! ```no_run
//! use uncloseai::{UncloseAI, ChatMessage};
//!
//! #[tokio::main]
//! async fn main() -> Result<(), Box<dyn std::error::Error>> {
//!     // Initialize client (auto-discovers from environment variables)
//!     let client = UncloseAI::new(None).await?;
//!
//!     // Non-streaming chat
//!     let response = client.chat(
//!         "auto",
//!         vec![ChatMessage::user("Hello!")],
//!         None
//!     ).await?;
//!
//!     println!("{}", response.choices[0].message.content);
//!     Ok(())
//! }
//! ```

use reqwest::{Client, Response};
use serde::{Deserialize, Serialize};
use std::env;
use std::error::Error as StdError;
use std::fmt;
use futures_util::StreamExt;

/// Custom error types for UncloseAI
#[derive(Debug)]
pub enum UncloseAIError {
    /// Network connection errors
    ConnectionError(String),
    /// Requested model not available
    ModelNotFoundError(String),
    /// Errors during streaming
    StreamingError(String),
    /// General API errors
    ApiError(String),
}

impl fmt::Display for UncloseAIError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            UncloseAIError::ConnectionError(msg) => write!(f, "Connection error: {}", msg),
            UncloseAIError::ModelNotFoundError(msg) => write!(f, "Model not found: {}", msg),
            UncloseAIError::StreamingError(msg) => write!(f, "Streaming error: {}", msg),
            UncloseAIError::ApiError(msg) => write!(f, "API error: {}", msg),
        }
    }
}

impl StdError for UncloseAIError {}

impl From<reqwest::Error> for UncloseAIError {
    fn from(err: reqwest::Error) -> Self {
        UncloseAIError::ConnectionError(err.to_string())
    }
}

impl From<serde_json::Error> for UncloseAIError {
    fn from(err: serde_json::Error) -> Self {
        UncloseAIError::ApiError(err.to_string())
    }
}

/// Information about a discovered model
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ModelInfo {
    /// Model ID
    pub id: String,
    /// Endpoint URL
    pub endpoint: String,
    /// Maximum context length
    pub max_tokens: u32,
}

/// Chat message with role and content
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ChatMessage {
    /// Role: system, user, or assistant
    pub role: String,
    /// Message content
    pub content: String,
}

impl ChatMessage {
    /// Create a system message
    pub fn system(content: impl Into<String>) -> Self {
        Self {
            role: "system".to_string(),
            content: content.into(),
        }
    }

    /// Create a user message
    pub fn user(content: impl Into<String>) -> Self {
        Self {
            role: "user".to_string(),
            content: content.into(),
        }
    }

    /// Create an assistant message
    pub fn assistant(content: impl Into<String>) -> Self {
        Self {
            role: "assistant".to_string(),
            content: content.into(),
        }
    }
}

/// Chat completion request
#[derive(Clone, Debug, Serialize)]
struct ChatRequest {
    model: String,
    messages: Vec<ChatMessage>,
    #[serde(skip_serializing_if = "Option::is_none")]
    max_tokens: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    temperature: Option<f32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    top_p: Option<f32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    stream: Option<bool>,
}

/// Chat completion response
#[derive(Clone, Debug, Deserialize)]
pub struct ChatResponse {
    pub id: String,
    pub choices: Vec<ChatChoice>,
}

/// Individual choice in chat response
#[derive(Clone, Debug, Deserialize)]
pub struct ChatChoice {
    pub index: u32,
    pub message: ChatMessage,
    pub finish_reason: Option<String>,
}

/// Streaming chat chunk
#[derive(Clone, Debug, Deserialize)]
pub struct ChatChunk {
    pub id: String,
    pub choices: Vec<ChatChunkChoice>,
}

/// Individual choice in streaming chunk
#[derive(Clone, Debug, Deserialize)]
pub struct ChatChunkChoice {
    pub index: u32,
    pub delta: ChatDelta,
    pub finish_reason: Option<String>,
}

/// Delta content in streaming chunk
#[derive(Clone, Debug, Deserialize)]
pub struct ChatDelta {
    #[serde(default)]
    pub role: Option<String>,
    #[serde(default)]
    pub content: Option<String>,
}

/// TTS request
#[derive(Clone, Debug, Serialize)]
struct TTSRequest {
    model: String,
    voice: String,
    input: String,
}

/// Models list response
#[derive(Deserialize)]
struct ModelsResponse {
    data: Vec<ModelData>,
}

/// Individual model data
#[derive(Deserialize)]
struct ModelData {
    id: String,
    #[serde(default)]
    max_model_len: Option<u32>,
}

/// Configuration options for UncloseAI client
#[derive(Clone, Debug)]
pub struct UncloseAIConfig {
    /// Model endpoints
    pub endpoints: Option<Vec<String>>,
    /// TTS endpoints
    pub tts_endpoints: Option<Vec<String>>,
    /// API key for authentication
    pub api_key: Option<String>,
    /// Request timeout in seconds
    pub timeout: u64,
    /// Enable debug logging
    pub debug: bool,
}

impl Default for UncloseAIConfig {
    fn default() -> Self {
        Self {
            endpoints: None,
            tts_endpoints: None,
            api_key: None,
            timeout: 30,
            debug: false,
        }
    }
}

/// Main UncloseAI client
pub struct UncloseAI {
    client: Client,
    models: Vec<ModelInfo>,
    tts_endpoints: Vec<String>,
    api_key: Option<String>,
    debug: bool,
}

impl UncloseAI {
    /// Create a new UncloseAI client
    ///
    /// # Arguments
    ///
    /// * `config` - Optional configuration. If None, uses defaults and auto-discovers from environment
    ///
    /// # Example
    ///
    /// ```no_run
    /// use uncloseai::{UncloseAI, UncloseAIConfig};
    ///
    /// #[tokio::main]
    /// async fn main() -> Result<(), Box<dyn std::error::Error>> {
    ///     // Auto-discover from environment
    ///     let client = UncloseAI::new(None).await?;
    ///
    ///     // Explicit configuration
    ///     let config = UncloseAIConfig {
    ///         endpoints: Some(vec!["https://api.example.com/v1".to_string()]),
    ///         tts_endpoints: Some(vec!["https://tts.example.com/v1".to_string()]),
    ///         api_key: Some("your-key".to_string()),
    ///         timeout: 30,
    ///         debug: true,
    ///     };
    ///     let client = UncloseAI::new(Some(config)).await?;
    ///     Ok(())
    /// }
    /// ```
    pub async fn new(config: Option<UncloseAIConfig>) -> Result<Self, UncloseAIError> {
        let config = config.unwrap_or_default();

        let client = Client::builder()
            .timeout(std::time::Duration::from_secs(config.timeout))
            .build()?;

        // Auto-discover or use provided endpoints
        let endpoints = config.endpoints.unwrap_or_else(|| Self::discover_endpoints_from_env("MODEL_ENDPOINT"));
        let tts_endpoints = config.tts_endpoints.unwrap_or_else(|| Self::discover_endpoints_from_env("TTS_ENDPOINT"));

        if config.debug {
            println!("[DEBUG] Initialized with {} model endpoint(s) and {} TTS endpoint(s)",
                     endpoints.len(), tts_endpoints.len());
        }

        // Discover models
        let models = Self::discover_models(&client, &endpoints, config.debug).await;

        Ok(Self {
            client,
            models,
            tts_endpoints,
            api_key: config.api_key,
            debug: config.debug,
        })
    }

    /// Discover endpoints from environment variables
    fn discover_endpoints_from_env(prefix: &str) -> Vec<String> {
        let mut endpoints = Vec::new();
        for i in 1..10000 {
            match env::var(format!("{}_{}", prefix, i)) {
                Ok(endpoint) => endpoints.push(endpoint),
                Err(_) => break,
            }
        }
        endpoints
    }

    /// Discover available models from endpoints
    async fn discover_models(client: &Client, endpoints: &[String], debug: bool) -> Vec<ModelInfo> {
        let mut models = Vec::new();

        for endpoint in endpoints {
            if debug {
                println!("[DEBUG] Discovering models from: {}", endpoint);
            }

            match client.get(format!("{}/models", endpoint)).send().await {
                Ok(resp) => {
                    match resp.json::<ModelsResponse>().await {
                        Ok(models_resp) => {
                            for model in models_resp.data {
                                let model_info = ModelInfo {
                                    id: model.id.clone(),
                                    endpoint: endpoint.clone(),
                                    max_tokens: model.max_model_len.unwrap_or(8192),
                                };
                                models.push(model_info);

                                if debug {
                                    println!("[DEBUG] Discovered: {}", model.id);
                                }
                            }
                        }
                        Err(e) => {
                            if debug {
                                println!("[DEBUG] Error parsing models from {}: {}", endpoint, e);
                            }
                        }
                    }
                }
                Err(e) => {
                    if debug {
                        println!("[DEBUG] Error discovering from {}: {}", endpoint, e);
                    }
                }
            }
        }

        models
    }

    /// List all discovered models
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use uncloseai::UncloseAI;
    /// # #[tokio::main]
    /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// let client = UncloseAI::new(None).await?;
    /// let models = client.list_models();
    /// for model in models {
    ///     println!("{} - {}", model.id, model.max_tokens);
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub fn list_models(&self) -> &[ModelInfo] {
        &self.models
    }

    /// Resolve model ID to ModelInfo
    fn resolve_model(&self, model: &str) -> Result<&ModelInfo, UncloseAIError> {
        if self.models.is_empty() {
            return Err(UncloseAIError::ModelNotFoundError(
                "No models available. Check endpoints.".to_string()
            ));
        }

        if model == "auto" {
            return Ok(&self.models[0]);
        }

        self.models
            .iter()
            .find(|m| m.id == model)
            .ok_or_else(|| UncloseAIError::ModelNotFoundError(
                format!("Model '{}' not found in registry", model)
            ))
    }

    /// Send a chat completion request (non-streaming)
    ///
    /// # Arguments
    ///
    /// * `model` - Model ID or "auto" for first available
    /// * `messages` - Vector of ChatMessage objects
    /// * `options` - Optional ChatOptions for temperature, max_tokens, etc.
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use uncloseai::{UncloseAI, ChatMessage};
    /// # #[tokio::main]
    /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// let client = UncloseAI::new(None).await?;
    /// let response = client.chat(
    ///     "auto",
    ///     vec![
    ///         ChatMessage::system("You are a helpful assistant."),
    ///         ChatMessage::user("Hello!")
    ///     ],
    ///     Some(ChatOptions {
    ///         max_tokens: Some(100),
    ///         temperature: Some(0.7),
    ///         ..Default::default()
    ///     })
    /// ).await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn chat(
        &self,
        model: &str,
        messages: Vec<ChatMessage>,
        options: Option<ChatOptions>,
    ) -> Result<ChatResponse, UncloseAIError> {
        let model_info = self.resolve_model(model)?;
        let options = options.unwrap_or_default();

        let request = ChatRequest {
            model: model_info.id.clone(),
            messages,
            max_tokens: options.max_tokens,
            temperature: options.temperature,
            top_p: options.top_p,
            stream: None,
        };

        let url = format!("{}/chat/completions", model_info.endpoint);

        let mut req = self.client.post(&url).json(&request);
        if let Some(api_key) = &self.api_key {
            req = req.bearer_auth(api_key);
        }

        let response = req.send().await?;
        let chat_response = response.json::<ChatResponse>().await?;

        Ok(chat_response)
    }

    /// Send a streaming chat completion request
    ///
    /// Returns a stream of ChatChunk objects
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use uncloseai::{UncloseAI, ChatMessage};
    /// # use futures_util::StreamExt;
    /// # #[tokio::main]
    /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// let client = UncloseAI::new(None).await?;
    /// let mut stream = client.chat_stream(
    ///     "auto",
    ///     vec![ChatMessage::user("Write a haiku")],
    ///     None
    /// ).await?;
    ///
    /// while let Some(chunk) = stream.next().await {
    ///     if let Ok(chunk) = chunk {
    ///         if let Some(content) = &chunk.choices[0].delta.content {
    ///             print!("{}", content);
    ///         }
    ///     }
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub async fn chat_stream(
        &self,
        model: &str,
        messages: Vec<ChatMessage>,
        options: Option<ChatOptions>,
    ) -> Result<impl futures_core::Stream<Item = Result<ChatChunk, UncloseAIError>>, UncloseAIError> {
        let model_info = self.resolve_model(model)?;
        let options = options.unwrap_or_default();

        let request = ChatRequest {
            model: model_info.id.clone(),
            messages,
            max_tokens: options.max_tokens,
            temperature: options.temperature,
            top_p: options.top_p,
            stream: Some(true),
        };

        let url = format!("{}/chat/completions", model_info.endpoint);

        let mut req = self.client.post(&url).json(&request);
        if let Some(api_key) = &self.api_key {
            req = req.bearer_auth(api_key);
        }

        let response = req.send().await?;

        // Parse SSE stream
        let stream = response.bytes_stream().map(|result| {
            match result {
                Ok(bytes) => {
                    let text = String::from_utf8_lossy(&bytes);
                    // Parse SSE format
                    for line in text.lines() {
                        if let Some(data) = line.strip_prefix("data: ") {
                            if data.trim() == "[DONE]" {
                                continue;
                            }
                            match serde_json::from_str::<ChatChunk>(data) {
                                Ok(chunk) => return Ok(chunk),
                                Err(_) => continue,
                            }
                        }
                    }
                    Err(UncloseAIError::StreamingError("No valid chunk in response".to_string()))
                }
                Err(e) => Err(UncloseAIError::StreamingError(e.to_string())),
            }
        });

        Ok(stream)
    }

    /// Generate speech from text
    ///
    /// # Arguments
    ///
    /// * `text` - Text to convert to speech
    /// * `voice` - Voice to use (alloy, echo, fable, onyx, nova, shimmer)
    /// * `model` - TTS model (tts-1 or tts-1-hd)
    ///
    /// # Example
    ///
    /// ```no_run
    /// # use uncloseai::UncloseAI;
    /// # use std::fs::File;
    /// # use std::io::Write;
    /// # #[tokio::main]
    /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// let client = UncloseAI::new(None).await?;
    /// let audio = client.tts("Hello!", "alloy", "tts-1").await?;
    /// let mut file = File::create("speech.mp3")?;
    /// file.write_all(&audio)?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn tts(
        &self,
        text: &str,
        voice: &str,
        model: &str,
    ) -> Result<Vec<u8>, UncloseAIError> {
        if self.tts_endpoints.is_empty() {
            return Err(UncloseAIError::ApiError("No TTS endpoints available".to_string()));
        }

        let endpoint = &self.tts_endpoints[0];
        let url = format!("{}/audio/speech", endpoint);

        let request = TTSRequest {
            model: model.to_string(),
            voice: voice.to_string(),
            input: text.to_string(),
        };

        let mut req = self.client.post(&url).json(&request);
        if let Some(api_key) = &self.api_key {
            req = req.bearer_auth(api_key);
        }

        let response = req.send().await?;
        let bytes = response.bytes().await?;

        Ok(bytes.to_vec())
    }
}

/// Options for chat completions
#[derive(Clone, Debug, Default)]
pub struct ChatOptions {
    /// Maximum tokens to generate
    pub max_tokens: Option<u32>,
    /// Sampling temperature (0.0 - 2.0)
    pub temperature: Option<f32>,
    /// Nucleus sampling parameter (0.0 - 1.0)
    pub top_p: Option<f32>,
}
