wryme 1.7.4

// Application state. The UI is a pure function of this.

use std::sync::{Arc, Mutex};

use crate::api::{ApiMessage, ApiToolCall};
use crate::book;
use crate::popup::Popup;
use crate::shop::Shop;
use crate::station::Station;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Role {
    User,
    Assistant,
}

/// What kind of chunk the model most recently sent us during a stream.
/// Drives the dim header indicator next to the role label: "thinking…"
/// vs "writing…" vs "tinkering…". Only meaningful while `streaming`.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Phase {
    Streaming,
    Thinking,
    Tinkering,
    Writing,
}

/// How the message area handles overflow.
///
/// Page is the default. Content is shown in discrete viewport-sized chunks
/// and navigation snaps between them. Scroll is the alternative: a smooth
/// row-by-row offset, more like a traditional terminal pager.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ViewMode {
    Page,
    Scroll,
}

/// One tool call the model made within an assistant turn, together with
/// the result we fed back to it. Persisted on the Message so the wire
/// transcript between user turns carries the full tool-call/result pair —
/// without it, small models "forget" they have tools and stop using them.
#[derive(Debug, Clone)]
pub struct ToolEvent {
    pub call_id: String,
    pub name: String,
    pub arguments: String,
    pub result: String,
}

#[derive(Debug)]
pub struct Message {
    pub role: Role,
    pub content: String,
    /// Image file paths attached to this user message. Sent upstream as
    /// image inputs; rendered as a small attachment note.
    pub images: Vec<String>,
    /// Reasoning / chain-of-thought from the model, if it sent any.
    /// Streamed before `content`, rendered below the reply (older in time)
    /// as a dimmer "brain" block.
    pub brain: String,
    /// True while the model is still streaming into this message.
    pub streaming: bool,
    /// Local time the message was created, formatted "HH:mm". Set once
    /// when the message is pushed and never updated. We don't persist
    /// across sessions, so the date is always today and not stored.
    pub timestamp: String,
    /// What the model is currently doing (writing, thinking, calling a
    /// tool). Drives the dim header label. Stops being displayed once
    /// streaming ends.
    pub phase: Phase,
    /// Name of the tool the model is currently calling, if any. Surfaced
    /// to the right of the phase indicator on the header line, just left
    /// of the timestamp. Only displayed while streaming.
    pub current_tool: Option<String>,
    /// Logical turn this message belongs to. A single assistant reply can
    /// span several messages (one per contiguous non-thinking cluster) —
    /// they all share the same `turn_id` so the wire transcript and the
    /// book record still treat them as one turn. User messages set 0.
    pub turn_id: u64,
    /// Tool calls made during this assistant turn and their results, in
    /// order. Replayed into the wire history on the next request so the
    /// model sees its own tool use and keeps using tools.
    pub tool_events: Vec<ToolEvent>,
}

fn now_hhmm() -> String {
    chrono::Local::now().format("%H:%M").to_string()
}

