//! Memory: the arguments of `recall`, `memory` and `ask_memory`, their shape checks, and the
//! messages `memoryd` exchanges with `loopd`, `brokerd` and `bxctl` (M5a spec, sections 5 and 7;
//! M5c spec, section 9). Content rules (an entry matching the file, a file's limit, the check on
//! characters, which workspaces `[honcho]` has) are `memoryd`'s.

use serde::{Deserialize, Serialize};

use crate::class::DataClass;
use crate::ids::SessionId;
use crate::wire::ToolRequest;

/// The tool names.
pub const RECALL: &str = "recall";
pub const MEMORY: &str = "memory";
pub const ASK_MEMORY: &str = "ask_memory";

/// The most characters a `recall` query may have.
pub const MAX_QUERY: usize = 500;
/// The most hits a `recall` may ask for, and how many it gets when it does not say.
pub const MAX_K: u64 = 10;
pub const DEFAULT_K: u64 = 5;
/// The most characters a `note` may have.
pub const MAX_NOTE: usize = 1_000;
/// The most characters of any other entry: the larger file's limit.
pub const MAX_ENTRY: usize = 2_200;
/// The line that separates entries in `memory.md`, `user.md` and `notes.md`.
pub const SEPARATOR: &str = "§";
/// The most characters an `ask_memory` question may have.
pub const MAX_QUESTION: usize = 2_000;
/// The most workspaces one `ask_memory` may name: both are asked at once, inside `brokerd`'s wait.
pub const MAX_WORKSPACES: usize = 2;
/// The most characters of a workspace name in `ask_memory`. It matches `memoryd`'s rule for
/// Honcho's names (`honcho::config::MAX_NAME`), so every agent `[[honcho.read]]` allows can be
/// asked.
pub const MAX_WORKSPACE_NAME: usize = 512;
/// The workspace `ask_memory` asks when it names none: the agent's own trusted one.
pub const OWN_WORKSPACE: &str = "own";

/// The arguments for `recall`.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct RecallArgs {
    pub query: String,
    #[serde(default = "default_k")]
    pub k: u64,
}

fn default_k() -> u64 {
    DEFAULT_K
}

/// The arguments for `ask_memory` (M5c spec, section 9). An empty `workspaces` asks `own`.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct AskMemoryArgs {
    pub question: String,
    #[serde(default)]
    pub workspaces: Vec<String>,
}

/// Which identity file an edit is for.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum MemoryFile {
    Memory,
    User,
}

impl MemoryFile {
    /// The file's name under `memory/`.
    pub fn name(self) -> &'static str {
        match self {
            MemoryFile::Memory => "memory.md",
            MemoryFile::User => "user.md",
        }
    }

    /// The file's limit in characters.
    pub fn limit(self) -> usize {
        match self {
            MemoryFile::Memory => 2_200,
            MemoryFile::User => 1_375,
        }
    }
}

/// The arguments for `memory`: one action. `note` records experience; the others edit identity.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "action", rename_all = "snake_case", deny_unknown_fields)]
pub enum MemoryArgs {
    Add {
        file: MemoryFile,
        text: String,
    },
    Replace {
        file: MemoryFile,
        old: String,
        text: String,
    },
    Remove {
        file: MemoryFile,
        old: String,
    },
    Consolidate {
        file: MemoryFile,
        entries: Vec<ConsolidateEntry>,
    },
    Note {
        text: String,
    },
}

/// One line of a `consolidate`: an existing entry (`old`) with exactly one of `keep`, `text`,
/// `move`, `drop`; or a new entry (`text` alone).
#[derive(Debug, Clone, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct ConsolidateEntry {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub old: Option<String>,
    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
    pub keep: bool,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub text: Option<String>,
    #[serde(default, rename = "move", skip_serializing_if = "std::ops::Not::not")]
    pub moved: bool,
    #[serde(default, skip_serializing_if = "std::ops::Not::not")]
    pub drop: bool,
}

