kcode_kweb_context/

lib.rs

#![forbid(unsafe_code)]

use std::{
    collections::{BTreeMap, HashMap, HashSet},
    error, fmt,
};

use kcode_kweb_db::NodeId;
use kcode_session_history::{
    Session as HistorySession,
    chatend::{BoxContent, BoxId, BoxState, EventId, PendingId, ToolSlotInput},
};
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use sha2::{Digest, Sha256};

const KWEB_TOOL_INSTANCE: &str = "kweb";
const RECENT_CONNECTIONS_PER_BOX: usize = 8;
const RECENT_CONNECTIONS_LOGICAL_SLOT: &str = "recent-connections";
const RECENT_CONNECTION_IDS_METADATA: &str = "kwebRecentConnectionIds";

pub type Result<T> = std::result::Result<T, Error>;

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct Error {
    message: String,
}

impl Error {
    fn new(message: impl Into<String>) -> Self {
        Self {
            message: message.into(),
        }
    }
}

impl fmt::Display for Error {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter.write_str(&self.message)
    }
}

impl error::Error for Error {}

#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Connection {
    pub id: String,
    pub short_name: String,
    pub short_description: String,
}

#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct Node {
    pub id: String,
    pub short_name: String,
    pub short_description: String,
    pub long_description: String,
    pub owner: String,
    #[serde(default)]
    pub fixed_connections: Vec<Connection>,
    #[serde(default)]
    pub recent_connections: Vec<Connection>,
    #[serde(default)]
    pub objects: Vec<String>,
    #[serde(default)]
    pub last_modified_by: String,
    #[serde(default)]
    pub last_modified_at: Option<String>,
}

impl Node {
    pub fn from_kweb_value(value: &Value) -> Result<Self> {
        let id = required_string(value, "id")?;
        canonical_node_id(&id)?;
        let owner = value
            .get("owner_node_id")
            .or_else(|| value.get("owner_root_node_id"))
            .and_then(Value::as_str)
            .unwrap_or("unowned")
            .to_owned();
        if !matches!(owner.as_str(), "self" | "unowned") {
            canonical_node_id(&owner)?;
        }
        let summaries = value
            .get("connection_summaries")
            .and_then(Value::as_array)
            .into_iter()
            .flatten()
            .filter_map(|summary| Some((summary.get("id")?.as_str()?.to_owned(), summary)))
            .collect::<HashMap<_, _>>();
        let fixed_connections = connections(
            value.get("fixed_connections"),
            &summaries,
            "fixed connection",
        )?;
        let recent_connections = connections(
            value.get("recent_connections"),
            &summaries,
            "recent connection",
        )?;
        let objects = string_ids(value.get("objects"), "object")?;
        Ok(Self {
            id,
            short_name: optional_string(value, "short_name"),
            short_description: optional_string(value, "short_description"),
            long_description: optional_string(value, "long_description"),
            owner,
            fixed_connections,
            recent_connections,
            objects,
            last_modified_by: optional_string(value, "last_modified_by"),
            last_modified_at: value
                .get("last_modified_at")
                .and_then(Value::as_str)
                .map(str::to_owned),
        })
    }

    pub fn draft(&self) -> NodeDraft {
        NodeDraft {
            short_name: self.short_name.clone(),
            short_description: self.short_description.clone(),
            long_description: self.long_description.clone(),
            owner: self.owner.clone(),
            fixed_connections: self
                .fixed_connections
                .iter()
                .map(|connection| connection.id.clone())
                .collect(),
            recent_connections: self
                .recent_connections
                .iter()
                .map(|connection| connection.id.clone())
                .collect(),
            objects: self.objects.clone(),
        }
    }
}

#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct NodeDraft {
    pub short_name: String,
    pub short_description: String,
    pub long_description: String,
    pub owner: String,
    #[serde(default)]
    pub fixed_connections: Vec<String>,
    #[serde(default)]
    pub recent_connections: Vec<String>,
    #[serde(default)]
    pub objects: Vec<String>,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct StagedCreate {
    pub pending_id: String,
    pub data: NodeDraft,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum BoxKind {
    Loaded,
    Fixed,
    Staged,
    Recent,
}

impl BoxKind {
    pub const fn name(self) -> &'static str {
        match self {
            Self::Loaded => "Kweb loaded node",
            Self::Fixed => "Kweb fixed connection",
            Self::Staged => "Kweb staged node",
            Self::Recent => "Kweb recent connections",
        }
    }

    pub const fn metadata_name(self) -> &'static str {
        match self {
            Self::Loaded => "loaded",
            Self::Fixed => "fixed",
            Self::Staged => "staged",
            Self::Recent => "recent",
        }
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
struct BoxSpec {
    logical_slot: String,
    kind: BoxKind,
    text: String,
    stored_node: Option<Node>,
    staged_node: Option<NodeDraft>,
}

#[derive(Clone, Debug, Eq, PartialEq, Serialize)]
#[serde(rename_all = "camelCase")]
pub struct LoadReport {
    pub requested_id: String,
    pub newly_loaded: bool,
    pub promoted_from_fixed: bool,
    pub new_fixed_ids: Vec<String>,
}

/// One current, effect-free provider-facing Kweb context section.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ProjectionItem {
    /// Stable logical identity used to replace a prior rendering.
    pub key: String,
    /// Human-readable section name.
    pub name: String,
    /// Complete current section text.
    pub text: String,
}

#[derive(Clone, Debug)]
pub struct Context {
    root_node_ids: Vec<String>,
    loaded_node_ids: Vec<String>,
    fixed_node_ids: Vec<String>,
    nodes_by_id: BTreeMap<String, Node>,
}

impl Context {
    pub fn new(root_node_ids: Vec<String>) -> Result<Self> {
        if root_node_ids.is_empty() {
            return Err(Error::new("Kweb context requires at least one root node"));
        }
        let mut seen = HashSet::new();
        for id in &root_node_ids {
            canonical_node_id(id)?;
            if !seen.insert(id.clone()) {
                return Err(Error::new("Kweb root node IDs must be distinct"));
            }
        }
        Ok(Self {
            root_node_ids,
            loaded_node_ids: Vec::new(),
            fixed_node_ids: Vec::new(),
            nodes_by_id: BTreeMap::new(),
        })
    }

