//! Inference on co/core (https://cocore.dev/docs/inference): OpenAI-compatible chat completions at
//! https://cocore.dev/api/v1 through the host's com.fedfork.exec.http.fetch. The module sends no
//! credentials: the host adds them for that origin, a service-auth token it mints from the user's PDS
//! (aud did:web:cocore.dev, lxm dev.cocore.inference.dispatch) or a co/core API key.
use abi::{Call, Event, Handle, Run};
use cocore::{Args, field};
use serde_json::{Value, json};

pub const API: &str = "https://cocore.dev/api/v1";
const MAX_TOKENS: u64 = 1024;
const ONBOARDING: &str = "this account has not connected co/core yet: sign in at https://cocore.dev once, then try again (onboarding_required)";

fn number(args: &Args, key: &str, default: u64) -> Result<u64, String> {
    let n = field(args, key).map_or(Ok(default), |v| v.parse().map_err(|_| format!("{key} must be a whole number")))?;
    if n == 0 {
        return Err(format!("{key} must be at least 1"));
    }
    Ok(n)
}

/// A JSON field given as itself or as its text (form fields are strings).
fn json_field(args: &Args, key: &str) -> Result<Option<Value>, String> {
    match args.get(key) {
        None | Some(Value::Null) => Ok(None),
        Some(Value::String(s)) if s.trim().is_empty() => Ok(None),
        Some(Value::String(s)) => serde_json::from_str(s).map(Some).map_err(|e| format!("{key}: {e}")),
        Some(v) => Ok(Some(v.clone())),
    }
}

fn truthy(args: &Args, key: &str) -> bool {
    matches!(field(args, key).as_deref(), Some("true" | "1" | "yes" | "on"))
}

/// The route: chat/completions, or its private (friends), verified (attested) or probono variant.
fn url(args: &Args) -> Result<String, String> {
    match field(args, "route").as_deref() {
        None | Some("open") => Ok(format!("{API}/chat/completions")),
        Some(r @ ("private" | "verified" | "probono")) => Ok(format!("{API}/{r}/chat/completions")),
        Some(r) => Err(format!("route {r}: open, private, verified or probono")),
    }
}

fn request(args: &Args) -> Result<Value, String> {
    let mut messages = match json_field(args, "messages")? {
        None => vec![],
        Some(Value::Array(m)) => m,
        Some(_) => return Err("messages must be a list of chat messages".into()),
    };
    if let Some(system) = field(args, "system") {
        messages.insert(0, json!({ "role": "system", "content": system }));
    }
    if let Some(prompt) = field(args, "prompt") {
        messages.push(json!({ "role": "user", "content": prompt }));
    }
    if !messages.iter().any(|m| m["role"] == "user") {
        return Err("prompt (or messages with a user message) is required".into());
    }
    let mut body = json!({
        "model": field(args, "model").unwrap_or_else(|| "stub".into()),
        "messages": messages,
        "max_tokens": number(args, "maxTokens", MAX_TOKENS)?,
        "stream": truthy(args, "stream"),
    });
    for (key, wire) in [("country", "country"), ("minProviderVersion", "min_provider_version"), ("minTrust", "min_trust")] {
        if let Some(v) = field(args, key) {
            body[wire] = json!(v);
        }
    }
    Ok(body)
}

/// A non-2xx answer as the run's error: co/core's { error: { type, code, message } } envelope.
fn http_error(output: &Value) -> Option<String> {
    let status = output["status"].as_u64().unwrap_or(0);
    if (200..300).contains(&status) {
        return None;
    }
    let body: Value = serde_json::from_str(output["body"].as_str().unwrap_or_default()).unwrap_or(Value::Null);
    let error = &body["error"];
    if error["code"] == "onboarding_required" {
        return Some(format!("co/core: {ONBOARDING}"));
    }
    let kind = error["code"].as_str().or(error["type"].as_str()).unwrap_or("error");
    let message = error["message"].as_str().map_or_else(|| output["body"].as_str().unwrap_or_default().chars().take(300).collect(), String::from);
    Some(format!("co/core HTTP {status} {kind}: {message}"))
}