pub struct App {
    /// Messages in chronological order (index 0 = oldest).
    /// The UI renders them in reverse so the newest sits at the top.
    pub messages: Vec<Message>,
    pub system: Option<String>,
    /// True while a request is in flight.
    pub in_flight: bool,
    pub status: String,
    pub should_quit: bool,
    /// Which page of the message stack the user is currently viewing.
    /// 0 = the live page (newest content at top). Higher = further back
    /// in the conversation. The renderer clamps this to the number of
    /// pages actually available given the current viewport. Only meaningful
    /// in `ViewMode::Page`.
    pub current_page: usize,
    /// Row-level scroll offset, used only in `ViewMode::Scroll`. 0 = newest
    /// content visible at the top, higher = scrolled into older content.
    pub scroll_row: usize,
    /// Mouse wheel accumulator. Ticks add up here; once the magnitude
    /// crosses a threshold we move a page and subtract the threshold.
    /// Keeps trackpad scrolling from blowing through pages instantly.
    pub wheel_accum: i32,
    /// Current view mode. Page is the default; Scroll is the alternative.
    pub view_mode: ViewMode,
    /// Last rendered viewport height (rows). Recorded by the renderer so
    /// the key handlers can step by viewport when the user hits PgUp/PgDn
    /// in scroll mode.
    pub last_viewport_h: usize,
    /// Most recent response.id seen from a Responses-protocol station.
    /// Captured for the UI; never replayed — requests are stateless
    /// (`store: false`, full transcript every turn). Reset to None on
    /// launch; we don't persist across runs.
    pub last_response_id: Option<String>,
    /// Running token usage for this window. `usage_ctx` is the latest
    /// reported prompt size (replaced each turn — every request re-reports
    /// the whole transcript, so summing it would explode past the real
    /// context). `usage_out` accumulates generated tokens across turns.
    /// Rendered as a gray K-count, bottom-right.
    pub usage_ctx: u64,
    pub usage_out: u64,
    pub voice_on: bool,
    pub voice_muted: bool,
    pub voice_speaker: Option<crate::voice::Speaker>,
    pub voice_buffer: String,
    /// All shops loaded at startup. Read-only after that. Used by the
    /// popup to list every model any shop can run.
    pub shops: Vec<Shop>,
    /// Every saved station loaded at startup, plus any the user saves
    /// during this session via the popup. Mutable.
    pub stations: Vec<Station>,
    /// The station currently in effect. The popup mutates this when the
    /// user adjusts dials, picks a different model, or loads a saved
    /// station. Cloned into each in-flight request.
    pub active_station: Station,
    /// The shop currently in effect. Re-resolved every time
    /// active_station.model changes.
    pub active_shop: Shop,
    /// Name of the saved station this session traces back to, if any.
    /// None when the session was synthesized (untitled, demo). Used
    /// alongside `active_station` to compute whether the active config
    /// is "dirty" (modified vs. its saved form).
    pub active_origin: Option<String>,
    /// The popup overlay state (closed, browsing, or entering a name).
    pub popup: Popup,
    /// The book engine: wryme's memory, a columnar Parquet store of
    /// open-ended compartments. Shared across turns as an Arc<Mutex<..>>
    /// so the streaming protocol and the background delivery can both
    /// reach it (the invisible `book` tool locks it in the flow).
    pub engine: Arc<Mutex<book::Engine>>,
    /// Monotonic counter for logical turns; bumped by begin_assistant and
    /// stamped onto every cluster of that turn.
    turn_counter: u64,
    /// True if the last thing streamed into the assistant was thinking
    /// (a Brain delta). A content delta arriving right after this starts
    /// a fresh cluster so each contiguous non-thinking run is its own entry.
    last_stream_was_brain: bool,
}

/// The book lives at `~/.config/wryme/book`, like the shops and
/// stations files — one folder for everything wryme keeps.
fn book_dir() -> std::path::PathBuf {
    let home = std::env::var("HOME").unwrap_or_else(|_| ".".to_string());
    std::path::PathBuf::from(home)
        .join(".config")
        .join("wryme")
        .join("book")
}

impl App {
    pub fn new(
        system: Option<String>,
        shops: Vec<Shop>,
        stations: Vec<Station>,
        active_station: Station,
        active_shop: Shop,
        active_origin: Option<String>,
    ) -> Self {
        Self {
            messages: Vec::new(),
            system,
            in_flight: false,
            status: String::new(),
            should_quit: false,
            current_page: 0,
            scroll_row: 0,
            wheel_accum: 0,
            view_mode: ViewMode::Page,
            last_viewport_h: 0,
            last_response_id: None,
            usage_ctx: 0,
            usage_out: 0,
            voice_on: false,
            voice_muted: false,
            voice_speaker: None,
            voice_buffer: String::new(),
            shops,
            stations,
            active_station,
            active_shop,
            active_origin,
            popup: Popup::default(),
            engine: Arc::new(Mutex::new(
                book::open_engine(&book_dir()).expect("open book"),
            )),
            turn_counter: 0,
            last_stream_was_brain: false,
        }
    }

