//! AI assistants (hub spec §13.4, #97): chat-completion endpoints the hub

//! may ask to review a proposal or draft its description. Assistants are

//! kept at three levels, site, organisation and repository, and combined

//! downwards: a repository sees its organisation's and the site's beside

//! its own, the innermost of a name wins, and one disabled at an outer

//! level is off for everything inside it. Each level adds a prompt; the

//! prompt in force for a repository is the site's, the organisation's and

//! the repository's own, each under a heading. Tokens are sealed with

//! ChaCha20-Poly1305 under `keys/assistants.key`, shown masked, and never

//! leave the server.

use crate::db::{now_ms, Pool};

use crate::Shared;

use rand::RngCore;

use ring::aead;

use sqlx::FromRow;

use std::path::Path;

use std::time::Duration;

/// The request shapes an assistant may speak.

pub const APIS: [&str; 2] = ["chat-completions", "anthropic-messages"];

/// Characters of context (prompt, description, diffs) a call may carry

/// unless the assistant says otherwise.

pub const DEFAULT_BUDGET: i64 = 120_000;

/// How long one call may take.

pub const CALL_TIMEOUT: Duration = Duration::from_secs(60);

/// The most an answer is kept.

const MAX_ANSWER: usize = 64 * 1024;

pub use crate::levels::Scope;

#[derive(Clone, Debug, FromRow)]

pub struct Assistant {

    pub id: i64,

    pub scope: String,

    pub scope_id: i64,

    pub name: String,

    pub endpoint: String,

    /// One of [`APIS`].

    pub api: String,

    pub model: String,

    /// Extra request headers, one `Name: value` per line.

    pub headers: String,

    /// The sealed token; empty when the endpoint needs none.

    pub token: Vec<u8>,

    pub enabled: bool,

    pub budget: i64,

    pub created_ms: i64,

    pub updated_ms: i64,

}

impl Assistant {

    pub fn scope(&self) -> Scope {

        match self.scope.as_str() {

            "org" => Scope::Org(self.scope_id),

            "repo" => Scope::Repo(self.scope_id),

            _ => Scope::Site,

        }

    }

    pub fn has_token(&self) -> bool {

        !self.token.is_empty()

    }

    /// The extra headers as pairs; malformed lines are skipped.

    pub fn header_list(&self) -> Vec<(String, String)> {

        parse_headers(&self.headers).unwrap_or_default()

    }

}

/// What a settings form supplies for an assistant.

#[derive(Clone, Debug, Default)]

pub struct Fields {

    pub name: String,

    pub endpoint: String,

    pub api: String,

    pub model: String,

    pub headers: String,

    pub enabled: bool,

    pub budget: i64,

}

/// `Name: value` lines to pairs: `Err(line)` names the first bad one.

pub fn parse_headers(text: &str) -> Result<Vec<(String, String)>, String> {

    let mut out = vec![];

    for line in text.lines().map(str::trim).filter(|l| !l.is_empty()) {

        let Some((name, value)) = line.split_once(':') else { return Err(line.to_string()) };

        let (name, value) = (name.trim(), value.trim());

        let name_ok = !name.is_empty() && name.chars().all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_');

        if !name_ok || value.chars().any(|c| c.is_control()) {

            return Err(line.to_string());

        }

        out.push((name.to_string(), value.to_string()));

    }

    Ok(out)

}

// -------------------------------------------------------------- sealing

/// The key that seals tokens, kept in `<data>/keys/assistants.key`.

#[derive(Clone)]

pub struct SecretKey([u8; 32]);

impl SecretKey {