/// What a `consolidate` line does.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum EntryAction {
    Keep(String),
    Rewrite { old: String, text: String },
    Move(String),
    Drop(String),
    New(String),
}

impl ConsolidateEntry {
    /// The line's action, or why it has none.
    pub fn action(&self) -> Result<EntryAction, &'static str> {
        let flags = usize::from(self.keep) + usize::from(self.moved) + usize::from(self.drop);
        match (&self.old, &self.text, flags) {
            (None, Some(text), 0) => Ok(EntryAction::New(text.clone())),
            (None, _, _) => Err("an entry without old must have text and nothing else"),
            (Some(old), Some(text), 0) => Ok(EntryAction::Rewrite {
                old: old.clone(),
                text: text.clone(),
            }),
            (Some(old), None, 1) => {
                if self.keep {
                    Ok(EntryAction::Keep(old.clone()))
                } else if self.moved {
                    Ok(EntryAction::Move(old.clone()))
                } else {
                    Ok(EntryAction::Drop(old.clone()))
                }
            }
            (Some(_), _, _) => {
                Err("an entry with old must have exactly one of keep, text, move or drop")
            }
        }
    }
}

/// Why an entry's text is not a well-formed entry: empty once trimmed, too long, or holding a line
/// that starts with the separator, spaces aside. In `notes.md` such a line would read as a header,
/// and a header says where an entry came from. `None` when it is fine.
pub fn entry_problem(text: &str, max: usize) -> Option<String> {
    if text.trim().is_empty() {
        return Some("an entry cannot be empty".to_string());
    }
    let chars = text.chars().count();
    if chars > max {
        return Some(format!(
            "an entry has at most {max} characters; this has {chars}"
        ));
    }
    for line in text.lines() {
        if line.trim_start().starts_with(SEPARATOR) {
            return Some(format!(
                "an entry cannot hold a line that starts with {SEPARATOR}"
            ));
        }
    }
    None
}

/// `entry_problem` as a `Result`, so a check can use `?`.
fn entry(text: &str, max: usize) -> Result<(), String> {
    match entry_problem(text, max) {
        Some(problem) => Err(problem),
        None => Ok(()),
    }
}

impl RecallArgs {
    /// The shape rules of `recall` (M5a spec, section 5).
    pub fn check(&self) -> Result<(), String> {
        if self.query.trim().is_empty() || self.query.chars().count() > MAX_QUERY {
            return Err(format!("query must have 1 to {MAX_QUERY} characters"));
        }
        if self.k == 0 || self.k > MAX_K {
            return Err(format!("k must be 1 to {MAX_K}"));
        }
        Ok(())
    }
}

impl AskMemoryArgs {
    /// The shape rules of `ask_memory` (M5c spec, section 9), in this order: the question, the
    /// number of names, then each name in turn. Only a name's form is checked here. Each sentence
    /// reaches the model, so it says what to do instead.
    pub fn check(&self) -> Result<(), String> {
        if self.question.trim().is_empty() {
            return Err("question is empty; write the question to ask".to_string());
        }
        let chars = self.question.chars().count();
        if chars > MAX_QUESTION {
            return Err(format!(
                "question has {chars} characters, over the limit of {MAX_QUESTION}; ask a shorter question"
            ));
        }
        let names = self.workspaces.len();
        if names > MAX_WORKSPACES {
            return Err(format!(
                "workspaces names {names}, over the limit of {MAX_WORKSPACES}; ask about the others in another call"
            ));
        }
        for (i, name) in self.workspaces.iter().enumerate() {
            if name.is_empty() {
                return Err(
                    "a workspace name is empty; name a workspace, or leave workspaces out to ask own"
                        .to_string(),
                );
            }
            if !workspace_name(name) {
                return Err(format!(
                    "a workspace name has 1 to {MAX_WORKSPACE_NAME} letters, digits, _ or -; use own, untrusted, fleet or another agent's name"
                ));
            }
            if self.workspaces.iter().take(i).any(|w| w == name) {
                return Err(format!(
                    "workspaces names {name} twice; name each workspace once"
                ));
            }
        }
        Ok(())
    }