    pub fn root_node_ids(&self) -> &[String] {
        &self.root_node_ids
    }

    pub fn loaded_node_ids(&self) -> &[String] {
        &self.loaded_node_ids
    }

    pub fn fixed_node_ids(&self) -> &[String] {
        &self.fixed_node_ids
    }

    pub fn full_node_ids(&self) -> Vec<&str> {
        self.loaded_node_ids
            .iter()
            .chain(&self.fixed_node_ids)
            .map(String::as_str)
            .collect()
    }

    pub fn contains_full_node(&self, id: &str) -> bool {
        self.loaded_node_ids.iter().any(|candidate| candidate == id)
            || self.fixed_node_ids.iter().any(|candidate| candidate == id)
    }

    pub fn node(&self, id: &str) -> Option<&Node> {
        self.nodes_by_id.get(id)
    }

    pub fn apply_load(&mut self, requested: Node, fixed: Vec<Node>) -> Result<LoadReport> {
        let requested_id = requested.id.clone();
        let was_loaded = self
            .loaded_node_ids
            .iter()
            .any(|candidate| candidate == &requested_id);
        let was_fixed = self
            .fixed_node_ids
            .iter()
            .any(|candidate| candidate == &requested_id);
        let previous_full = self
            .full_node_ids()
            .into_iter()
            .map(str::to_owned)
            .collect::<HashSet<_>>();
        let expected_fixed = requested
            .fixed_connections
            .iter()
            .map(|connection| connection.id.as_str())
            .filter(|id| *id != requested_id)
            .collect::<HashSet<_>>();
        let provided_fixed = fixed
            .iter()
            .map(|node| node.id.as_str())
            .collect::<HashSet<_>>();
        if expected_fixed != provided_fixed {
            return Err(Error::new(format!(
                "load for {requested_id} did not provide exactly its fixed connections"
            )));
        }
        self.nodes_by_id.insert(requested_id.clone(), requested);
        for node in fixed {
            self.nodes_by_id.insert(node.id.clone(), node);
        }
        if !was_loaded {
            self.loaded_node_ids.push(requested_id.clone());
        }
        self.rebuild_fixed();
        let new_fixed_ids = self
            .fixed_node_ids
            .iter()
            .filter(|id| !previous_full.contains(*id))
            .cloned()
            .collect();
        Ok(LoadReport {
            requested_id,
            newly_loaded: !was_loaded && !was_fixed,
            promoted_from_fixed: !was_loaded && was_fixed,
            new_fixed_ids,
        })
    }

    pub fn refresh(&mut self, nodes: impl IntoIterator<Item = Node>) -> Result<()> {
        for node in nodes {
            if !self.contains_full_node(&node.id) {
                return Err(Error::new(format!(
                    "cannot refresh unloaded Kweb node {}",
                    node.id
                )));
            }
            self.nodes_by_id.insert(node.id.clone(), node);
        }
        self.rebuild_fixed();
        Ok(())
    }

    pub fn restore(
        &mut self,
        nodes: impl IntoIterator<Item = Node>,
        directly_loaded: Vec<String>,
    ) -> Result<()> {
        self.nodes_by_id.clear();
        for node in nodes {
            self.nodes_by_id.insert(node.id.clone(), node);
        }
        let mut seen = HashSet::new();
        self.loaded_node_ids = directly_loaded
            .into_iter()
            .filter(|id| self.nodes_by_id.contains_key(id) && seen.insert(id.clone()))
            .collect();
        if !self.nodes_by_id.is_empty() && self.loaded_node_ids.is_empty() {
            return Err(Error::new(
                "restored Kweb context contains nodes but no loaded node",
            ));
        }
        self.rebuild_fixed();
        Ok(())
    }

    fn box_specs(
        &self,
        updates: &BTreeMap<String, NodeDraft>,
        creates: &[StagedCreate],
    ) -> Result<Vec<BoxSpec>> {
        for id in updates.keys() {
            if !self.contains_full_node(id) {
                return Err(Error::new(format!(
                    "staged update targets unloaded Kweb node {id}"
                )));
            }
        }
        let mut specs = Vec::new();
        for id in &self.loaded_node_ids {
            specs.push(self.full_box(id, BoxKind::Loaded, updates.get(id))?);
        }
        for id in &self.fixed_node_ids {
            specs.push(self.full_box(id, BoxKind::Fixed, updates.get(id))?);
        }
        for create in creates {
            specs.push(BoxSpec {
                logical_slot: create.pending_id.clone(),
                kind: BoxKind::Staged,
                text: format_node(&create.pending_id, &create.data),
                stored_node: None,
                staged_node: Some(create.data.clone()),
            });
        }
        specs.push(BoxSpec {
            logical_slot: "recent-connections".into(),
            kind: BoxKind::Recent,
            text: self.format_recent_connections(updates, creates)?,
            stored_node: None,
            staged_node: None,
        });
        Ok(specs)
    }

    /// Renders the complete current Kweb projection without touching Chatend.
    ///
    /// Unlike [`Self::sync_chatend`], this view has no durable representation
    /// history. It emits one current recent-connections section and is intended
    /// for ephemeral consumers such as box-free subagents.
    pub fn projection(
        &self,
        updates: &BTreeMap<String, NodeDraft>,
        creates: &[StagedCreate],
    ) -> Result<Vec<ProjectionItem>> {
        self.box_specs(updates, creates)?
            .into_iter()
            .map(|spec| {
                Ok(ProjectionItem {
                    key: spec.logical_slot,
                    name: spec.kind.name().to_owned(),
                    text: spec.text,
                })
            })
            .collect()
    }