    pub fn load_or_create(keys: &Path) -> anyhow::Result<SecretKey> {

        std::fs::create_dir_all(keys)?;

        let path = keys.join("assistants.key");

        if let Ok(bytes) = std::fs::read(&path) {

            if bytes.len() == 32 {

                return Ok(SecretKey(bytes.try_into().unwrap()));

            }

        }

        let mut k = [0u8; 32];

        rand::rngs::OsRng.fill_bytes(&mut k);

        std::fs::write(&path, k)?;

        #[cfg(unix)]

        {

            use std::os::unix::fs::PermissionsExt;

            std::fs::set_permissions(&path, std::fs::Permissions::from_mode(0o600))?;

        }

        Ok(SecretKey(k))

    }

    fn key(&self) -> aead::LessSafeKey {

        aead::LessSafeKey::new(aead::UnboundKey::new(&aead::CHACHA20_POLY1305, &self.0).expect("key length"))

    }

    /// Seal a token: a random nonce, then the ciphertext and tag. An

    /// empty token seals to nothing, meaning "no token".

    pub fn seal(&self, plain: &str) -> Vec<u8> {

        if plain.is_empty() {

            return vec![];

        }

        let mut nonce = [0u8; 12];

        rand::rngs::OsRng.fill_bytes(&mut nonce);

        let mut out = plain.as_bytes().to_vec();

        self.key().seal_in_place_append_tag(aead::Nonce::assume_unique_for_key(nonce), aead::Aad::empty(), &mut out).expect("seal");

        let mut sealed = nonce.to_vec();

        sealed.extend(out);

        sealed

    }

    /// The token back, or `None` when the bytes were sealed under another

    /// key or damaged. Nothing opens to an empty token.

    pub fn open(&self, sealed: &[u8]) -> Option<String> {

        if sealed.is_empty() {

            return Some(String::new());

        }

        let (nonce, rest) = sealed.split_at_checked(12)?;

        let mut buf = rest.to_vec();

        let nonce = aead::Nonce::try_assume_unique_for_key(nonce).ok()?;

        let plain = self.key().open_in_place(nonce, aead::Aad::empty(), &mut buf).ok()?;

        String::from_utf8(plain.to_vec()).ok()

    }

}

// ------------------------------------------------------------- storage

pub async fn list(pool: &Pool, scope: Scope) -> sqlx::Result<Vec<Assistant>> {

    sqlx::query_as("SELECT * FROM assistants WHERE scope = ? AND scope_id = ? ORDER BY name")

        .bind(scope.kind())

        .bind(scope.id())

        .fetch_all(pool)

        .await

}

pub async fn get(pool: &Pool, id: i64) -> sqlx::Result<Option<Assistant>> {

    sqlx::query_as("SELECT * FROM assistants WHERE id = ?").bind(id).fetch_optional(pool).await

}

/// Add an assistant; `None` when the level already has one of that name.

pub async fn create(pool: &Pool, scope: Scope, f: &Fields, token: &[u8]) -> sqlx::Result<Option<Assistant>> {

    let now = now_ms();

    let r = sqlx::query_as(

        "INSERT INTO assistants (scope, scope_id, name, endpoint, api, model, headers, token, enabled, budget, created_ms, updated_ms) \

         VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT DO NOTHING RETURNING *",

    )

    .bind(scope.kind())

    .bind(scope.id())

    .bind(&f.name)

    .bind(&f.endpoint)

    .bind(&f.api)

    .bind(&f.model)

    .bind(&f.headers)

    .bind(token)

    .bind(f.enabled)

    .bind(f.budget)

    .bind(now)

    .bind(now)

    .fetch_optional(pool)

    .await?;

    Ok(r)

}

/// Change an assistant's details; `token` is `None` to keep the stored

/// one. False when the name is taken at that level.

