settled-reach/server/sr-voice/src/server.rs

use std::time::Instant;

use tiny_http::{Header, Method, Response, Server};

use crate::inference::InferenceEngine;
use crate::prompt::PromptPayload;

const MAX_TOKENS: u32 = 64;
const TEMPERATURE: f32 = 0.7;
const TOP_P: f32 = 0.9;

#[derive(serde::Deserialize)]
struct GenerateRequest {
    prompt: String,
    seed: Option<u64>,
}

pub fn run_server(
    engine: InferenceEngine,
    port: u16,
    model_name: &str,
) -> Result<(), Box<dyn std::error::Error>> {
    let addr = format!("127.0.0.1:{}", port);
    let server = Server::http(&addr)
        .map_err(|e| format!("Failed to bind {}: {}", addr, e))?;

    let start = Instant::now();
    eprintln!("sr-voice server ready on http://{}", addr);
    eprintln!("  model: {}", model_name);
    eprintln!("  POST /generate  POST /batch  GET /health");

    for request in server.incoming_requests() {
        let path = request.url().to_string();
        let method = request.method().clone();

        match (method, path.as_str()) {
            (Method::Get, "/health") => {
                let body = serde_json::json!({
                    "status": "ready",
                    "model": model_name,
                    "uptime_secs": start.elapsed().as_secs(),
                });
                respond(request, 200, &body.to_string());
            }
            (Method::Post, "/generate") => handle_generate(&engine, request),
            (Method::Post, "/batch") => handle_batch(&engine, request),
            _ => {
                respond(request, 404, &serde_json::json!({"error": "not found"}).to_string());
            }
        }
    }

    Ok(())
}

fn handle_generate(engine: &InferenceEngine, mut request: tiny_http::Request) {
    let mut body = String::new();
    if std::io::Read::read_to_string(request.as_reader(), &mut body).is_err() {
        respond(request, 400, r#"{"error":"failed to read body"}"#);
        return;
    }

    let req: GenerateRequest = match serde_json::from_str(&body) {
        Ok(r) => r,
        Err(e) => {
            let msg = serde_json::json!({"error": format!("invalid JSON: {}", e)});
            respond(request, 400, &msg.to_string());
            return;
        }
    };

    eprintln!("  generate: {} chars", req.prompt.len());
    match engine.generate(&req.prompt, MAX_TOKENS, TEMPERATURE, TOP_P, req.seed.map(|s| s as u32)) {
        Ok(result) => {
            eprintln!("  -> {} tokens, {:.1} t/s", result.tokens_generated, result.tokens_per_sec);
            respond(request, 200, &serde_json::to_string(&result).unwrap());
        }
        Err(e) => {
            let msg = serde_json::json!({"error": e.to_string()});
            respond(request, 500, &msg.to_string());
        }
    }
}

fn handle_batch(engine: &InferenceEngine, mut request: tiny_http::Request) {
    let mut body = String::new();
    if std::io::Read::read_to_string(request.as_reader(), &mut body).is_err() {
        respond(request, 400, r#"{"error":"failed to read body"}"#);
        return;
    }

    let payloads: Vec<PromptPayload> = match serde_json::from_str(&body) {
        Ok(p) => p,
        Err(e) => {
            let msg = serde_json::json!({"error": format!("invalid JSON: {}", e)});
            respond(request, 400, &msg.to_string());
            return;
        }
    };

    eprintln!("  batch: {} prompts", payloads.len());
    let mut output = String::new();
    for payload in &payloads {
        match engine.generate(&payload.prompt, MAX_TOKENS, TEMPERATURE, TOP_P, None) {
            Ok(result) => {
                eprintln!("  -> {}: {} tokens, {:.1} t/s", payload.id, result.tokens_generated, result.tokens_per_sec);
                #[derive(serde::Serialize)]
                struct BatchLine<'a> {
                    id: &'a str,
                    #[serde(flatten)]
                    result: &'a crate::inference::GenerationResult,
                }
                let line = serde_json::to_string(&BatchLine { id: &payload.id, result: &result }).unwrap();
                output.push_str(&line);
                output.push('\n');
            }
            Err(e) => {
                let line = serde_json::json!({"id": payload.id, "error": e.to_string()});
                output.push_str(&line.to_string());
                output.push('\n');
            }
        }
    }

    respond(request, 200, &output);
}

fn respond(request: tiny_http::Request, status: u16, body: &str) {
    let header = Header::from_bytes("Content-Type", "application/json").unwrap();
    let response = Response::from_string(body)
        .with_status_code(status)
        .with_header(header);
    let _ = request.respond(response);
}