    /// Reconcile this context's complete provider-facing projection through an
    /// already-open durable Session History handle.
    ///
    /// The returned IDs are the active Kweb boxes whose name or canonical
    /// revision changed, in current projection order.
    pub fn sync_chatend(
        &self,
        journal: &mut HistorySession,
        recorded_at: impl Into<String>,
        updates: &BTreeMap<String, NodeDraft>,
        creates: &[StagedCreate],
    ) -> Result<Vec<BoxId>> {
        let recorded_at = recorded_at.into();
        let previous = kweb_box_versions(journal);
        let specs = self.box_specs(updates, creates)?;
        let mut desired = Vec::with_capacity(specs.len());
        let mut recent = None;
        for spec in specs {
            let mut metadata = json!({
                "revisionHash": revision_hash(&spec.text),
            });
            if let Some(node) = spec.stored_node {
                metadata["canonicalNodeId"] = json!(node.id);
                metadata["storedNode"] = serde_json::to_value(node).map_err(|error| {
                    Error::new(format!("serializing stored Kweb node: {error}"))
                })?;
            }
            if let Some(node) = spec.staged_node {
                metadata["staged"] = json!(true);
                metadata["nodeData"] = serde_json::to_value(node).map_err(|error| {
                    Error::new(format!("serializing staged Kweb node: {error}"))
                })?;
            }
            let mut content = BoxContent {
                text: spec.text,
                objects: Vec::new(),
                metadata,
            };
            content.use_concise_header();
            mark_kweb_content(&mut content, &spec.logical_slot, spec.kind.metadata_name());
            let entry = DesiredKwebBox {
                logical_slot: spec.logical_slot,
                name: spec.kind.name().into(),
                content,
            };
            if spec.kind == BoxKind::Recent {
                if recent.replace(entry).is_some() {
                    return Err(Error::new(
                        "Kweb context produced more than one recent-connections candidate",
                    ));
                }
            } else {
                desired.push(entry);
            }
        }
        let recent = recent
            .ok_or_else(|| Error::new("Kweb context produced no recent-connections candidate"))?;
        desired.extend(desired_recent_connection_boxes(journal, recent)?);
        reconcile_kweb_slots(journal, &recorded_at, desired)?;
        Ok(changed_kweb_box_ids(journal, &previous))
    }

    fn full_box(&self, id: &str, kind: BoxKind, update: Option<&NodeDraft>) -> Result<BoxSpec> {
        let node = self
            .nodes_by_id
            .get(id)
            .ok_or_else(|| Error::new(format!("missing full Kweb node {id}")))?;
        let data = update.cloned().unwrap_or_else(|| node.draft());
        Ok(BoxSpec {
            logical_slot: id.to_owned(),
            kind,
            text: format_node(id, &data),
            stored_node: Some(node.clone()),
            staged_node: update.cloned(),
        })
    }

    fn format_recent_connections(
        &self,
        updates: &BTreeMap<String, NodeDraft>,
        creates: &[StagedCreate],
    ) -> Result<String> {
        let creates_by_id = creates
            .iter()
            .map(|create| (create.pending_id.as_str(), &create.data))
            .collect::<HashMap<_, _>>();
        let mut summaries = HashMap::new();
        for node in self.nodes_by_id.values() {
            summaries.insert(
                node.id.as_str(),
                (node.short_name.as_str(), node.short_description.as_str()),
            );
            for connection in node
                .fixed_connections
                .iter()
                .chain(&node.recent_connections)
            {
                summaries.entry(connection.id.as_str()).or_insert((
                    connection.short_name.as_str(),
                    connection.short_description.as_str(),
                ));
            }
        }
        let mut recent_ids = Vec::new();
        let mut seen = HashSet::new();
        for id in self.full_node_ids() {
            let node = self
                .nodes_by_id
                .get(id)
                .ok_or_else(|| Error::new(format!("missing full Kweb node {id}")))?;
            let recent = updates
                .get(id)
                .map(|draft| draft.recent_connections.as_slice())
                .unwrap_or_else(|| &[]);
            if updates.contains_key(id) {
                for connection_id in recent {
                    if seen.insert(connection_id.clone()) {
                        recent_ids.push(connection_id.clone());
                    }
                }
            } else {
                for connection in &node.recent_connections {
                    if seen.insert(connection.id.clone()) {
                        recent_ids.push(connection.id.clone());
                    }
                }
            }
        }
        for create in creates {
            for connection_id in &create.data.recent_connections {
                if seen.insert(connection_id.clone()) {
                    recent_ids.push(connection_id.clone());
                }
            }
        }
        let mut lines = vec!["Recent connections".to_owned()];
        for id in recent_ids {
            let staged_summary = updates
                .get(&id)
                .or_else(|| creates_by_id.get(id.as_str()).copied())
                .map(|node| (node.short_name.as_str(), node.short_description.as_str()));
            let (name, description) = staged_summary
                .or_else(|| summaries.get(id.as_str()).copied())
                .ok_or_else(|| {
                    Error::new(format!(
                        "recent connection {id} must resolve to a nonempty short name and short description"
                    ))
                })?;
            if name.trim().is_empty() || description.trim().is_empty() {
                return Err(Error::new(format!(
                    "recent connection {id} must resolve to a nonempty short name and short description"
                )));
            }
            lines.push(format!("{id} · {name}: {description}"));
        }
        if lines.len() == 1 {
            lines.push("None.".into());
        }
        Ok(lines.join("\n"))
    }

    fn rebuild_fixed(&mut self) {
        let loaded = self.loaded_node_ids.iter().cloned().collect::<HashSet<_>>();
        let mut seen = loaded.clone();
        let mut fixed = Vec::new();
        for id in &self.loaded_node_ids {
            let Some(node) = self.nodes_by_id.get(id) else {
                continue;
            };
            for connection in &node.fixed_connections {
                if self.nodes_by_id.contains_key(&connection.id)
                    && seen.insert(connection.id.clone())
                {
                    fixed.push(connection.id.clone());
                }
            }
        }
        self.fixed_node_ids = fixed;
        self.nodes_by_id
            .retain(|id, _| loaded.contains(id) || seen.contains(id));
    }
}