pub async fn update(pool: &Pool, id: i64, f: &Fields, token: Option<&[u8]>) -> sqlx::Result<bool> {

    let r = match token {

        Some(t) => sqlx::query(

            "UPDATE assistants SET name = ?, endpoint = ?, api = ?, model = ?, headers = ?, enabled = ?, budget = ?, token = ?, updated_ms = ? \

             WHERE id = ? AND NOT EXISTS (SELECT 1 FROM assistants o WHERE o.scope = assistants.scope AND o.scope_id = assistants.scope_id AND o.name = ? AND o.id <> ?)",

        )

        .bind(&f.name)

        .bind(&f.endpoint)

        .bind(&f.api)

        .bind(&f.model)

        .bind(&f.headers)

        .bind(f.enabled)

        .bind(f.budget)

        .bind(t)

        .bind(now_ms())

        .bind(id)

        .bind(&f.name)

        .bind(id)

        .execute(pool)

        .await?,

        None => sqlx::query(

            "UPDATE assistants SET name = ?, endpoint = ?, api = ?, model = ?, headers = ?, enabled = ?, budget = ?, updated_ms = ? \

             WHERE id = ? AND NOT EXISTS (SELECT 1 FROM assistants o WHERE o.scope = assistants.scope AND o.scope_id = assistants.scope_id AND o.name = ? AND o.id <> ?)",

        )

        .bind(&f.name)

        .bind(&f.endpoint)

        .bind(&f.api)

        .bind(&f.model)

        .bind(&f.headers)

        .bind(f.enabled)

        .bind(f.budget)

        .bind(now_ms())

        .bind(id)

        .bind(&f.name)

        .bind(id)

        .execute(pool)

        .await?,

    };

    Ok(r.rows_affected() == 1)

}

pub async fn set_enabled(pool: &Pool, id: i64, enabled: bool) -> sqlx::Result<()> {

    sqlx::query("UPDATE assistants SET enabled = ?, updated_ms = ? WHERE id = ?")

        .bind(enabled)

        .bind(now_ms())

        .bind(id)

        .execute(pool)

        .await?;

    Ok(())

}

pub async fn delete(pool: &Pool, id: i64) -> sqlx::Result<()> {

    sqlx::query("DELETE FROM assistants WHERE id = ?").bind(id).execute(pool).await?;

    Ok(())

}

/// Forget a level's assistants and prompt, when the organisation or

/// repository goes.

pub async fn delete_scope(pool: &Pool, scope: Scope) -> sqlx::Result<()> {

    sqlx::query("DELETE FROM assistants WHERE scope = ? AND scope_id = ?").bind(scope.kind()).bind(scope.id()).execute(pool).await?;

    sqlx::query("DELETE FROM assistant_prompts WHERE scope = ? AND scope_id = ?").bind(scope.kind()).bind(scope.id()).execute(pool).await?;

    Ok(())

}

/// The prompt a level adds, empty when none.

pub async fn prompt(pool: &Pool, scope: Scope) -> sqlx::Result<String> {

    Ok(sqlx::query_scalar("SELECT prompt FROM assistant_prompts WHERE scope = ? AND scope_id = ?")

        .bind(scope.kind())

        .bind(scope.id())

        .fetch_optional(pool)

        .await?

        .unwrap_or_default())

}

pub async fn set_prompt(pool: &Pool, scope: Scope, text: &str) -> sqlx::Result<()> {

    if text.trim().is_empty() {

        sqlx::query("DELETE FROM assistant_prompts WHERE scope = ? AND scope_id = ?")

            .bind(scope.kind())

            .bind(scope.id())

            .execute(pool)

            .await?;

    } else {

        sqlx::query(

            "INSERT INTO assistant_prompts (scope, scope_id, prompt, updated_ms) VALUES (?, ?, ?, ?) \

             ON CONFLICT (scope, scope_id) DO UPDATE SET prompt = excluded.prompt, updated_ms = excluded.updated_ms",

        )

        .bind(scope.kind())

        .bind(scope.id())

        .bind(text.trim())

        .bind(now_ms())

        .execute(pool)

        .await?;

    }

    Ok(())

}

// ----------------------------------------------------- what applies where

/// An assistant seen from a level: where it comes from, and whether it is

/// on once outer levels have had their say.

#[derive(Clone, Debug)]