    /// True when the active station differs from the saved entry it was
    /// loaded from. False when there is no origin (untitled / demo) or
    /// when active matches its saved entry exactly.
    pub fn is_dirty(&self) -> bool {
        let Some(origin) = &self.active_origin else {
            return false;
        };
        let Some(saved) = self.stations.iter().find(|s| s.name == *origin) else {
            // Origin set but saved entry missing. Should not happen in
            // normal use; treat as dirty so the user notices.
            return true;
        };
        saved.model != self.active_station.model
            || saved.dials.boldness != self.active_station.dials.boldness
            || saved.dials.patience != self.active_station.dials.patience
            || saved.dials.verbosity != self.active_station.dials.verbosity
            || saved.dials.tinker_keep != self.active_station.dials.tinker_keep
            || saved.dials.tinker_clip != self.active_station.dials.tinker_clip
            || saved.voice != self.active_station.voice
    }

    pub fn note(&mut self, msg: impl Into<String>) {
        self.status = msg.into();
    }

    pub fn stop_voice(&mut self) {
        if let Some(speaker) = self.voice_speaker.as_mut() {
            speaker.stop();
        }
        self.voice_buffer.clear();
    }

    /// Quiet for the rest of this turn: kills current speech and drops
    /// anything queued, and new deltas stay silent until the next turn.
    pub fn mute_voice(&mut self) {
        self.stop_voice();
        self.voice_muted = true;
    }

    pub fn unmute_voice(&mut self) {
        self.voice_muted = false;
    }

    pub fn voice_is_active(&self) -> bool {
        self.voice_speaker
            .as_ref()
            .map(|s| s.is_active())
            .unwrap_or(false)
    }

    pub fn shutdown_voice(&mut self) {
        if let Some(mut speaker) = self.voice_speaker.take() {
            speaker.shutdown();
        }
        self.voice_buffer.clear();
    }

    pub fn ensure_speaker(&mut self, voice: Option<String>) {
        let same = self
            .voice_speaker
            .as_ref()
            .map(|s| s.name == voice)
            .unwrap_or(false);
        if !same {
            self.shutdown_voice();
            self.voice_speaker = Some(crate::voice::Speaker::new(voice));
        }
    }

    pub fn push_user(&mut self, content: String, images: Vec<String>) {
        self.messages.push(Message {
            role: Role::User,
            content,
            images,
            brain: String::new(),
            streaming: false,
            timestamp: now_hhmm(),
            phase: Phase::Streaming,
            current_tool: None,
            turn_id: 0,
            tool_events: Vec::new(),
        });
        // The engine records every turn into the continuous stream, and
        // re-checks whether an unattributed thread is weightful enough to
        // prod the agent into deeming it.
        if let Some(last) = self.messages.last() {
            if last.role == Role::User {
                if let Ok(mut e) = self.engine.lock() {
                    e.record_turn("user", &last.content);
                }
            }
        }
    }

    pub fn begin_assistant(&mut self) {
        self.turn_counter += 1;
        let tid = self.turn_counter;
        self.messages.push(Message {
            role: Role::Assistant,
            content: String::new(),
            images: Vec::new(),
            brain: String::new(),
            streaming: true,
            timestamp: now_hhmm(),
            phase: Phase::Streaming,
            current_tool: None,
            turn_id: tid,
            tool_events: Vec::new(),
        });
    }

    pub fn append_to_last_assistant(&mut self, delta: &str) {
        if delta.is_empty() {
            return;
        }
        let idx = self
            .messages
            .iter()
            .rposition(|m| m.role == Role::Assistant && m.streaming);
        let Some(idx) = idx else { return };

        // A writing delta right after thinking starts a fresh cluster, so
        // each contiguous non-thinking run shows as its own assistant entry.
        // (An empty placeholder that already holds the leading brain just
        // absorbs the first content — that is still the first cluster.)
        if self.last_stream_was_brain && !self.messages[idx].content.is_empty() {
            let tid = self.messages[idx].turn_id;
            // Close the previous cluster: only the LATEST cluster of a turn
            // stays `streaming`, so the phase label, tool name and cursor
            // show on the newest entry only. Older clusters render as
            // finished text without live indicators.
            self.messages[idx].streaming = false;
            self.messages.push(Message {
                role: Role::Assistant,
                content: String::new(),
                images: Vec::new(),
                brain: String::new(),
                streaming: true,
                timestamp: now_hhmm(),
                phase: Phase::Streaming,
                current_tool: None,
                turn_id: tid,
                tool_events: Vec::new(),
            });
        }
        self.last_stream_was_brain = false;

        let m = match self.messages.last_mut() {
            Some(m) => m,
            None => return,
        };
        // Verbatim concatenation (issue #16): stream deltas are appended
        // exactly as they arrive. An earlier heuristic inserted a space
        // whenever both sides were non-whitespace (to fix "sentence.Next"
        // splits), but that corrupts well-formed streams — "Hello" + ","
        // became "Hello ,", subword splits became "un der". Both the Chat
        // and Responses specs define deltas as exact slices; any spacing
        // the model intends already rides inside them.
        m.content.push_str(delta);
        m.phase = Phase::Writing;
    }