struct DesiredKwebBox {
    logical_slot: String,
    name: String,
    content: BoxContent,
}

#[derive(Clone, Debug, Eq, PartialEq)]
struct RecentConnectionEntry {
    id: String,
    text: String,
}

fn mark_kweb_content(content: &mut BoxContent, logical_slot: &str, role: &str) {
    if !content.metadata.is_object() {
        content.metadata = json!({});
    }
    content.metadata["kwebLogicalSlot"] = json!(logical_slot);
    content.metadata["kwebRole"] = json!(role);
}

fn kweb_logical_slot(state: &BoxState, actual_slot: &str) -> String {
    state
        .canonical
        .content
        .metadata
        .get("kwebLogicalSlot")
        .and_then(Value::as_str)
        .unwrap_or(actual_slot)
        .to_owned()
}

fn recent_connection_entries(content: &BoxContent) -> Result<Vec<RecentConnectionEntry>> {
    let body = content
        .text
        .strip_prefix("Recent connections")
        .ok_or_else(|| Error::new("Kweb recent-connections box has an invalid heading"))?;
    let body = body
        .strip_prefix('\n')
        .ok_or_else(|| Error::new("Kweb recent-connections box has no body"))?;
    if body == "None." {
        return Ok(Vec::new());
    }

    let mut entries: Vec<RecentConnectionEntry> = Vec::new();
    for line in body.split('\n') {
        let identifier = line.split_once(" · ").and_then(|(identifier, _)| {
            let canonical = identifier.parse::<NodeId>().is_ok();
            let pending = PendingId::parse(identifier.to_owned()).is_ok();
            (canonical || pending).then_some(identifier)
        });
        if let Some(identifier) = identifier {
            entries.push(RecentConnectionEntry {
                id: identifier.to_owned(),
                text: line.to_owned(),
            });
        } else {
            let entry = entries.last_mut().ok_or_else(|| {
                Error::new("Kweb recent-connections box starts with invalid entry text")
            })?;
            entry.text.push('\n');
            entry.text.push_str(line);
        }
    }

    if let Some(expected) = content
        .metadata
        .get(RECENT_CONNECTION_IDS_METADATA)
        .and_then(Value::as_array)
    {
        let expected = expected
            .iter()
            .map(|value| {
                value.as_str().ok_or_else(|| {
                    Error::new("Kweb recent-connection IDs metadata contains a non-string value")
                })
            })
            .collect::<Result<Vec<_>>>()?;
        if expected
            != entries
                .iter()
                .map(|entry| entry.id.as_str())
                .collect::<Vec<_>>()
        {
            return Err(Error::new(
                "Kweb recent-connection IDs metadata does not match its canonical text",
            ));
        }
    }
    Ok(entries)
}

fn format_recent_connection_entries(entries: &[RecentConnectionEntry]) -> String {
    if entries.is_empty() {
        return "Recent connections\nNone.".into();
    }
    format!(
        "Recent connections\n{}",
        entries
            .iter()
            .map(|entry| entry.text.as_str())
            .collect::<Vec<_>>()
            .join("\n")
    )
}

fn revision_hash(text: &str) -> String {
    hex::encode(Sha256::digest(text.as_bytes()))
}

fn update_recent_connection_content(
    content: &mut BoxContent,
    logical_slot: &str,
    entries: &[RecentConnectionEntry],
) {
    content.text = format_recent_connection_entries(entries);
    mark_kweb_content(content, logical_slot, "recent");
    content.metadata["revisionHash"] = json!(revision_hash(&content.text));
    content.metadata[RECENT_CONNECTION_IDS_METADATA] = json!(
        entries
            .iter()
            .map(|entry| entry.id.as_str())
            .collect::<Vec<_>>()
    );
}

fn desired_recent_connection_boxes(
    journal: &HistorySession,
    fresh: DesiredKwebBox,
) -> Result<Vec<DesiredKwebBox>> {
    let mut boxes = Vec::new();
    let mut seen = HashSet::new();
    let mut used_logical_slots = HashSet::new();

    if let Some(tool) = journal.state().tools.get(KWEB_TOOL_INSTANCE) {
        for slot in &tool.slots {
            let state = journal
                .state()
                .box_state(slot.box_id)
                .ok_or_else(|| Error::new("Kweb tool slot box is missing"))?;
            let logical_slot = kweb_logical_slot(state, &slot.slot);
            used_logical_slots.insert(logical_slot.clone());
            if slot.retired
                || state
                    .canonical
                    .content
                    .metadata
                    .get("kwebRole")
                    .and_then(Value::as_str)
                    != Some("recent")
            {
                continue;
            }
            let entries = recent_connection_entries(&state.canonical.content)?;
            for entry in &entries {
                seen.insert(entry.id.clone());
            }
            boxes.push((
                DesiredKwebBox {
                    logical_slot,
                    name: state.name.clone(),
                    content: state.canonical.content.clone(),
                },
                entries,
            ));
        }
    }

    let additions = recent_connection_entries(&fresh.content)?
        .into_iter()
        .filter(|entry| seen.insert(entry.id.clone()))
        .collect::<Vec<_>>();
    let mut next_addition = 0;

    if let Some((last, entries)) = boxes.last_mut()
        && entries.len() < RECENT_CONNECTIONS_PER_BOX
    {
        let available = RECENT_CONNECTIONS_PER_BOX - entries.len();
        let end = additions.len().min(available);
        entries.extend_from_slice(&additions[..end]);
        next_addition = end;
        if end > 0 {
            update_recent_connection_content(&mut last.content, &last.logical_slot, entries);
        }
    }

    while next_addition < additions.len() || boxes.is_empty() {
        let end = (next_addition + RECENT_CONNECTIONS_PER_BOX).min(additions.len());
        let entries = additions[next_addition..end].to_vec();
        let mut sequence = boxes.len() + 1;
        let logical_slot = loop {
            let candidate = if sequence == 1 {
                RECENT_CONNECTIONS_LOGICAL_SLOT.to_owned()
            } else {
                format!("{RECENT_CONNECTIONS_LOGICAL_SLOT}:{sequence}")
            };
            if used_logical_slots.insert(candidate.clone()) {
                break candidate;
            }
            sequence += 1;
        };
        let mut content = fresh.content.clone();
        update_recent_connection_content(&mut content, &logical_slot, &entries);
        boxes.push((
            DesiredKwebBox {
                logical_slot,
                name: fresh.name.clone(),
                content,
            },
            entries,
        ));
        next_addition = end;
    }

    Ok(boxes.into_iter().map(|(box_spec, _)| box_spec).collect())
}

