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

// UncloseAI - Rust client for OpenAI-compatible APIs with streaming support

use reqwest::Client;
use serde::{Deserialize, Serialize};
use std::env;
use std::error::Error;
use std::fs::File;
use std::io::Write;
use futures_util::StreamExt;

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct ModelInfo {
    pub id: String,
    pub endpoint: String,
    pub max_tokens: usize,
}

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

#[derive(Serialize)]
struct ChatRequest {
    model: String,
    messages: Vec<ChatMessage>,
    max_tokens: usize,
    temperature: f32,
    #[serde(skip_serializing_if = "Option::is_none")]
    stream: Option<bool>,
}

#[derive(Serialize)]
struct TtsRequest {
    model: String,
    voice: String,
    input: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    response_format: Option<String>,
}

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

#[derive(Deserialize)]
struct Choice {
    #[serde(default)]
    message: Option<Message>,
    #[serde(default)]
    delta: Option<Message>,
}

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

#[derive(Deserialize)]
struct ModelData {
    id: String,
    #[serde(default)]
    max_model_len: Option<usize>,
}

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

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

impl UncloseAI {
    pub async fn new(
        model_endpoints: Option<Vec<String>>,
        tts_endpoints: Option<Vec<String>>,
        api_key: Option<String>,
        timeout: Option<u64>,
        debug: bool,
    ) -> Result<Self, Box<dyn Error>> {
        let client = Client::builder()
            .timeout(std::time::Duration::from_secs(timeout.unwrap_or(30)))
            .build()?;

        let model_ends = model_endpoints.unwrap_or_else(|| Self::discover_env_endpoints("MODEL_ENDPOINT"));
        let tts_ends = tts_endpoints.unwrap_or_else(|| Self::discover_env_endpoints("TTS_ENDPOINT"));

        if debug {
            println!("[DEBUG] Initialized with {} endpoint(s)", model_ends.len());
        }

        let models = Self::discover_models(&client, &model_ends, debug).await;

        Ok(Self {
            models,
            tts_endpoints: tts_ends,
            api_key,
            timeout: timeout.unwrap_or(30),
            debug,
            client,
        })
    }

