//! 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 a call may take while a person waits for it (a draft, a

/// connection test): inside Cloudflare's 100 seconds for a proxied request.

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

/// How long a call answered in the background may take (a review, a

/// summary): a model that thinks first needs minutes on a large proposal

/// (oxygenhub#201).

pub const BACKGROUND_TIMEOUT: Duration = Duration::from_secs(300);

/// The longest answer asked for: room for a model that thinks before it

/// writes (oxygenhub#201: 4096 was spent on thinking alone).

pub const MAX_ANSWER_TOKENS: u32 = 16_000;

/// The most an answer is kept.

const MAX_ANSWER: usize = 64 * 1024;

/// The longest a streamed answer may take in all, however steadily it

/// comes: the limit a call is given is for silence (oxygenhub#248).

const STREAM_LIMIT: Duration = Duration::from_secs(15 * 60);

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> {

        SecretKey::load_or_create_named(keys, "assistants.key")

    }

    /// The same, for another key file beside it (`secrets.key` seals Flow

    /// secrets, #107).

    pub fn load_or_create_named(keys: &Path, file: &str) -> anyhow::Result<SecretKey> {

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

        let path = keys.join(file);

        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 = $1 AND scope_id = $2 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 = $1").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 ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12) 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 = $1, endpoint = $2, api = $3, model = $4, headers = $5, enabled = $6, budget = $7, token = $8, updated_ms = $9 \

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

        )

        .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 = $1, endpoint = $2, api = $3, model = $4, headers = $5, enabled = $6, budget = $7, updated_ms = $8 \

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

        )

        .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)

}

/// A renamed assistant keeps its name on the reviews it answered and the

/// events that mention it, for the repositories its level covers: the

/// name is what people read, not a key.

pub async fn rename_in_history(pool: &Pool, scope: Scope, owner_hint: &str, old: &str, new: &str) -> sqlx::Result<()> {

    if old == new || old.is_empty() || new.is_empty() {

        return Ok(());

    }

    let like = format!("%{old}%");

    match scope {

        Scope::Site => {

            for table in ["proposal_reviews", "proposal_comments"] {

                sqlx::query(&format!("UPDATE {table} SET via = $1 WHERE via = $2")).bind(new).bind(old).execute(pool).await?;

            }

            sqlx::query("UPDATE proposal_events SET event = replace(event, $1, $2) WHERE event LIKE $3")

                .bind(old)

                .bind(new)

                .bind(&like)

                .execute(pool)

                .await?;

        }

        Scope::Org(_) => {

            let in_org = "proposal_id IN (SELECT p.id FROM proposals p JOIN repositories r ON r.id = p.repo_id WHERE r.owner = $1)";

            for table in ["proposal_reviews", "proposal_comments"] {

                sqlx::query(&format!("UPDATE {table} SET via = $1 WHERE via = $2 AND {in_org}"))

                    .bind(new)

                    .bind(old)

                    .bind(owner_hint)

                    .execute(pool)

                    .await?;

            }

            sqlx::query(&format!("UPDATE proposal_events SET event = replace(event, $1, $2) WHERE event LIKE $3 AND {in_org}"))

                .bind(old)

                .bind(new)

                .bind(&like)

                .bind(owner_hint)

                .execute(pool)

                .await?;

        }

        Scope::Repo(id) => {

            let in_repo = "proposal_id IN (SELECT id FROM proposals WHERE repo_id = $1)";

            for table in ["proposal_reviews", "proposal_comments"] {

                sqlx::query(&format!("UPDATE {table} SET via = $1 WHERE via = $2 AND {in_repo}"))

                    .bind(new)

                    .bind(old)

                    .bind(id)

                    .execute(pool)

                    .await?;

            }

            sqlx::query(&format!("UPDATE proposal_events SET event = replace(event, $1, $2) WHERE event LIKE $3 AND {in_repo}"))

                .bind(old)

                .bind(new)

                .bind(&like)

                .bind(id)

                .execute(pool)

                .await?;

        }

    }

    Ok(())

}

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

    sqlx::query("UPDATE assistants SET enabled = $1, updated_ms = $2 WHERE id = $3")

        .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 = $1").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 = $1 AND scope_id = $2").bind(scope.kind()).bind(scope.id()).execute(pool).await?;

    sqlx::query("DELETE FROM assistant_prompts WHERE scope = $1 AND scope_id = $2")

        .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 = $1 AND scope_id = $2")

        .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 = $1 AND scope_id = $2")

            .bind(scope.kind())

            .bind(scope.id())

            .execute(pool)

            .await?;

    } else {

        sqlx::query(

            "INSERT INTO assistant_prompts (scope, scope_id, prompt, updated_ms) VALUES ($1, $2, $3, $4) \

             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. The answer is

/// asked for as a stream, and `wait` is the longest the provider may be

/// silent, before it starts or between parts, so a long answer that keeps

/// coming is not cut off (oxygenhub#248); a provider that answers whole

/// instead is read whole within `wait`.

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

    ask_streaming(app, a, system, user, wait, |_| {}).await

}

/// [`ask`], handing each part of the answer to `part` as it arrives.

pub async fn ask_streaming(

    app: &Shared,

    a: &Assistant,

    system: &str,

    user: &str,

    wait: Duration,

    mut part: impl FnMut(&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 = crate::net::client(&a.endpoint, &app.config.outbound.allow, STREAM_LIMIT.max(wait)).await?;

    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": MAX_ANSWER_TOKENS,

                "system": system,

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

                "stream": true,

            })

        }

        _ => {

            if !token.is_empty() {

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

            }

            // No length here: providers differ (`max_tokens`, or

            // `max_completion_tokens` for newer OpenAI models, with

            // different ceilings), and their defaults are generous.

            serde_json::json!({

                "model": a.model,

                "messages": [

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

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

                ],

                "stream": true,

            })

        }

    };

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

        req = req.header(name, value);

    }

    let quiet = || format!("no answer within {} seconds", wait.as_secs());

    let r = tokio::time::timeout(wait, req.json(&body).send()).await.map_err(|_| quiet())?.map_err(|e| {

        if e.is_timeout() {

            quiet()

        } else {

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

        }

    })?;

    let status = r.status();

    let streamed =

        r.headers().get(reqwest::header::CONTENT_TYPE).and_then(|v| v.to_str().ok()).is_some_and(|t| t.starts_with("text/event-stream"));

    if status.is_success() && streamed {