zesdex/src/service/provider.rs

use std::time::Duration;
use anyhow::Result;

use crate::app::runtime::stream::{SseParser, StreamEvent};
use crate::app::runtime::stream::turn::StreamedTurn;
use crate::dto::chat::message::ChatMessage;
use crate::dto::provider::request::{ChatRequest, StreamOptions, ToolDef};

const DEFAULT_BASE_URL: &str = "https://opencode.ai/zen/v1";
const DEFAULT_MODEL: &str = "deepseek-v4-flash-free";
const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
const REQUEST_TIMEOUT: Duration = Duration::from_secs(60);

pub struct LlmClient {
    pub client: reqwest::blocking::Client,
    pub api_key: String,
    pub base_url: String,
    pub model: String,
}

impl LlmClient {
    pub fn new(api_key: String, model: String, base_url: Option<String>) -> Self {
        let model = if model.is_empty() {
            DEFAULT_MODEL.to_string()
        } else {
            model
        };
        let client = reqwest::blocking::Client::builder()
            .timeout(REQUEST_TIMEOUT)
            .connect_timeout(CONNECT_TIMEOUT)
            .build()
            .unwrap_or_else(|_| reqwest::blocking::Client::new());
        LlmClient {
            client,
            api_key,
            base_url: base_url.filter(|s| !s.is_empty()).unwrap_or_else(|| DEFAULT_BASE_URL.to_string()),
            model,
        }
    }

    pub fn chat(&self, messages: &[ChatMessage]) -> Result<String> {
        let response = self.chat_with_tools(messages, None)?;
        Ok(response.content.unwrap_or_default())
    }

    pub fn chat_with_tools(
        &self,
        messages: &[ChatMessage],
        tools: Option<Vec<ToolDef>>,
    ) -> Result<ChatMessage> {
        let req = ChatRequest {
            model: self.model.clone(),
            messages: messages.to_vec(),
            max_tokens: Some(4096),
            temperature: Some(0.7),
            tools,
            stream: Some(false),
            top_p: None,
            stop: None,
            stream_options: None,
        };

        let url = format!("{}/chat/completions", self.base_url);
        let max_retries = 10;
        let mut attempt = 0;

        loop {
            attempt += 1;

            let mut http_req = self.client
                .post(&url)
                .header("Content-Type", "application/json");

            if !self.api_key.is_empty() {
                http_req = http_req.header("Authorization", format!("Bearer {}", self.api_key));
            }

            let result = (|| -> Result<ChatMessage> {
                let resp = http_req.json(&req).send().map_err(|e| {
                    if e.is_timeout() {
                        anyhow::anyhow!("API request timed out after {:?}. Check your network or try again.", REQUEST_TIMEOUT)
                    } else if e.is_connect() {
                        anyhow::anyhow!("Could not connect to {}. Is the URL correct and is the service reachable?", self.base_url)
                    } else {
                        anyhow::anyhow!("API request failed: {}", e)
                    }
                })?;

                if !resp.status().is_success() {
                    let status = resp.status();
                    let body = resp.text().unwrap_or_default();
                    anyhow::bail!("API error {} from {}: {}", status, self.base_url, body);
                }

                let data: crate::dto::provider::response::ChatResponse = resp.json()?;
                let message = data
                    .choices
                    .into_iter()
                    .next()
                    .map(|c| c.message)
                    .ok_or_else(|| anyhow::anyhow!("API response had no choices"))?;
                Ok(message)
            })();

            match result {
                Ok(msg) => return Ok(msg),
                Err(e) => {
                    if attempt >= max_retries {
                        return Err(e);
                    }
                    eprintln!("Warning: {}. Retrying {}/{}...", e, attempt, max_retries);
                    std::thread::sleep(Duration::from_secs(2));
                }
            }
        }
    }