pub struct Available {

    pub assistant: Assistant,

    /// `the site`, `the acme organisation`, or empty for the level's own.

    pub from: String,

    /// Off because a level above disabled an assistant of this name.

    pub off_above: bool,

}

impl Available {

    pub fn on(&self) -> bool {

        self.assistant.enabled && !self.off_above

    }

}

/// Every assistant a level sees, outermost first, one per name: the

/// innermost definition of a name wins, and a name disabled at an outer

/// level is off inside it whatever the inner level says.

pub async fn available(pool: &Pool, scope: Scope, owner: &str) -> sqlx::Result<Vec<Available>> {

    let mut levels = crate::levels::above(pool, scope, owner).await?;

    levels.push((scope, String::new()));

    let mut out: Vec<Available> = vec![];

    let mut off: Vec<String> = vec![];

    for (level, from) in levels {

        for a in list(pool, level).await? {

            let off_above = off.iter().any(|n| n.eq_ignore_ascii_case(&a.name));

            if !a.enabled && !off.iter().any(|n| n.eq_ignore_ascii_case(&a.name)) {

                off.push(a.name.clone());

            }

            out.retain(|o| !o.assistant.name.eq_ignore_ascii_case(&a.name));

            out.push(Available { assistant: a, from: from.clone(), off_above });

        }

    }

    Ok(out)

}

/// The assistants a repository may ask, by name.

pub async fn enabled_for_repo(pool: &Pool, repo: &crate::db::repos::RepoRow) -> sqlx::Result<Vec<Assistant>> {

    Ok(available(pool, Scope::Repo(repo.id), &repo.owner).await?.into_iter().filter(Available::on).map(|a| a.assistant).collect())

}

/// The prompts a level inherits, outermost first: (where from, text),

/// only the levels that set one.

pub async fn inherited_prompts(pool: &Pool, scope: Scope, owner: &str) -> sqlx::Result<Vec<(String, String)>> {

    let mut out = vec![];

    for (level, from) in crate::levels::above(pool, scope, owner).await? {

        let text = prompt(pool, level).await?;

        if !text.is_empty() {

            out.push((from, text));

        }

    }

    Ok(out)

}

/// The prompt in force for a repository: each level's text under a

/// heading, outermost first, so a project adds to what the site says.

pub async fn prompt_for_repo(pool: &Pool, repo: &crate::db::repos::RepoRow) -> sqlx::Result<String> {

    let scope = Scope::Repo(repo.id);

    let mut parts = inherited_prompts(pool, scope, &repo.owner).await?;

    let own = prompt(pool, scope).await?;

    if !own.is_empty() {

        parts.push(("this repository".to_string(), own));

    }

    Ok(parts.into_iter().map(|(from, text)| format!("## Instructions from {from}\n\n{text}\n")).collect::<Vec<_>>().join("\n"))

}

// -------------------------------------------------------------- calling

/// Ask an assistant: the system prompt and the request, the answer's

/// text. Errors are for the person who asked: a status and the first

/// line of what the endpoint said, or why no answer came.

