zeroski 0.0.1

zero.ski CLI — @-protocol dispatch, streaming chat, and trace feed for the Zero runtime
//! `zeroski chat` — stream /api/chat to stdout.
//!
//! Server emits OpenAI-style SSE frames:
//!   data: {"content": "Hello"}
//!   data: {"done": true, "model": "...", "tokens_in": 9, "tokens_out": 1, ...}
//!   data: [DONE]
//!
//! We print `content` deltas straight to stdout (no newline between frames)
//! so it looks like a typewriter, and emit a one-line summary to stderr
//! once the `done` frame arrives.

use std::io::{self, IsTerminal, Read, Write};

use anyhow::{anyhow, Result};
use futures_util::StreamExt;
use serde_json::{json, Value};

use crate::api::Client;
use crate::config::Config;

pub async fn run(cfg: &Config, text: Vec<String>, model: Option<String>) -> Result<()> {
    let prompt = resolve_prompt(text)?;
    if prompt.trim().is_empty() {
        return Err(anyhow!("empty prompt — pass text as arg or pipe via stdin"));
    }

    let client = Client::new(cfg)?;
    let mut body = json!({
        "messages": [{ "role": "user", "content": prompt }],
    });
    if let Some(m) = model {
        body["model"] = Value::String(m);
    }

    let resp = client.post_stream("/api/chat", &body).await?;
    let mut stream = resp.bytes_stream();
    let mut buffer = String::new();
    let mut stdout = io::stdout().lock();
    let mut saw_done = false;
    let mut printed_any = false;

    while let Some(chunk) = stream.next().await {
        let bytes = chunk?;
        buffer.push_str(&String::from_utf8_lossy(&bytes));
        while let Some(idx) = buffer.find("\n\n") {
            let frame: String = buffer.drain(..idx + 2).collect();
            let payload = frame
                .lines()
                .find_map(|l| l.strip_prefix("data: ").or_else(|| l.strip_prefix("data:")))
                .unwrap_or("")
                .trim();
            if payload.is_empty() || payload == "[DONE]" {
                continue;
            }
            let parsed: Value = match serde_json::from_str(payload) {
                Ok(v) => v,
                Err(_) => continue,
            };
            if let Some(err) = parsed.get("error").and_then(Value::as_str) {
                writeln!(stdout, "\n[error] {err}").ok();
                return Err(anyhow!("server returned error in stream: {err}"));
            }
            if let Some(delta) = parsed.get("content").and_then(Value::as_str) {
                write!(stdout, "{delta}")?;
                stdout.flush()?;
                printed_any = true;
            }
            if parsed.get("done").and_then(Value::as_bool).unwrap_or(false) {
                saw_done = true;
                if printed_any {
                    writeln!(stdout)?;
                }
                let model = parsed.get("model").and_then(Value::as_str).unwrap_or("?");
                let tin = parsed.get("tokens_in").and_then(Value::as_u64).unwrap_or(0);
                let tout = parsed.get("tokens_out").and_then(Value::as_u64).unwrap_or(0);
                let cost = parsed.get("cost").and_then(Value::as_f64).unwrap_or(0.0);
                eprintln!("[{model}] in={tin} out={tout} cost=${cost:.6}");
            }
        }
    }

    if !saw_done {
        return Err(anyhow!("stream ended without [DONE]"));
    }
    Ok(())
}

fn resolve_prompt(text: Vec<String>) -> Result<String> {
    if !text.is_empty() {
        return Ok(text.join(" "));
    }
    let mut stdin = io::stdin();
    if stdin.is_terminal() {
        return Err(anyhow!("no prompt — pass text as arg or pipe via stdin"));
    }
    let mut buf = String::new();
    stdin.read_to_string(&mut buf)?;
    Ok(buf.trim().to_string())
}