type KwebBoxVersions = BTreeMap<BoxId, (String, EventId)>;

fn kweb_box_versions(journal: &HistorySession) -> KwebBoxVersions {
    journal
        .state()
        .tool_layouts
        .get(KWEB_TOOL_INSTANCE)
        .into_iter()
        .flatten()
        .filter_map(|box_id| {
            let state = journal.state().box_state(*box_id)?;
            state
                .active
                .then(|| (*box_id, (state.name.clone(), state.canonical.event_id)))
        })
        .collect()
}

fn changed_kweb_box_ids(journal: &HistorySession, previous: &KwebBoxVersions) -> Vec<BoxId> {
    journal
        .state()
        .tool_layouts
        .get(KWEB_TOOL_INSTANCE)
        .into_iter()
        .flatten()
        .filter_map(|box_id| {
            let state = journal.state().box_state(*box_id)?;
            let current = (state.name.as_str(), state.canonical.event_id);
            let changed = previous
                .get(box_id)
                .map(|(name, revision)| (name.as_str(), *revision) != current)
                .unwrap_or(true);
            (state.active && changed).then_some(*box_id)
        })
        .collect()
}

fn unique_slot(logical: &str, used: &mut HashSet<String>) -> String {
    if used.insert(logical.to_owned()) {
        return logical.to_owned();
    }
    let mut generation = 2_u64;
    loop {
        let candidate = format!("{logical}#generation-{generation}");
        if used.insert(candidate.clone()) {
            return candidate;
        }
        generation += 1;
    }
}

fn reconcile_kweb_slots(
    journal: &mut HistorySession,
    recorded_at: &str,
    desired: Vec<DesiredKwebBox>,
) -> Result<()> {
    let current = journal
        .state()
        .tools
        .get(KWEB_TOOL_INSTANCE)
        .cloned()
        .unwrap_or_default();
    let desired_by_logical = desired
        .iter()
        .enumerate()
        .map(|(index, entry)| (entry.logical_slot.as_str(), index))
        .collect::<BTreeMap<_, _>>();
    if desired_by_logical.len() != desired.len() {
        return Err(Error::new(
            "Kweb box layout contains duplicate logical slots",
        ));
    }
    let mut claimed = HashSet::new();
    let mut actual_by_desired = BTreeMap::new();
    let mut slots = Vec::with_capacity(current.slots.len() + desired.len());
    let mut used_actual = current
        .slots
        .iter()
        .map(|slot| slot.slot.clone())
        .collect::<HashSet<_>>();
    for slot in &current.slots {
        let state = journal
            .state()
            .box_state(slot.box_id)
            .ok_or_else(|| Error::new("Kweb tool slot box is missing"))?;
        let logical = kweb_logical_slot(state, &slot.slot);
        let selected = !slot.retired
            && desired_by_logical.contains_key(logical.as_str())
            && claimed.insert(logical.clone());
        if selected {
            let entry = &desired[desired_by_logical[logical.as_str()]];
            slots.push(ToolSlotInput {
                slot: slot.slot.clone(),
                name: entry.name.clone(),
                content: entry.content.clone(),
                retired: false,
            });
            actual_by_desired.insert(entry.logical_slot.clone(), slot.slot.clone());
        } else {
            slots.push(ToolSlotInput {
                slot: slot.slot.clone(),
                name: state.name.clone(),
                content: state.canonical.content.clone(),
                retired: slot.retired || !selected,
            });
        }
    }
    for entry in &desired {
        if actual_by_desired.contains_key(&entry.logical_slot) {
            continue;
        }
        let actual = unique_slot(&entry.logical_slot, &mut used_actual);
        slots.push(ToolSlotInput {
            slot: actual.clone(),
            name: entry.name.clone(),
            content: entry.content.clone(),
            retired: false,
        });
        actual_by_desired.insert(entry.logical_slot.clone(), actual);
    }
    let layout_slots = desired
        .iter()
        .map(|entry| actual_by_desired[&entry.logical_slot].clone())
        .collect::<Vec<_>>();
    journal
        .apply_tool_slots_with_layout(recorded_at, KWEB_TOOL_INSTANCE, slots, &layout_slots)
        .map_err(|error| Error::new(format!("applying Kweb projection: {error}")))?;
    Ok(())
}

fn connections(
    value: Option<&Value>,
    summaries: &HashMap<String, &Value>,
    label: &str,
) -> Result<Vec<Connection>> {
    let mut result = Vec::new();
    let mut seen = HashSet::new();
    for entry in value.and_then(Value::as_array).into_iter().flatten() {
        let id = entry
            .as_str()
            .or_else(|| entry.get("id").and_then(Value::as_str))
            .ok_or_else(|| Error::new(format!("{label} has no node ID")))?
            .to_owned();
        canonical_node_id(&id)?;
        if !seen.insert(id.clone()) {
            continue;
        }
        let summary = summaries.get(&id).copied();
        result.push(Connection {
            id,
            short_name: entry
                .get("short_name")
                .and_then(Value::as_str)
                .or_else(|| summary.and_then(|value| value.get("short_name")?.as_str()))
                .unwrap_or_default()
                .to_owned(),
            short_description: entry
                .get("short_description")
                .and_then(Value::as_str)
                .or_else(|| summary.and_then(|value| value.get("short_description")?.as_str()))
                .unwrap_or_default()
                .to_owned(),
        });
    }
    Ok(result)
}