pub async fn ask(app: &Shared, a: &Assistant, system: &str, user: &str) -> Result<String, String> {

    let Some(token) = app.assistant_key.open(&a.token) else {

        return Err("its token cannot be read: set it again".into());

    };

    let client = reqwest::Client::builder()

        .timeout(CALL_TIMEOUT)

        .user_agent(format!("oxygenhub/{}", env!("CARGO_PKG_VERSION")))

        .build()

        .map_err(|e| e.to_string())?;

    let mut req = client.post(&a.endpoint).header("content-type", "application/json");

    let body = match a.api.as_str() {

        "anthropic-messages" => {

            if !token.is_empty() {

                req = req.header("x-api-key", &token);

            }

            req = req.header("anthropic-version", "2023-06-01");

            serde_json::json!({

                "model": a.model,

                "max_tokens": 4096,

                "system": system,

                "messages": [{ "role": "user", "content": user }],

            })

        }

        _ => {

            if !token.is_empty() {

                req = req.header("authorization", format!("Bearer {token}"));

            }

            serde_json::json!({

                "model": a.model,

                "messages": [

                    { "role": "system", "content": system },

                    { "role": "user", "content": user },

                ],

            })

        }

    };

    for (name, value) in a.header_list() {

        req = req.header(name, value);

    }

    let r = req.json(&body).send().await.map_err(|e| {

        if e.is_timeout() {

            "no answer within 60 seconds".to_string()

        } else {

            scrub(&e.to_string(), &token)

        }

    })?;

    let status = r.status();

    let text: String = r.text().await.unwrap_or_default().chars().take(MAX_ANSWER).collect();

    if !status.is_success() {

        let why = serde_json::from_str::<serde_json::Value>(&text)

            .ok()

            .and_then(|v| v["error"]["message"].as_str().or(v["error"].as_str()).or(v["message"].as_str()).map(str::to_string))

            .unwrap_or_else(|| text.lines().next().unwrap_or("").chars().take(160).collect());

        return Err(scrub(&format!("{} {}", status.as_u16(), why.trim()), &token));

    }

    let v: serde_json::Value = serde_json::from_str(&text).map_err(|_| "the answer was not JSON".to_string())?;

    let answer = match a.api.as_str() {

        "anthropic-messages" => v["content"]

            .as_array()

            .map(|parts| parts.iter().filter_map(|p| p["text"].as_str()).collect::<Vec<_>>().join(""))

            .unwrap_or_default(),

        _ => match &v["choices"][0]["message"]["content"] {

            serde_json::Value::String(s) => s.clone(),

            serde_json::Value::Array(parts) => parts.iter().filter_map(|p| p["text"].as_str()).collect::<Vec<_>>().join(""),

            _ => String::new(),

        },

    };

    if answer.trim().is_empty() {

        return Err("the answer carried no text".into());

    }

    Ok(answer)

}

/// A tiny exchange, for the settings page's TEST: the start of the answer.

pub async fn test(app: &Shared, a: &Assistant) -> Result<String, String> {

    let answer = ask(app, a, "You are being connected to an OxygenHub site. Answer briefly.", "Reply with the single word OK.").await?;

    Ok(answer.trim().chars().take(80).collect())

}

/// Never let the token through in an error message.

fn scrub(text: &str, token: &str) -> String {

    if token.is_empty() {

        text.to_string()

    } else {

        text.replace(token, "***")

    }

}

#[cfg(test)]

mod tests {

    use super::*;

    #[test]

    fn tokens_seal_and_open_only_under_their_key() {

        let dir = tempfile::tempdir().unwrap();

        let k = SecretKey::load_or_create(dir.path()).unwrap();

        let sealed = k.seal("sk-very-secret");

        assert_ne!(sealed, b"sk-very-secret");

        assert!(!String::from_utf8_lossy(&sealed).contains("secret"));

        assert_eq!(k.open(&sealed).as_deref(), Some("sk-very-secret"));

        assert_eq!(SecretKey::load_or_create(dir.path()).unwrap().open(&sealed).as_deref(), Some("sk-very-secret"), "the key is kept");

        let other = SecretKey::load_or_create(tempfile::tempdir().unwrap().path()).unwrap();

        assert_eq!(other.open(&sealed), None);

        assert_eq!(k.open(&[1, 2, 3]), None);

        assert_eq!(k.seal(""), Vec::<u8>::new());

        assert_eq!(k.open(&[]).as_deref(), Some(""));

    }

    #[test]

    fn headers_parse_or_name_the_bad_line() {

        assert_eq!(

            parse_headers("X-Org: acme\n\n OpenAI-Beta : v1 ").unwrap(),

            vec![("X-Org".into(), "acme".into()), ("OpenAI-Beta".into(), "v1".into())]