/// The answer: a chat.completion, or an event stream of chat.completion.chunk folded into one.
fn answer(output: &Value) -> Result<Value, String> {
    if let Some(e) = http_error(output) {
        return Err(e);
    }
    let body = output["body"].as_str().unwrap_or_default();
    let streamed = output["headers"]["content-type"].as_str().is_some_and(|t| t.starts_with("text/event-stream"));
    if !streamed {
        let c: Value = serde_json::from_str(body).map_err(|e| format!("co/core answered no JSON: {e}"))?;
        let choice = &c["choices"][0];
        return Ok(json!({
            "text": choice["message"]["content"].as_str().unwrap_or_default(),
            "finishReason": choice["finish_reason"],
            "toolCalls": choice["message"].get("tool_calls").cloned().unwrap_or(json!([])),
            "model": c["model"],
            "usage": c["usage"],
            "cocore": c.get("x_cocore").cloned().unwrap_or(Value::Null),
        }));
    }
    let (mut text, mut finish, mut usage, mut model, mut cocore) = (String::new(), Value::Null, Value::Null, Value::Null, Value::Null);
    let mut tools: Vec<Value> = vec![];
    for data in body.replace("\r\n", "\n").split("\n\n").filter_map(|frame| frame.lines().find_map(|l| l.strip_prefix("data:")).map(str::trim)) {
        if data == "[DONE]" {
            continue;
        }
        let Ok(chunk) = serde_json::from_str::<Value>(data) else { continue };
        if chunk.get("error").is_some() {
            return Err(http_error(&json!({ "status": 502, "body": data })).unwrap_or_default());
        }
        let choice = &chunk["choices"][0];
        text.push_str(choice["delta"]["content"].as_str().unwrap_or_default());
        // tool calls stream in pieces by index: the name once, the arguments appended
        for part in choice["delta"]["tool_calls"].as_array().into_iter().flatten() {
            let i = part["index"].as_u64().unwrap_or(0) as usize;
            while tools.len() <= i {
                tools.push(json!({ "type": "function", "function": { "name": "", "arguments": "" } }));
            }
            let t = &mut tools[i];
            if let Some(id) = part["id"].as_str() {
                t["id"] = json!(id);
            }
            for key in ["name", "arguments"] {
                let piece = part["function"][key].as_str().unwrap_or_default();
                t["function"][key] = json!(format!("{}{piece}", t["function"][key].as_str().unwrap_or_default()));
            }
        }
        let take = |slot: &mut Value, v: &Value| {
            if !v.is_null() {
                *slot = v.clone();
            }
        };
        take(&mut finish, &choice["finish_reason"]);
        take(&mut usage, &chunk["usage"]);
        take(&mut model, &chunk["model"]);
        take(&mut cocore, &chunk["x_cocore"]);
    }
    Ok(json!({ "text": text, "finishReason": finish, "toolCalls": tools, "model": model, "usage": usage, "cocore": cocore }))
}

fn chat_handler(run: Run<Args>) -> Result<Handle, String> {
    Ok(match run.event {
        Event::Start => {
            let mut body = request(&run.args)?;
            let headers = json!({ "content-type": "application/json", "accept": if body["stream"] == true { "text/event-stream" } else { "application/json" } });
            let url = url(&run.args)?;
            if let Some(Value::Object(extra)) = json_field(&run.args, "extra")? {
                body.as_object_mut().unwrap().extend(extra);
            }
            Handle::new().call(Call::http("chat", "POST", &url, headers, Some(body.to_string())))
        }
        Event::Result { result, .. } => Handle::new().result(answer(&result.into_result()?)?).done(),
        _ => Handle::new(),
    })
}

/// The public model directory (view: openai, the default; summary or directory).
fn models_handler(run: Run<Args>) -> Result<Handle, String> {
    Ok(match run.event {
        Event::Start => {
            let url = field(&run.args, "view").map_or_else(|| format!("{API}/models"), |v| format!("{API}/models?view={v}"));
            Handle::new().call(Call::http("models", "GET", &url, json!({}), None))
        }
        Event::Result { result, .. } => {
            let output = result.into_result()?;
            if let Some(e) = http_error(&output) {
                return Err(e);
            }
            Handle::new().result(serde_json::from_str(output["body"].as_str().unwrap_or_default()).map_err(|e| format!("co/core answered no JSON: {e}"))?).done()
        }
        _ => Handle::new(),
    })
}

abi::entries!(chat => chat_handler, models => models_handler);