fn string_ids(value: Option<&Value>, label: &str) -> Result<Vec<String>> {
    let mut result = Vec::new();
    let mut seen = HashSet::new();
    for entry in value.and_then(Value::as_array).into_iter().flatten() {
        let id = entry
            .as_str()
            .ok_or_else(|| Error::new(format!("{label} ID must be a string")))?
            .to_owned();
        if seen.insert(id.clone()) {
            result.push(id);
        }
    }
    Ok(result)
}

fn canonical_node_id(value: &str) -> Result<()> {
    value
        .parse::<NodeId>()
        .map(|_| ())
        .map_err(|_| Error::new(format!("{value:?} is not a canonical Kweb node ID")))
}

fn required_string(value: &Value, key: &str) -> Result<String> {
    value
        .get(key)
        .and_then(Value::as_str)
        .map(str::to_owned)
        .ok_or_else(|| Error::new(format!("Kweb node has no string {key}")))
}

fn optional_string(value: &Value, key: &str) -> String {
    value
        .get(key)
        .and_then(Value::as_str)
        .unwrap_or_default()
        .to_owned()
}

fn format_node(identifier: &str, node: &NodeDraft) -> String {
    [
        format!("Node ID: {identifier}"),
        format!("Node name: {}", fallback(&node.short_name)),
        format!("Node summary: {}", fallback(&node.short_description)),
        format!("Node owner ID: {}", fallback(&node.owner)),
        "Node long description:".into(),
        indent(&node.long_description),
        format!(
            "Fixed connection IDs: {}",
            list_or_none(&node.fixed_connections)
        ),
        format!(
            "Recent connection IDs: {}",
            list_or_none(&node.recent_connections)
        ),
    ]
    .join("\n")
}

fn indent(value: &str) -> String {
    fallback(value)
        .lines()
        .map(|line| format!("  {line}"))
        .collect::<Vec<_>>()
        .join("\n")
}

fn fallback(value: &str) -> &str {
    if value.trim().is_empty() {
        "(none)"
    } else {
        value
    }
}

fn list_or_none(values: &[String]) -> String {
    if values.is_empty() {
        "none".into()
    } else {
        values.join(", ")
    }
}

#[cfg(test)]
mod tests {
    use std::path::PathBuf;
    use std::time::{SystemTime, UNIX_EPOCH};

    use super::*;
    use kcode_session_history::{
        Config as HistoryConfig, NewSession, SessionHistory,
        chatend::{Representation, SessionKind},
    };
    use serde_json::json;

    fn id(index: u8) -> String {
        NodeId::from_bytes([0, 0, 0, 0, 0, index])
            .unwrap()
            .to_string()
    }

    fn connection(index: u8) -> Connection {
        Connection {
            id: id(index),
            short_name: format!("Node {index}"),
            short_description: format!("Summary {index}"),
        }
    }

    fn node(index: u8, fixed: &[u8], recent: &[u8]) -> Node {
        Node {
            id: id(index),
            short_name: format!("Node {index}"),
            short_description: format!("Summary {index}"),
            long_description: format!("Long description {index}"),
            owner: id(1),
            fixed_connections: fixed.iter().copied().map(connection).collect(),
            recent_connections: recent.iter().copied().map(connection).collect(),
            objects: vec![],
            last_modified_by: "test-model-high".into(),
            last_modified_at: Some("2026-07-28T00:00:00Z".into()),
        }
    }

    fn node_with_recent_description(
        index: u8,
        fixed: &[u8],
        recent: &[u8],
        description: &str,
    ) -> Node {
        let mut node = node(index, fixed, recent);
        for (connection, connection_index) in node.recent_connections.iter_mut().zip(recent) {
            connection.short_description = format!("{description} {connection_index}");
        }
        node
    }

    fn draft(index: u8, recent: &[u8]) -> NodeDraft {
        NodeDraft {
            short_name: format!("Node {index}"),
            short_description: format!("Summary {index}"),
            long_description: format!("Long description {index}"),
            owner: id(1),
            fixed_connections: Vec::new(),
            recent_connections: recent.iter().map(|value| id(*value)).collect(),
            objects: Vec::new(),
        }
    }

    fn test_journal(label: &str) -> (PathBuf, HistorySession) {
        let root = std::env::temp_dir().join(format!(
            "kcode-kweb-context-{label}-{}-{}",
            std::process::id(),
            SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap()
                .as_nanos()
        ));
        let history = SessionHistory::open(HistoryConfig {
            directory: root.join("sessions"),
            completed_list: root.join("completed.jsonl"),
            provider_cost_compatibility: None,
        })
        .unwrap();
        let journal = history
            .create_session(NewSession {
                kind: SessionKind::Conversation,
                created_at: "2026-07-29T00:00:00Z".into(),
                effective_context_tokens: 10_000,
                channel: Value::Null,
            })
            .unwrap();
        (root, journal)
    }

    fn recent_box_ids(journal: &HistorySession) -> Vec<BoxId> {
        journal
            .state()
            .tool_layouts
            .get(KWEB_TOOL_INSTANCE)
            .into_iter()
            .flatten()
            .copied()
            .filter(|box_id| {
                journal
                    .state()
                    .box_state(*box_id)
                    .and_then(|state| state.canonical.content.metadata.get("kwebRole"))
                    .and_then(Value::as_str)
                    == Some("recent")
            })
            .collect()
    }

    #[test]
    fn parses_the_kweb_wire_shape_into_typed_connections() {
        let parsed = Node::from_kweb_value(&json!({
            "id": id(1),
            "owner_node_id": id(1),
            "short_name": "Root",
            "short_description": "Root summary",
            "long_description": "Root details",
            "fixed_connections": [id(2)],
            "recent_connections": [id(3), id(3)],
            "objects": [],
            "connection_summaries": [
                {"id":id(2),"short_name":"Fixed","short_description":"Fixed summary"},
                {"id":id(3),"short_name":"Recent","short_description":"Recent summary"}
            ]
        }))
        .unwrap();
        assert_eq!(
            parsed.fixed_connections,
            vec![Connection {
                id: id(2),
                short_name: "Fixed".into(),
                short_description: "Fixed summary".into(),
            }]
        );
        assert_eq!(parsed.recent_connections.len(), 1);
        assert_eq!(parsed.recent_connections[0].short_name, "Recent");
    }

