//! 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;

/// Which level an assistant or a prompt belongs to.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Scope {
    Site,
    Org(i64),
    Repo(i64),
}

impl Scope {
    pub fn kind(self) -> &'static str {
        match self {
            Scope::Site => "site",
            Scope::Org(_) => "org",
            Scope::Repo(_) => "repo",
        }
    }

    pub fn id(self) -> i64 {
        match self {
            Scope::Site => 0,
            Scope::Org(id) | Scope::Repo(id) => id,
        }
    }
}

#[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

/// The levels above a scope, outermost first: the site, then the
/// organisation when the owner is one.
pub async fn levels_above(pool: &Pool, scope: Scope, owner: &str) -> sqlx::Result<Vec<(Scope, String)>> {
    let mut out = vec![];
    match scope {
        Scope::Site => {}
        Scope::Org(_) => out.push((Scope::Site, "the site".to_string())),
        Scope::Repo(_) => {
            out.push((Scope::Site, "the site".to_string()));
            if let Some(org) = crate::db::accounts::org_by_name(pool, owner).await? {
                out.push((Scope::Org(org.id), format!("the {} organisation", org.name)));
            }
        }
    }
    Ok(out)
}

/// 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 = 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 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())]
        );
        assert_eq!(parse_headers("no colon").unwrap_err(), "no colon");
        assert_eq!(parse_headers("bad name: x").unwrap_err(), "bad name: x");
    }
}