    /// Streaming variant of `chat_with_tools`. Feeds SSE chunks into an `SseParser` /
    /// `StreamedTurn` and invokes `on_event` for every parsed `StreamEvent` as it arrives,
    /// so the caller can push incremental UI updates in real time. Returns the fully
    /// assembled assistant message plus token usage (prompt, completion) if the server
    /// reported it. Retries the whole request only if no event has been observed yet
    /// (once tokens start arriving, a partial turn cannot be safely replayed).
    pub fn chat_with_tools_streaming(
        &self,
        messages: &[ChatMessage],
        tools: Option<Vec<ToolDef>>,
        temperature: Option<f32>,
        max_tokens: Option<u32>,
        mut on_event: impl FnMut(&StreamEvent),
    ) -> Result<(ChatMessage, Option<(u64, u64)>)> {
        let req = ChatRequest {
            model: self.model.clone(),
            messages: messages.to_vec(),
            max_tokens: Some(max_tokens.unwrap_or(4096)),
            temperature: Some(temperature.unwrap_or(0.7)),
            tools,
            stream: Some(true),
            top_p: None,
            stop: None,
            stream_options: Some(StreamOptions { include_usage: true }),
        };

        let url = format!("{}/chat/completions", self.base_url);
        let max_retries = 10;
        let mut attempt = 0;
        let mut started = false;

        loop {
            attempt += 1;
            let mut wrapped = |event: &StreamEvent| {
                started = true;
                on_event(event);
            };
            match self.try_stream_once(&req, &url, &mut wrapped) {
                Ok(result) => return Ok(result),
                Err(e) => {
                    if started || attempt >= max_retries {
                        return Err(e);
                    }
                    eprintln!("Warning: {}. Retrying {}/{}...", e, attempt, max_retries);
                    std::thread::sleep(Duration::from_secs(2));
                }
            }
        }
    }

    fn try_stream_once(
        &self,
        req: &ChatRequest,
        url: &str,
        on_event: &mut dyn FnMut(&StreamEvent),
    ) -> Result<(ChatMessage, Option<(u64, u64)>)> {
        use std::io::Read;

        let mut http_req = self.client
            .post(url)
            .header("Content-Type", "application/json");
        if !self.api_key.is_empty() {
            http_req = http_req.header("Authorization", format!("Bearer {}", self.api_key));
        }

        let resp = http_req.json(req).send().map_err(|e| {
            if e.is_timeout() {
                anyhow::anyhow!("API request timed out after {:?}. Check your network or try again.", REQUEST_TIMEOUT)
            } else if e.is_connect() {
                anyhow::anyhow!("Could not connect to {}. Is the URL correct and is the service reachable?", self.base_url)
            } else {
                anyhow::anyhow!("API request failed: {}", e)
            }
        })?;

        if !resp.status().is_success() {
            let status = resp.status();
            let body = resp.text().unwrap_or_default();
            anyhow::bail!("API error {} from {}: {}", status, self.base_url, body);
        }

        let mut turn = StreamedTurn::new();
        let mut usage: Option<(u64, u64)> = None;
        let mut parser = SseParser::new();
        let mut reader = resp;
        let mut byte_buf: Vec<u8> = Vec::new();
        let mut chunk_buf = [0u8; 4096];

        loop {
            let n = reader.read(&mut chunk_buf)
                .map_err(|e| anyhow::anyhow!("stream read error: {}", e))?;
            if n == 0 {
                break;
            }
            byte_buf.extend_from_slice(&chunk_buf[..n]);
            let valid_len = match std::str::from_utf8(&byte_buf) {
                Ok(s) => s.len(),
                Err(e) => e.valid_up_to(),
            };
            if valid_len == 0 {
                continue;
            }
            let text = String::from_utf8_lossy(&byte_buf[..valid_len]).into_owned();
            byte_buf.drain(..valid_len);

            for event in parser.feed(&text) {
                on_event(&event);
                match &event {
                    StreamEvent::Usage { prompt_tokens, completion_tokens, .. } => {
                        usage = Some((*prompt_tokens, *completion_tokens));
                    }
                    StreamEvent::Error(msg) => {
                        anyhow::bail!("stream error: {}", msg);
                    }
                    StreamEvent::Done => {
                        turn.apply_event(&event);
                        return Ok((turn.build_assistant_message(), usage));
                    }
                    _ => turn.apply_event(&event),
                }
            }
        }

        turn.is_complete = true;
        Ok((turn.build_assistant_message(), usage))
    }
}