    /// The workspaces to ask: those named, in order, or `own` when none is.
    pub fn asked(&self) -> Vec<&str> {
        if self.workspaces.is_empty() {
            vec![OWN_WORKSPACE]
        } else {
            self.workspaces.iter().map(String::as_str).collect()
        }
    }
}

/// A workspace name's form: at most `MAX_WORKSPACE_NAME` characters, each an ASCII letter, a
/// digit, `_` or `-`, as Honcho's names are. `check` refuses an empty name first.
fn workspace_name(name: &str) -> bool {
    name.len() <= MAX_WORKSPACE_NAME
        && name
            .bytes()
            .all(|b| b.is_ascii_alphanumeric() || b == b'_' || b == b'-')
}

impl MemoryArgs {
    /// The shape rules of `memory` (M5a spec, section 5): every text is a well-formed entry, and
    /// every `consolidate` line has one action and the list is not empty.
    pub fn check(&self) -> Result<(), String> {
        match self {
            MemoryArgs::Add { text, .. } => entry(text, MAX_ENTRY),
            MemoryArgs::Replace { old, text, .. } => {
                entry(old, MAX_ENTRY)?;
                entry(text, MAX_ENTRY)
            }
            MemoryArgs::Remove { old, .. } => entry(old, MAX_ENTRY),
            MemoryArgs::Note { text } => entry(text, MAX_NOTE),
            MemoryArgs::Consolidate { entries, .. } => {
                if entries.is_empty() {
                    return Err("consolidate needs every entry of the file".to_string());
                }
                for line in entries {
                    match line.action().map_err(str::to_string)? {
                        EntryAction::Keep(old)
                        | EntryAction::Move(old)
                        | EntryAction::Drop(old) => entry(&old, MAX_ENTRY)?,
                        EntryAction::Rewrite { old, text } => {
                            entry(&old, MAX_ENTRY)?;
                            entry(&text, MAX_ENTRY)?;
                        }
                        EntryAction::New(text) => entry(&text, MAX_ENTRY)?,
                    }
                }
                Ok(())
            }
        }
    }
}

/// `memoryd` → `loopd`: the memory block for a new session's baseline.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct MemoryBlockReply {
    pub text: String,
}

/// `loopd` → `memoryd`: this session's log has new records, those of a turn that has ended.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct SessionUpdated {
    pub session: SessionId,
    /// The ended turn's number in the session (M5c spec, section 5). It has no default: a frame
    /// without it is refused, never read as turn 0.
    pub turn: u64,
}

/// `memoryd` → `brokerd`, on `state.sock`: the session's state now (M5c spec, section 7).
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct SessionState {
    pub session: SessionId,
}

/// `brokerd` → `memoryd`: the session's class, and whether it has read untrusted content. Neither
/// field has a default: a reply without one is refused, never read as a clean session.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct SessionStateReply {
    pub class: DataClass,
    pub untrusted: bool,
}

/// `brokerd` → `memoryd`: a decided `recall`, `memory` or `ask_memory` call, with the session's
/// state when it was decided.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct MemoryCall {
    pub request: ToolRequest,
    pub taint: DataClass,
    pub untrusted: bool,
}

/// `memoryd` → `bxctl`: what a reindex read.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct Reindexed {
    pub logs: u64,
    pub records: u64,
    pub notes: u64,
}

/// The owner's memory commands (M5a spec, section 10).
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum MemoryCommand {
    Log,
    Undo,
    Pending,
    Accept,
    Reject,
    /// What the Honcho ledger holds: its counts and failed items (M5c spec, section 5).
    Sync,
}

/// `bxctl` → `memoryd`: one command; `arg` is the count for `log`, the edit or pending id for
/// `undo`, `accept` and `reject`, and empty otherwise.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct MemoryAdmin {
    pub command: MemoryCommand,
    pub arg: String,
}

/// `memoryd` → `bxctl`: the text to print.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct MemoryAdminReply {
    pub text: String,
}