    #[test]
    fn loaded_nodes_take_precedence_over_the_fixed_role() {
        let mut context = Context::new(vec![id(1)]).unwrap();
        context
            .apply_load(node(1, &[2], &[]), vec![node(2, &[], &[])])
            .unwrap();
        let report = context.apply_load(node(2, &[], &[]), Vec::new()).unwrap();
        assert!(report.promoted_from_fixed);
        assert_eq!(context.loaded_node_ids(), &[id(1), id(2)]);
        assert!(context.fixed_node_ids().is_empty());
        assert_eq!(
            context
                .box_specs(&BTreeMap::new(), &[])
                .unwrap()
                .iter()
                .map(|spec| spec.kind)
                .collect::<Vec<_>>(),
            vec![BoxKind::Loaded, BoxKind::Loaded, BoxKind::Recent]
        );
    }

    #[test]
    fn effect_free_projection_has_stable_keys_and_complete_current_text() {
        let mut context = Context::new(vec![id(1)]).unwrap();
        context
            .apply_load(node(1, &[2], &[3]), vec![node(2, &[], &[])])
            .unwrap();

        let projected = context.projection(&BTreeMap::new(), &[]).unwrap();
        assert_eq!(
            projected
                .iter()
                .map(|item| item.key.as_str())
                .collect::<Vec<_>>(),
            vec![id(1), id(2), "recent-connections".to_owned()]
        );
        assert_eq!(projected[0].name, "Kweb loaded node");
        assert!(projected[0].text.contains("Long description 1"));
        assert_eq!(projected[2].name, "Kweb recent connections");
        assert!(projected[2].text.contains(&id(3)));
    }

    #[test]
    fn full_node_kinds_share_one_body_format_without_active_connections() {
        let mut context = Context::new(vec![id(1)]).unwrap();
        context
            .apply_load(node(1, &[2], &[]), vec![node(2, &[], &[])])
            .unwrap();
        let create = StagedCreate {
            pending_id: "pending:1".into(),
            data: draft(3, &[]),
        };
        let specs = context.box_specs(&BTreeMap::new(), &[create]).unwrap();
        assert_eq!(specs[0].kind.name(), "Kweb loaded node");
        assert_eq!(specs[1].kind.name(), "Kweb fixed connection");
        assert_eq!(specs[2].kind.name(), "Kweb staged node");
        for spec in &specs[..3] {
            assert!(spec.text.contains("Node ID:"));
            assert!(spec.text.contains("Node name:"));
            assert!(spec.text.contains("Node owner ID:"));
            assert!(spec.text.contains("Fixed connection IDs:"));
            assert!(spec.text.contains("Recent connection IDs:"));
            assert!(!spec.text.contains("Active"));
        }
        assert_eq!(
            specs[0].text,
            concat!(
                "Node ID: AAAAAAAB\n",
                "Node name: Node 1\n",
                "Node summary: Summary 1\n",
                "Node owner ID: AAAAAAAB\n",
                "Node long description:\n",
                "  Long description 1\n",
                "Fixed connection IDs: AAAAAAAC\n",
                "Recent connection IDs: none"
            )
        );
    }

    #[test]
    fn all_recent_connections_share_one_globally_deduplicated_box() {
        let mut context = Context::new(vec![id(1)]).unwrap();
        context
            .apply_load(node(1, &[2], &[4, 5]), vec![node(2, &[], &[5, 6, 7])])
            .unwrap();
        let creates = vec![StagedCreate {
            pending_id: "pending:1".into(),
            data: draft(3, &[6, 7]),
        }];
        let specs = context.box_specs(&BTreeMap::new(), &creates).unwrap();
        let recent = specs
            .iter()
            .filter(|spec| spec.kind == BoxKind::Recent)
            .collect::<Vec<_>>();
        assert_eq!(recent.len(), 1);
        assert_eq!(
            recent[0].text,
            format!(
                concat!(
                    "Recent connections\n",
                    "{} · Node 4: Summary 4\n",
                    "{} · Node 5: Summary 5\n",
                    "{} · Node 6: Summary 6\n",
                    "{} · Node 7: Summary 7"
                ),
                id(4),
                id(5),
                id(6),
                id(7)
            )
        );
    }

    #[test]
    fn empty_recent_projection_is_still_one_exact_box() {
        let mut context = Context::new(vec![id(1)]).unwrap();
        context.apply_load(node(1, &[], &[]), Vec::new()).unwrap();
        let specs = context.box_specs(&BTreeMap::new(), &[]).unwrap();
        assert_eq!(specs.len(), 2);
        assert_eq!(specs[1].kind, BoxKind::Recent);
        assert_eq!(specs[1].text, "Recent connections\nNone.");
    }

    #[test]
    fn recent_projection_rejects_missing_name_or_description() {
        for missing_name in [true, false] {
            let mut source = node(1, &[], &[2]);
            if missing_name {
                source.recent_connections[0].short_name.clear();
            } else {
                source.recent_connections[0].short_description.clear();
            }
            let mut context = Context::new(vec![id(1)]).unwrap();
            context.apply_load(source, Vec::new()).unwrap();
            assert_eq!(
                context
                    .box_specs(&BTreeMap::new(), &[])
                    .unwrap_err()
                    .to_string(),
                format!(
                    "recent connection {} must resolve to a nonempty short name and short description",
                    id(2)
                )
            );
        }
    }