    pub fn append_to_last_brain(&mut self, delta: &str) {
        // All reasoning of a turn lands on ONE message — the turn's first
        // cluster — so the brain shows as a single UI entry per turn rather
        // than a footnote on every cluster.
        let tid = self.turn_counter;
        if let Some(m) = self
            .messages
            .iter_mut()
            .find(|m| m.role == Role::Assistant && m.turn_id == tid)
        {
            m.brain.push_str(delta);
        }
        // The "thinking…" indicator belongs on the latest visual entry only.
        if let Some(m) = self
            .messages
            .iter_mut()
            .rev()
            .find(|m| m.role == Role::Assistant && m.streaming)
        {
            m.phase = Phase::Thinking;
        }
        self.last_stream_was_brain = true;
    }

    pub fn record_tool_call(&mut self, name: Option<String>) {
        if let Some(m) = self
            .messages
            .iter_mut()
            .rev()
            .find(|m| m.role == Role::Assistant && m.streaming)
        {
            m.phase = Phase::Tinkering;
            if let Some(n) = name {
                m.current_tool = Some(n);
            }
        }
    }

    /// Persist a tool call/result pair onto the streaming assistant message,
    /// so the next user turn's wire history carries the full tool transcript.
    pub fn record_tool_result(
        &mut self,
        call_id: String,
        name: String,
        arguments: String,
        result: String,
    ) {
        if let Some(m) = self
            .messages
            .iter_mut()
            .rev()
            .find(|m| m.role == Role::Assistant && m.streaming)
        {
            m.tool_events.push(ToolEvent {
                call_id,
                name,
                arguments,
                result,
            });
        }
    }

    pub fn finish_streaming(&mut self) {
        self.in_flight = false;

        // Find the streaming assistant message (most recent) and its turn_id.
        let mut just_finished: Option<usize> = None;
        for i in (0..self.messages.len()).rev() {
            if self.messages[i].streaming {
                just_finished = Some(i);
                break;
            }
        }

        if let Some(i) = just_finished {
            let tid = self.messages[i].turn_id;

            // Mark every cluster of this turn done.
            for m in self.messages.iter_mut() {
                if m.streaming && m.turn_id == tid {
                    m.streaming = false;
                }
            }

            // Record the whole logical turn (all clusters) into the book once.
            let joined: String = self
                .messages
                .iter()
                .filter(|m| m.role == Role::Assistant && m.turn_id == tid && !m.content.is_empty())
                .map(|m| m.content.as_str())
                .collect::<Vec<_>>()
                .join("\n");
            if !joined.is_empty() {
                if let Ok(mut e) = self.engine.lock() {
                    e.record_turn("assistant", &joined);
                }
            }

            // Drop empty clusters of this turn so the screen does not show a
            // confusing empty bubble. Server hiccups and pre-delta errors are
            // common causes. If upstream sent an error, the status bar already
            // explains what happened. If not, leave a short note.
            let any_nonempty = self.messages.iter().any(|m| {
                m.role == Role::Assistant
                    && m.turn_id == tid
                    && (!m.content.is_empty() || !m.brain.is_empty() || m.current_tool.is_some())
            });
            if !any_nonempty {
                self.messages
                    .retain(|m| !(m.role == Role::Assistant && m.turn_id == tid));
                if self.status.is_empty() {
                    self.note("empty reply");
                }
            }
        }
    }