    fn discover_env_endpoints(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
    }

    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 from: {}", endpoint);
            }

            match client.get(format!("{}/models", endpoint)).send().await {
                Ok(resp) => {
                    if let Ok(models_resp) = resp.json::<ModelsResponse>().await {
                        for model in models_resp.data {
                            // Skip permission entries
                            if model.id.starts_with("modelperm-") || model.id.starts_with("chatcmpl-") {
                                continue;
                            }

                            models.push(ModelInfo {
                                id: model.id.clone(),
                                endpoint: endpoint.clone(),
                                max_tokens: model.max_model_len.unwrap_or(8192),
                            });

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

        models
    }

    pub fn list_models(&self) -> &[ModelInfo] {
        &self.models
    }

    fn resolve_model(&self, model: Option<&str>) -> Result<&ModelInfo, Box<dyn Error>> {
        if self.models.is_empty() {
            return Err("No models available".into());
        }

        match model {
            None => Ok(&self.models[0]),
            Some(id) => self
                .models
                .iter()
                .find(|m| m.id == id)
                .ok_or_else(|| format!("Model '{}' not found", id).into()),
        }
    }

    pub async fn chat(
        &self,
        messages: Vec<ChatMessage>,
        model: Option<&str>,
        max_tokens: Option<usize>,
        temperature: Option<f32>,
    ) -> Result<ChatResponse, Box<dyn Error>> {
        let model_info = self.resolve_model(model)?;

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

        let mut req = self
            .client
            .post(format!("{}/chat/completions", model_info.endpoint))
            .json(&request);

        if let Some(api_key) = &self.api_key {
            req = req.header("Authorization", format!("Bearer {}", api_key));
        }

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

    pub async fn chat_stream<F>(
        &self,
        messages: Vec<ChatMessage>,
        model: Option<&str>,
        max_tokens: Option<usize>,
        temperature: Option<f32>,
        mut callback: F,
    ) -> Result<(), Box<dyn Error>>
    where
        F: FnMut(String),
    {
        let model_info = self.resolve_model(model)?;

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

        let mut req = self
            .client
            .post(format!("{}/chat/completions", model_info.endpoint))
            .json(&request);

        if let Some(api_key) = &self.api_key {
            req = req.header("Authorization", format!("Bearer {}", api_key));
        }

        let response = req.send().await?;
        let mut stream = response.bytes_stream();
        let mut buffer = String::new();

        while let Some(chunk_result) = stream.next().await {
            let chunk = chunk_result?;
            buffer.push_str(&String::from_utf8_lossy(&chunk));

            // Process complete lines
            while let Some(newline_pos) = buffer.find('\n') {
                let line = buffer[..newline_pos].trim().to_string();
                buffer = buffer[newline_pos + 1..].to_string();

                if line.starts_with("data: ") {
                    let data = &line[6..];
                    if data.trim() == "[DONE]" {
                        return Ok(());
                    }

                    if let Ok(parsed) = serde_json::from_str::<ChatResponse>(data) {
                        if let Some(choice) = parsed.choices.first() {
                            if let Some(delta) = &choice.delta {
                                if let Some(content) = &delta.content {
                                    if !content.is_empty() {
                                        callback(content.clone());
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }

        Ok(())
    }

    pub async fn tts(
        &self,
        text: &str,
        voice: Option<&str>,
        model: Option<&str>,
        response_format: Option<&str>,
    ) -> Result<Vec<u8>, Box<dyn Error>> {
        if self.tts_endpoints.is_empty() {
            return Err("No TTS endpoints available".into());
        }

        let endpoint = &self.tts_endpoints[0];

        let request = TtsRequest {
            model: model.unwrap_or("tts-1").to_string(),
            voice: voice.unwrap_or("alloy").to_string(),
            input: text.to_string(),
            response_format: response_format.map(|s| s.to_string()),
        };

        let mut req = self
            .client
            .post(format!("{}/audio/speech", endpoint))
            .json(&request);

        if let Some(api_key) = &self.api_key {
            req = req.header("Authorization", format!("Bearer {}", api_key));
        }

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

// Demo when run as application
#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    println!("=== UncloseAI Rust Client (with Streaming) ===\n");

    let client = UncloseAI::new(None, None, None, None, true).await?;

    if client.list_models().is_empty() {
        eprintln!("ERROR: No models discovered. Set environment variables:");
        eprintln!("  MODEL_ENDPOINT_1, MODEL_ENDPOINT_2, etc.");
        std::process::exit(1);
    }

    println!("\nDiscovered {} model(s):", client.list_models().len());
    for model in client.list_models() {
        println!("  - {} (max_tokens: {})", model.id, model.max_tokens);
    }
    println!();

    // Non-streaming chat
    println!("=== Non-Streaming Chat ===");
    match client
        .chat(
            vec![
                ChatMessage {
                    role: "system".to_string(),
                    content: "You are a helpful AI assistant.".to_string(),
                },
                ChatMessage {
                    role: "user".to_string(),
                    content: "Explain quantum computing in one sentence.".to_string(),
                },
            ],
            None,
            Some(100),
            None,
        )
        .await
    {
        Ok(response) => {
            if let Some(choice) = response.choices.first() {
                if let Some(message) = &choice.message {
                    if let Some(content) = &message.content {
                        println!("Response: {}\n", content);
                    }
                }
            }
        }
        Err(e) => eprintln!("Error: {}\n", e),
    }

    // Streaming chat
    println!("=== Streaming Chat ===");
    let model_id = if client.list_models().len() > 1 {
        Some(client.list_models()[1].id.as_str())
    } else {
        None
    };
    let model_name = model_id.unwrap_or(&client.list_models()[0].id);
    println!("Model: {}", model_name);
    print!("Response: ");
    std::io::stdout().flush()?;

    match client
        .chat_stream(
            vec![
                ChatMessage {
                    role: "system".to_string(),
                    content: "You are a coding assistant.".to_string(),
                },
                ChatMessage {
                    role: "user".to_string(),
                    content: "Write a Rust function to check if a number is prime".to_string(),
                },
            ],
            model_id,
            Some(200),
            None,
            |content| {
                print!("{}", content);
                std::io::stdout().flush().ok();
            },
        )
        .await
    {
        Ok(_) => println!("\n"),
        Err(e) => eprintln!("\nError: {}\n", e),
    }

    // TTS
    if !client.tts_endpoints.is_empty() {
        println!("=== TTS Speech Generation ===");
        match client
            .tts(
                "Hello from UncloseAI Rust client! This demonstrates streaming support.",
                None,
                None,
                None,
            )
            .await
        {
            Ok(audio_data) => {
                let mut file = File::create("speech.mp3")?;
                file.write_all(&audio_data)?;
                println!("[OK] Speech file created: speech.mp3 ({} bytes)\n", audio_data.len());
            }
            Err(e) => eprintln!("[ERROR] TTS Error: {}\n", e),
        }
    }

    println!("=== Examples Complete ===");
    Ok(())
}
