toast-api 0.1.4

use anyhow::{anyhow, Result};
use brotli;
use flate2::read::{DeflateDecoder, GzDecoder};
use rquest::Response;
use serde::Deserialize;
use std::io::Read;
// For JSON parsing
use chrono::Local;
use lazy_static::lazy_static;
use regex::Regex;
use serde_json;

/// Decode response body based on Content-Encoding
pub async fn decode_body(res: Response) -> Result<Vec<u8>> {
    // Extract the content encoding header as an owned String to avoid borrowing `res`
    let encoding = res
        .headers()
        .get("Content-Encoding")
        .and_then(|v| v.to_str().ok())
        .unwrap_or("")
        .to_string();

    // Consume response to get bytes
    let bytes_result = res.bytes().await;

    // Handle possible empty body or connection errors
    let buf = match bytes_result {
        Ok(bytes) if bytes.is_empty() => {
            return Err(anyhow!("Empty response body received from server"))
        }
        Ok(bytes) => bytes.to_vec(),
        Err(e) => return Err(anyhow!("Failed to read response body: {}", e)),
    };

    // Process based on encoding
    let out = match encoding.as_str() {
        "gzip" => {
            let mut d = GzDecoder::new(&buf[..]);
            let mut v = Vec::new();
            d.read_to_end(&mut v)?;
            v
        }
        "deflate" => {
            let mut d = DeflateDecoder::new(&buf[..]);
            let mut v = Vec::new();
            d.read_to_end(&mut v)?;
            v
        }
        "br" => {
            let mut v = Vec::new();
            brotli::Decompressor::new(&buf[..], 4096).read_to_end(&mut v)?;
            v
        }
        _ => buf,
    };
    Ok(out)
}

/// Parsed line from Claude stream
#[derive(Debug, Deserialize)]
struct SendLine {
    #[serde(default)]
    completion: String,
    #[serde(default)]
    error: Option<ClaudeErrorInfo>,
}

/// Info about errors from Claude
#[derive(Debug, Deserialize)]
struct ClaudeErrorInfo {
    #[serde(default)]
    #[allow(non_snake_case)]
    r#type: String,
    #[serde(default)]
    message: String,
    #[serde(default)]
    resets_at: Option<i64>,
}

/// Parse the event-stream bytes into a completion string
pub fn parse_stream(bytes: &[u8]) -> Result<String> {
    let text = String::from_utf8_lossy(bytes);
    let mut completions = String::new();
    for line in text.lines() {
        if let Some(start) = line.find('{') {
            if let Ok(obj) = serde_json::from_str::<SendLine>(&line[start..]) {
                if let Some(err) = obj.error {
                    match err.r#type.as_str() {
                        "message_rate_limit_error" => {
                            let secs = err.resets_at.unwrap_or_default() - Local::now().timestamp();
                            return Err(anyhow!(
                                "Rate-limit: wait {} seconds – {}",
                                secs,
                                err.message
                            ));
                        }
                        "overloaded_error" => {
                            return Err(anyhow!("Server overloaded: {}", err.message));
                        }
                        _ => return Err(anyhow!("Claude error: {}", err.message)),
                    }
                }
                completions.push_str(&obj.completion);
            }
        }
    }
    Ok(completions.trim().to_owned())
}

lazy_static! {
    pub static ref READ_RE: Regex = Regex::new(r"(?i)^#?\s*read_file\s+(.+)").unwrap();
    pub static ref EXEC_RE: Regex = Regex::new(r"(?i)^#?\s*exec\s+(.+)").unwrap();
}

/// Format a line with ANSI colors for read_file or exec commands
pub fn format_tool_line(line: &str) -> String {
    if READ_RE.is_match(line) {
        format!("\x1b[34m{}\x1b[0m", line.trim_start_matches('#').trim())
    } else if EXEC_RE.is_match(line) {
        format!(
            "\x1b[38;5;208m{}\x1b[0m",
            line.trim_start_matches('#').trim()
        )
    } else {
        line.to_owned()
    }
}

/// Extract read_file and exec commands from a Claude answer
pub fn extract_commands(answer: &str) -> (Vec<String>, Vec<String>) {
    let mut reads = Vec::new();
    let mut execs = Vec::new();
    let mut lines = answer.lines().peekable();
    while lines.peek().is_some() {
        let line = lines.next().unwrap();
        if let Some(caps) = READ_RE.captures(line) {
            reads.extend(caps[1].split_whitespace().map(|s| s.to_string()));
        } else if let Some(caps) = EXEC_RE.captures(line) {
            let first = caps[1].to_string();
            if let Some(hd) = first.split("<<").nth(1) {
                let delim = hd.trim().trim_matches(|c| c == '\'' || c == '"');
                let mut block = vec![first.clone()];
                for l in lines.by_ref() {
                    block.push(l.to_string());
                    if l.trim() == delim {
                        break;
                    }
                }
                execs.push(block.join("\n"));
            } else {
                execs.push(first);
            }
        }
    }
    (reads, execs)
}

/// Apply ANSI formatting to an answer
pub fn prettify(answer: &str) -> String {
    answer
        .lines()
        .map(format_tool_line)
        .collect::<Vec<_>>()
        .join("\n")
}