    /// Build the wire-format message list to send upstream. The base
    /// system prompt, then the established compartments' rendered
    /// bookmarks as one system message each — the preamble — then the
    /// live turns.
    pub fn api_messages(&self) -> Vec<ApiMessage> {
        let mut out = Vec::with_capacity(self.messages.len() + 1);
        if let Some(sys) = &self.system {
            out.push(ApiMessage {
                role: "system".into(),
                content: sys.clone(),
                images: Vec::new(),
                tool_calls: Vec::new(),
                tool_call_id: String::new(),
                tool_result: String::new(),
            });
        }
        if let Ok(mut engine) = self.engine.lock() {
            for preamble in engine.preamble() {
                out.push(ApiMessage {
                    role: "system".into(),
                    content: preamble,
                    images: Vec::new(),
                    tool_calls: Vec::new(),
                    tool_call_id: String::new(),
                    tool_result: String::new(),
                });
            }
            // The book-writing prod: a quiet system reminder the engine
            // slips the agent once when an unattributed thread is weightful.
            if let Some(prod) = engine.take_prod() {
                out.push(ApiMessage {
                    role: "system".into(),
                    content: prod,
                    images: Vec::new(),
                    tool_calls: Vec::new(),
                    tool_call_id: String::new(),
                    tool_result: String::new(),
                });
            }
        }
        let mut last_asst_turn: Option<u64> = None;
        for m in &self.messages {
            // Skip an empty streaming placeholder. We send the history
            // BEFORE the assistant turn we're about to fill.
            if m.streaming && m.content.is_empty() {
                continue;
            }
            // A cluster of an assistant turn still in flight: skip it too —
            // we only send history up to the turn we're about to fill.
            if m.role == Role::Assistant && m.streaming {
                continue;
            }
            // Consecutive assistant clusters of the SAME logical turn merge
            // into one ApiMessage, so the wire sees one assistant turn.
            if m.role == Role::Assistant
                && last_asst_turn == Some(m.turn_id)
                && out.last().map(|a| a.role.as_str()) == Some("assistant")
            {
                let mut merged = false;
                if let Some(last) = out.last_mut() {
                    if !last.content.is_empty() && !m.content.is_empty() {
                        last.content.push('\n');
                    }
                    last.content.push_str(&m.content);
                    for ev in &m.tool_events {
                        last.tool_calls.push(ApiToolCall {
                            id: ev.call_id.clone(),
                            name: ev.name.clone(),
                            arguments: ev.arguments.clone(),
                        });
                    }
                    merged = true;
                }
                // Emit the tool-role results after releasing the borrow.
                for ev in &m.tool_events {
                    out.push(ApiMessage {
                        role: "tool".into(),
                        content: ev.result.clone(),
                        images: Vec::new(),
                        tool_calls: Vec::new(),
                        tool_call_id: ev.call_id.clone(),
                        tool_result: ev.result.clone(),
                    });
                }
                if !merged {
                    // Should not happen (we checked out.last() is assistant),
                    // but keep a fresh assistant message so history is sane.
                    out.push(ApiMessage {
                        role: "assistant".into(),
                        content: m.content.clone(),
                        images: Vec::new(),
                        tool_calls: m
                            .tool_events
                            .iter()
                            .map(|ev| ApiToolCall {
                                id: ev.call_id.clone(),
                                name: ev.name.clone(),
                                arguments: ev.arguments.clone(),
                            })
                            .collect(),
                        tool_call_id: String::new(),
                        tool_result: String::new(),
                    });
                    for ev in &m.tool_events {
                        out.push(ApiMessage {
                            role: "tool".into(),
                            content: ev.result.clone(),
                            images: Vec::new(),
                            tool_calls: Vec::new(),
                            tool_call_id: ev.call_id.clone(),
                            tool_result: ev.result.clone(),
                        });
                    }
                }
                last_asst_turn = Some(m.turn_id);
                continue;
            }
            if m.role == Role::Assistant {
                last_asst_turn = Some(m.turn_id);
            }
            let (tool_calls, mut results) = if m.role == Role::Assistant {
                let mut tcs = Vec::new();
                let mut res = Vec::new();
                for ev in &m.tool_events {
                    tcs.push(ApiToolCall {
                        id: ev.call_id.clone(),
                        name: ev.name.clone(),
                        arguments: ev.arguments.clone(),
                    });
                    res.push(ApiMessage {
                        role: "tool".into(),
                        content: ev.result.clone(),
                        images: Vec::new(),
                        tool_calls: Vec::new(),
                        tool_call_id: ev.call_id.clone(),
                        tool_result: ev.result.clone(),
                    });
                }
                (tcs, res)
            } else {
                (Vec::new(), Vec::new())
            };
            out.push(ApiMessage {
                role: match m.role {
                    Role::User => "user",
                    Role::Assistant => "assistant",
                }
                .into(),
                content: m.content.clone(),
                images: if m.role == Role::User {
                    m.images.clone()
                } else {
                    Vec::new()
                },
                tool_calls,
                tool_call_id: String::new(),
                tool_result: String::new(),
            });
            // Emit the tool-role result messages right after their assistant
            // turn, so the wire sees the call/result pair together.
            out.append(&mut results);
        }
        // Tinker keep/clip: prune replay, keep pair atomic. Defaults keep=all/full.
        let clip = self.active_station.dials.tinker_clip;
        let total_pairs = out.iter().filter(|m| m.role == "tool").count();
        let keep_n = self.active_station.dials.tinker_keep.keep_n(total_pairs);
        if keep_n.is_some() || clip != crate::station::TinkerVal::All {
            // Collect pair indices: each assistant with tool_calls + following tool msgs.
            // We prune oldest pairs to keep last N.
            if let Some(n) = keep_n {
                if total_pairs > n {
                    let drop = total_pairs - n;
                    let mut to_drop = drop;
                    let mut pruned: Vec<ApiMessage> = Vec::with_capacity(out.len());
                    for msg in out {
                        if msg.role == "tool" && to_drop > 0 {
                            to_drop -= 1;
                            // also drop its paired call from preceding assistant
                            if let Some(last) = pruned.last_mut() {
                                if last.role == "assistant" && !last.tool_calls.is_empty() {
                                    // find matching call_id
                                    let id = &msg.tool_call_id;
                                    last.tool_calls.retain(|c| c.id != *id);
                                }
                            }
                            continue;
                        }
                        if msg.role == "assistant" && to_drop > 0 && msg.tool_calls.len() <= to_drop
                        {
                            // This assistant's calls are among dropped oldest - handled via tool drop above,
                            // but if assistant has no remaining calls keep content anyway
                            // (don't drop whole assistant turn)
                        }
                        pruned.push(msg);
                    }
                    out = pruned;
                } else {
                    // no keep-pruning needed, keep out as is
                }
            }
            if clip != crate::station::TinkerVal::All {
                for m in &mut out {
                    if m.role == "tool" {
                        let clipped = clip.clip(&m.content);
                        m.content = clipped.clone();
                        m.tool_result = clipped;
                    }
                }
            }
        }
        out
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn test_app() -> App {
        let mut app = App::new(
            None,
            vec![Shop::demo()],
            vec![Station::demo()],
            Station::demo(),
            Shop::demo(),
            None,
        );
        // Keep the test off the real book: swap in a temp engine.
        let dir = std::env::temp_dir().join(format!("wryme_app_{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&dir);
        app.engine = Arc::new(Mutex::new(book::open_engine(&dir).unwrap()));
        app
    }

    #[test]
    fn stream_deltas_concatenate_verbatim() {
        // Issue #16: no space insertion. Punctuation and subword splits
        // must land exactly as streamed.
        let mut app = test_app();
        app.begin_assistant();
        for d in ["Hello", ",", " world", "!", " un", "der"] {
            app.append_to_last_assistant(d);
        }
        let joined: String = app
            .messages
            .iter()
            .map(|m| m.content.as_str())
            .collect::<Vec<_>>()
            .join("");
        assert_eq!(joined, "Hello, world! under");
        let _ = std::fs::remove_dir_all(
            std::env::temp_dir().join(format!("wryme_app_{}", std::process::id())),
        );
    }
}