//! The context stack: builds a role's system prompt from a base
//! fragment, layered `system.md` fragments, and `context: true` entries.

use std::fs;
use std::path::PathBuf;
use std::sync::Arc;

use crate::knowledge::Mount;

/// A horizontal rule between major segments of the prompt: base, layers,
/// context entries.
const DIVIDER: &str = "\n\n---\n\n";

/// The heading every context entry gets when pushed into the prompt.
const ENTRY_HEADING: &str = "### ";

/// The blank line that separates one context entry from the next
/// entry's heading.
const ENTRY_SEPARATOR: &str = "\n\n";

/// What a system prompt is composed from: a base fragment, the layer
/// directories that may each add a `system.md`, and the mounted
/// knowledge.
///
/// [`ContextStack::prompt`] composes the prompt every time it is called,
/// and the DM calls it at the start of every turn. The two halves age
/// differently. Each layer's `system.md` is read from disk on every
/// call, so an edit to a fragment reaches the next turn. The entries
/// are the ones the mount holds, and the mount is a snapshot taken when
/// the session started, so a file added to a layer mid-session is not
/// mounted and does not reach the prompt.
#[derive(Debug)]
pub struct ContextStack {
    base: String,
    layers: Vec<PathBuf>,
    mount: Arc<Mount>,
}

impl ContextStack {
    /// Builds a context stack over a `base` fragment, `layers` (lowest
    /// first), and `mount`.
    ///
    /// The base fragment is what distinguishes roles: the DM base sets
    /// identity, narration style, and tool usage. A planner agent would
    /// pass a different base. What stays the same is how layers and
    /// context entries are assembled on top of it.
    pub fn new(base: &str, layers: &[PathBuf], mount: Arc<Mount>) -> Self {
        Self {
            base: base.to_string(),
            layers: layers.to_vec(),
            mount,
        }
    }

    /// Composes the system prompt, without a trailing newline: the base
    /// fragment first, then each layer's `system.md` in mount order,
    /// then the mount's knowledge directories, then every mounted entry
    /// with `context: true`, each separated from the next by a blank
    /// line.
    ///
    /// A layer with no `system.md` contributes nothing and is not a
    /// failure. A layer whose `system.md` cannot be read is: the caller
    /// turns that into a failed turn rather than sending the model a
    /// prompt that is quietly missing a layer.
    pub fn prompt(&self) -> Result<String, String> {
        let mut prompt = String::new();

        prompt.push_str(self.base.trim());

        for layer_root in &self.layers {
            let fragment_path = layer_root.join("system.md");
            match fs::read_to_string(&fragment_path) {
                Ok(text) => {
                    prompt.push_str(DIVIDER);
                    prompt.push_str(text.trim());
                }
                Err(err) if err.kind() == std::io::ErrorKind::NotFound => {
                    // A layer without a fragment contributes nothing.
                }
                Err(err) => {
                    return Err(format!(
                        "{}: cannot read system.md: {err}",
                        layer_root.display()
                    ));
                }
            }
        }

        let top_dirs = top_level_directories(&self.mount);
        if !top_dirs.is_empty() {
            prompt.push_str(DIVIDER);
            prompt.push_str("## Knowledge directories\n\n");
            prompt.push_str("Each scope is followed by the frontmatter keys you can filter on with `lookup`'s `filter` parameter.\n\n");
            prompt.push_str(&top_dirs);
        }

        let context_entries: Vec<_> = self.mount.entries().filter(|e| e.context).collect();
        if !context_entries.is_empty() {
            if !prompt.is_empty() {
                prompt.push_str(DIVIDER);
            }
            for entry in context_entries {
                prompt.push_str(ENTRY_HEADING);
                prompt.push_str(&entry.address);
                prompt.push_str(" — ");
                prompt.push_str(&entry.name);
                prompt.push('\n');
                prompt.push_str(entry.body.trim());
                prompt.push_str(ENTRY_SEPARATOR);
            }
        }

        Ok(prompt.trim().to_string())
    }
}

/// Collects the top two levels of mounted addresses as a sorted list
/// of scopes like `rules/spells` or `lore/places`, each with the
/// frontmatter keys seen on entries under that scope.
fn top_level_directories(mount: &Mount) -> String {
    let mut scope_keys: std::collections::BTreeMap<String, std::collections::BTreeSet<String>> =
        std::collections::BTreeMap::new();
    for entry in mount.entries() {
        let scope = entry
            .address
            .split('/')
            .take(2)
            .collect::<Vec<_>>()
            .join("/");
        let keys = scope_keys.entry(scope).or_default();
        let mapping_keys = entry
            .frontmatter
            .as_mapping()
            .into_iter()
            .flat_map(|mapping| mapping.keys());
        for key in mapping_keys {
            if let Some(key_str) = key.as_str() {
                // `name` is what keyword search matches, and `deleted`,
                // `context`, and `visibility` steer the mount itself, so
                // none of them is useful as a filter key.
                if key_str != "deleted"
                    && key_str != "context"
                    && key_str != "visibility"
                    && key_str != "name"
                {
                    keys.insert(key_str.to_string());
                }
            }
        }
    }
    scope_keys
        .into_iter()
        .map(|(scope, keys)| {
            if keys.is_empty() {
                format!("- `{scope}`")
            } else {
                let key_list: Vec<_> = keys.into_iter().collect();
                format!("- `{scope}` — {}", key_list.join(", "))
            }
        })
        .collect::<Vec<_>>()
        .join("\n")
}

#[cfg(test)]
#[path = "context_tests.rs"]
mod tests;