    #[test]
    fn staged_updates_drive_full_text_and_recent_projection() {
        let mut context = Context::new(vec![id(1)]).unwrap();
        context
            .apply_load(node(1, &[3], &[2]), vec![node(3, &[], &[])])
            .unwrap();
        let mut updates = BTreeMap::new();
        updates.insert(id(1), draft(9, &[3]));
        let specs = context.box_specs(&updates, &[]).unwrap();
        assert!(specs[0].text.contains("Node name: Node 9"));
        assert!(specs[0].staged_node.is_some());
        assert!(!specs.last().unwrap().text.contains(&id(2)));
        assert!(specs.last().unwrap().text.contains(&id(3)));
    }

    #[test]
    fn sync_fills_permanent_recent_boxes_eight_at_a_time() {
        let (root, mut journal) = test_journal("recent-boxes");
        let mut context = Context::new(vec![id(1)]).unwrap();
        let initial_indices = (2..=19).collect::<Vec<_>>();
        context
            .apply_load(
                node_with_recent_description(1, &[], &initial_indices, "old"),
                Vec::new(),
            )
            .unwrap();
        context
            .sync_chatend(&mut journal, "t1", &BTreeMap::new(), &[])
            .unwrap();

        let original_ids = recent_box_ids(&journal);
        assert_eq!(
            original_ids
                .iter()
                .map(|box_id| {
                    recent_connection_entries(
                        &journal
                            .state()
                            .box_state(*box_id)
                            .unwrap()
                            .canonical
                            .content,
                    )
                    .unwrap()
                    .len()
                })
                .collect::<Vec<_>>(),
            vec![8, 8, 2]
        );
        let original_revisions = original_ids
            .iter()
            .map(|box_id| {
                journal
                    .state()
                    .box_state(*box_id)
                    .unwrap()
                    .canonical
                    .event_id
            })
            .collect::<Vec<_>>();
        journal
            .summarize_box("t2", original_ids[0], "retained first box")
            .unwrap();
        journal.dehydrate_boxes("t3", &original_ids[1..=2]).unwrap();

        let expanded_indices = (2..=28).collect::<Vec<_>>();
        context
            .refresh([node_with_recent_description(
                1,
                &[],
                &expanded_indices,
                "new",
            )])
            .unwrap();
        let changed = context
            .sync_chatend(&mut journal, "t4", &BTreeMap::new(), &[])
            .unwrap();
        let current_ids = recent_box_ids(&journal);
        assert_eq!(
            current_ids
                .iter()
                .map(|box_id| {
                    recent_connection_entries(
                        &journal
                            .state()
                            .box_state(*box_id)
                            .unwrap()
                            .canonical
                            .content,
                    )
                    .unwrap()
                    .len()
                })
                .collect::<Vec<_>>(),
            vec![8, 8, 8, 3]
        );
        assert_eq!(&current_ids[..3], original_ids.as_slice());
        assert!(!changed.contains(&original_ids[0]));
        assert!(!changed.contains(&original_ids[1]));
        assert!(changed.contains(&original_ids[2]));
        assert!(changed.contains(&current_ids[3]));

        let first = journal.state().box_state(original_ids[0]).unwrap();
        assert_eq!(first.canonical.event_id, original_revisions[0]);
        assert!(matches!(
            first.representation,
            Representation::Summarized { based_on, .. } if based_on == first.canonical.event_id
        ));
        let second = journal.state().box_state(original_ids[1]).unwrap();
        assert_eq!(second.canonical.event_id, original_revisions[1]);
        assert!(matches!(
            second.representation,
            Representation::Dehydrated { based_on } if based_on == second.canonical.event_id
        ));
        let third = journal.state().box_state(original_ids[2]).unwrap();
        assert_ne!(third.canonical.event_id, original_revisions[2]);
        assert!(third.canonical.content.text.contains("old 19"));
        assert!(third.canonical.content.text.contains("new 25"));
        assert!(!third.canonical.content.text.contains("new 19"));
        assert!(matches!(
            third.representation,
            Representation::Dehydrated { based_on } if based_on == original_revisions[2]
        ));

        drop(journal);
        std::fs::remove_dir_all(root).unwrap();
    }

    #[test]
    fn sync_starts_one_empty_fillable_recent_box() {
        let (root, mut journal) = test_journal("empty-recent-box");
        let mut context = Context::new(vec![id(1)]).unwrap();
        context.apply_load(node(1, &[], &[]), Vec::new()).unwrap();
        context
            .sync_chatend(&mut journal, "t1", &BTreeMap::new(), &[])
            .unwrap();

        let boxes = recent_box_ids(&journal);
        assert_eq!(boxes.len(), 1);
        let content = &journal
            .state()
            .box_state(boxes[0])
            .unwrap()
            .canonical
            .content;
        assert!(recent_connection_entries(content).unwrap().is_empty());
        assert_eq!(content.metadata[RECENT_CONNECTION_IDS_METADATA], json!([]));

        drop(journal);
        std::fs::remove_dir_all(root).unwrap();
    }

    #[test]
    fn sync_reports_only_changed_boxes_in_projection_order() {
        let (root, mut journal) = test_journal("changed-boxes");
        let mut context = Context::new(vec![id(1)]).unwrap();
        context.apply_load(node(1, &[], &[]), Vec::new()).unwrap();
        let initial = context
            .sync_chatend(&mut journal, "t1", &BTreeMap::new(), &[])
            .unwrap();
        assert_eq!(initial.len(), 2);
        assert!(
            context
                .sync_chatend(&mut journal, "t2", &BTreeMap::new(), &[],)
                .unwrap()
                .is_empty()
        );

        context
            .apply_load(node(1, &[2], &[]), vec![node(2, &[], &[])])
            .unwrap();
        let changed = context
            .sync_chatend(&mut journal, "t3", &BTreeMap::new(), &[])
            .unwrap();
        assert_eq!(changed.len(), 2);
        assert_eq!(
            changed
                .iter()
                .map(|box_id| journal.state().box_state(*box_id).unwrap().name.as_str())
                .collect::<Vec<_>>(),
            vec!["Kweb loaded node", "Kweb fixed connection"]
        );

        drop(journal);
        std::fs::remove_dir_all(root).unwrap();
    }
}