agent-context 0.1.0

//! [`AgentContext`] Actor 及所有 [`Message`] 实现。
//!
//! 管理三区消息模型的状态，提供消息增删改查、模型对话、上下文压缩等操作。

use std::sync::Arc;

use kameo::prelude::*;

use super::event::{ChangeEvent, CompressStrategy};
use super::stream::AgentSendStream;
use super::types::{ContextBackend, ScratchOpts};
use crate::error::AgentError;
use crate::message::ContextMessage;
use crate::readonly::ReadOnly;
use crate::role::Role;

// ---------------------------------------------------------------------------
// AgentContext Actor
// ---------------------------------------------------------------------------

/// LLM 对话上下文管理器，kameo Actor。
///
/// 管理三区 + Scratch 消息模型（immutable → compressed → incremental → scratch），提供：
/// - 消息增删改查（[`AppendMsg`]、[`UpdateMsg`]、[`RemoveMsg`] 等）
/// - 对话发送（[`SendMsg`]、[`SendStreamMsg`]），支持通过 [`ScratchOpts`] 追加临时元数据
/// - 上下文压缩（[`CompressMsg`]）
/// - Token 估算和溢出检测（[`EstimateTokensMsg`]、[`IsFullMsg`]）
/// - 变更回调（[`ChangeEvent`]）
///
/// ## 构造
///
/// ```ignore
/// let ctx = AgentContext::new(backend, vec![])
///     .with_on_change(|event| { /* 处理变更 */ });
/// let actor = AgentContext::spawn(ctx);
/// ```
#[derive(Actor)]
pub struct AgentContext<B: ContextBackend> {
    backend: B,
    immutable: ReadOnly<B::Message>,
    compressed: Vec<B::Message>,
    incremental: Vec<B::Message>,
    #[expect(clippy::type_complexity, reason = "回调类型不可避免复杂")]
    on_change: Option<Arc<dyn Fn(ChangeEvent<B::Message>) + Send + Sync>>,
    #[expect(clippy::type_complexity, reason = "回调类型不可避免复杂")]
    on_compressed: Option<
        Arc<
            dyn Fn(Vec<B::Message>, Vec<B::Message>) -> (Vec<B::Message>, Vec<B::Message>)
                + Send
                + Sync,
        >,
    >,
}

impl<B: ContextBackend> AgentContext<B> {
    /// 创建新的上下文管理器。
    ///
    /// - `backend`: 实现了 [`ContextBackend`] 的 LLM 后端实例
    /// - `immutable`: 初始不可变消息（系统提示词等），放入 immutable 区
    pub fn new(backend: B, immutable: Vec<B::Message>) -> Self {
        Self {
            backend,
            immutable: ReadOnly::from(immutable),
            compressed: Vec::new(),
            incremental: Vec::new(),
            on_change: None,
            on_compressed: None,
        }
    }

    /// 注册增量区变更回调。
    ///
    /// 每次对 incremental 区的写操作（追加/更新/插入/移除/清空等）都会触发此回调。
    /// 用于 CLI 实时展示、日志记录等场景。
    pub fn with_on_change(
        mut self,
        f: impl Fn(ChangeEvent<B::Message>) + Send + Sync + 'static,
    ) -> Self {
        self.on_change = Some(Arc::new(f));
        self
    }

    /// 注册压缩后处理回调。
    ///
    /// 在 [`CompressMsg`] 生成摘要后、写入 compressed 区之前调用。
    /// 回调接收 `(摘要消息列表, 保留消息列表)`，返回 `(最终摘要, 最终保留)`。
    /// 用于自定义后处理（如过滤、重新排序摘要内容）。
    pub fn with_on_compressed(
        mut self,
        f: impl Fn(Vec<B::Message>, Vec<B::Message>) -> (Vec<B::Message>, Vec<B::Message>)
        + Send
        + Sync
        + 'static,
    ) -> Self {
        self.on_compressed = Some(Arc::new(f));
        self
    }

    fn default_summary_prompt() -> String {
        "请将以下对话历史压缩为简洁摘要，保留关键信息、决策和上下文。输出一条 system 消息。"
            .to_string()
    }
}

// ---------------------------------------------------------------------------
// Actor Messages
// ---------------------------------------------------------------------------

/// 追加一条消息到 incremental 区末尾。Reply = `()`。
///
/// 触发 [`ChangeEvent::Appended`] 回调。
pub struct AppendMsg<M> {
    /// 要追加的消息
    pub message: M,
}

impl<B: ContextBackend> Message<AppendMsg<B::Message>> for AgentContext<B> {
    type Reply = ();

    async fn handle(
        &mut self,
        msg: AppendMsg<B::Message>,
        _ctx: &mut Context<Self, Self::Reply>,
    ) -> Self::Reply {
        self.incremental.push(msg.message);
        if let Some(cb) = &self.on_change
            && let Some(last) = self.incremental.last().cloned()
        {
            cb(ChangeEvent::Appended(last));
        }
    }
}

/// 获取三区总消息数。Reply = `usize`。
pub struct Len;

impl<B: ContextBackend> Message<Len> for AgentContext<B> {
    type Reply = usize;

    async fn handle(&mut self, _msg: Len, _ctx: &mut Context<Self, Self::Reply>) -> Self::Reply {
        self.immutable.len() + self.compressed.len() + self.incremental.len()
    }
}

/// 检查三区是否全部为空。Reply = `bool`。
pub struct IsEmpty;

impl<B: ContextBackend> Message<IsEmpty> for AgentContext<B> {
    type Reply = bool;

    async fn handle(
        &mut self,
        _msg: IsEmpty,
        _ctx: &mut Context<Self, Self::Reply>,
    ) -> Self::Reply {
        self.immutable.is_empty() && self.compressed.is_empty() && self.incremental.is_empty()
    }
}

/// 批量追加消息到 incremental 区。Reply = `()`。
///
/// 每条消息单独触发 [`ChangeEvent::Appended`]。
pub struct ExtendMsg<M> {
    /// 要批量追加的消息列表
    pub messages: Vec<M>,
}

impl<B: ContextBackend> Message<ExtendMsg<B::Message>> for AgentContext<B> {
    type Reply = ();

    async fn handle(
        &mut self,
        msg: ExtendMsg<B::Message>,
        _ctx: &mut Context<Self, Self::Reply>,
    ) -> Self::Reply {
        for m in msg.messages {
            self.incremental.push(m);
            if let Some(cb) = &self.on_change
                && let Some(last) = self.incremental.last().cloned()
            {
                cb(ChangeEvent::Appended(last));
            }
        }
    }
}

/// 静默追加消息，不触发 [`ChangeEvent`] 回调。Reply = `()`。
///
/// 用于 [`AgentSendStream`] Drop 时自动存储消息，避免二次通知。
pub struct SilentAppendMsg<M> {
    /// 要静默追加的消息
    pub message: M,
}

impl<B: ContextBackend> Message<SilentAppendMsg<B::Message>> for AgentContext<B> {
    type Reply = ();

    async fn handle(
        &mut self,
        msg: SilentAppendMsg<B::Message>,
        _ctx: &mut Context<Self, Self::Reply>,
    ) -> Self::Reply {
        self.incremental.push(msg.message);
    }
}

/// 按全局索引获取消息。Reply = `Option<Message>`。
///
/// 索引按 immutable → compressed → incremental 顺序计算。
/// 越界返回 `None`。
pub struct Get(pub usize);

impl<B: ContextBackend> Message<Get> for AgentContext<B> {
    type Reply = Option<B::Message>;

    async fn handle(&mut self, msg: Get, _ctx: &mut Context<Self, Self::Reply>) -> Self::Reply {
        let idx = msg.0;
        let imm_len = self.immutable.len();
        let comp_len = self.compressed.len();
        if idx < imm_len {
            Some(self.immutable[idx].clone())
        } else if idx < imm_len + comp_len {
            Some(self.compressed[idx - imm_len].clone())
        } else if idx < imm_len + comp_len + self.incremental.len() {
            Some(self.incremental[idx - imm_len - comp_len].clone())
        } else {
            None
        }
    }
}

/// 获取三区全部消息的拼接结果。Reply = `Vec<Message>`。
///
/// 顺序：immutable → compressed → incremental。
pub struct MessagesMsg;

impl<B: ContextBackend> Message<MessagesMsg> for AgentContext<B> {
    type Reply = Vec<B::Message>;

    async fn handle(
        &mut self,
        _msg: MessagesMsg,
        _ctx: &mut Context<Self, Self::Reply>,
    ) -> Self::Reply {
        self.immutable
            .iter()
            .chain(self.compressed.iter())
            .chain(self.incremental.iter())
            .cloned()
            .collect()
    }
}

/// 获取 immutable 区的消息副本。Reply = `Vec<Message>`。
pub struct ImmutableMsg;

impl<B: ContextBackend> Message<ImmutableMsg> for AgentContext<B> {
    type Reply = Vec<B::Message>;

    async fn handle(
        &mut self,
        _msg: ImmutableMsg,
        _ctx: &mut Context<Self, Self::Reply>,
    ) -> Self::Reply {
        self.immutable.to_vec()
    }
}

/// 获取 compressed 区的消息副本。Reply = `Vec<Message>`。
pub struct CompressedMsg;

impl<B: ContextBackend> Message<CompressedMsg> for AgentContext<B> {
    type Reply = Vec<B::Message>;

    async fn handle(
        &mut self,
        _msg: CompressedMsg,
        _ctx: &mut Context<Self, Self::Reply>,
    ) -> Self::Reply {
        self.compressed.clone()
    }
}

/// 获取 incremental 区的消息副本。Reply = `Vec<Message>`。
pub struct IncrementalMsg;

impl<B: ContextBackend> Message<IncrementalMsg> for AgentContext<B> {
    type Reply = Vec<B::Message>;

    async fn handle(
        &mut self,
        _msg: IncrementalMsg,
        _ctx: &mut Context<Self, Self::Reply>,
    ) -> Self::Reply {
        self.incremental.clone()
    }
}

/// 按角色筛选三区全部消息。Reply = `Vec<Message>`。
pub struct FindByRoleMsg(pub Role);

impl<B: ContextBackend> Message<FindByRoleMsg> for AgentContext<B> {
    type Reply = Vec<B::Message>;

    async fn handle(
        &mut self,
        msg: FindByRoleMsg,
        _ctx: &mut Context<Self, Self::Reply>,
    ) -> Self::Reply {
        self.immutable
            .iter()
            .chain(self.compressed.iter())
            .chain(self.incremental.iter())
            .filter(|m| m.role() == msg.0)
            .cloned()
            .collect()
    }
}

/// 替换 incremental 区指定索引的消息。Reply = `Result<(), AgentError>`。
///
/// 仅对 incremental 区有效。触发 [`ChangeEvent::Updated`]。
pub struct UpdateMsg<M> {
    /// incremental 区索引
    pub index: usize,
    /// 新消息
    pub message: M,
}

impl<B: ContextBackend> Message<UpdateMsg<B::Message>> for AgentContext<B> {
    type Reply = Result<(), AgentError>;

    async fn handle(
        &mut self,
        msg: UpdateMsg<B::Message>,
        _ctx: &mut Context<Self, Self::Reply>,
    ) -> Self::Reply {
        if msg.index >= self.incremental.len() {
            return Err(AgentError::Context("索引越界".into()));
        }
        let old = std::mem::replace(&mut self.incremental[msg.index], msg.message);
        if let Some(cb) = &self.on_change {
            cb(ChangeEvent::Updated {
                index: msg.index,
                old,
                new: self.incremental[msg.index].clone(),
            });
        }
        Ok(())
    }
}

/// 在 incremental 区指定索引插入消息。Reply = `Result<(), AgentError>`。
///
/// 触发 [`ChangeEvent::Inserted`]。
pub struct InsertMsg<M> {
    /// incremental 区插入位置
    pub index: usize,
    /// 要插入的消息
    pub message: M,
}

impl<B: ContextBackend> Message<InsertMsg<B::Message>> for AgentContext<B> {
    type Reply = Result<(), AgentError>;

    async fn handle(
        &mut self,
        msg: InsertMsg<B::Message>,
        _ctx: &mut Context<Self, Self::Reply>,
    ) -> Self::Reply {
        if msg.index > self.incremental.len() {
            return Err(AgentError::Context("索引越界".into()));
        }
        self.incremental.insert(msg.index, msg.message);
        if let Some(cb) = &self.on_change {
            cb(ChangeEvent::Inserted {
                index: msg.index,
                message: self.incremental[msg.index].clone(),
            });
        }
        Ok(())
    }
}

/// 移除 incremental 区指定索引的消息。Reply = `Result<(), AgentError>`。
///
/// 触发 [`ChangeEvent::Removed`]。
pub struct RemoveMsg {
    /// incremental 区索引
    pub index: usize,
}

impl<B: ContextBackend> Message<RemoveMsg> for AgentContext<B> {
    type Reply = Result<(), AgentError>;

    async fn handle(
        &mut self,
        msg: RemoveMsg,
        _ctx: &mut Context<Self, Self::Reply>,
    ) -> Self::Reply {
        if msg.index >= self.incremental.len() {
            return Err(AgentError::Context("索引越界".into()));
        }
        let removed = self.incremental.remove(msg.index);
        if let Some(cb) = &self.on_change {
            cb(ChangeEvent::Removed {
                index: msg.index,
                message: removed,
            });
        }
        Ok(())
    }
}

/// 弹出 incremental 区最后一条消息。Reply = `Option<Message>`。
///
/// 触发 [`ChangeEvent::Popped`]。
pub struct PopMsg;

impl<B: ContextBackend> Message<PopMsg> for AgentContext<B> {
    type Reply = Option<B::Message>;

    async fn handle(&mut self, _msg: PopMsg, _ctx: &mut Context<Self, Self::Reply>) -> Self::Reply {
        let popped = self.incremental.pop();
        if let Some(ref msg) = popped
            && let Some(cb) = &self.on_change
        {
            cb(ChangeEvent::Popped(msg.clone()));
        }
        popped
    }
}

/// 按角色保留 incremental 区消息，其余移除。Reply = `()`。
///
/// 触发 [`ChangeEvent::Retained`]。
pub struct RetainMsg {
    /// 要保留的角色
    pub role: Role,
}

impl<B: ContextBackend> Message<RetainMsg> for AgentContext<B> {
    type Reply = ();

    async fn handle(
        &mut self,
        msg: RetainMsg,
        _ctx: &mut Context<Self, Self::Reply>,
    ) -> Self::Reply {
        let mut removed = Vec::new();
        let role = msg.role;
        self.incremental.retain(|m| {
            if m.role() == role {
                true
            } else {
                removed.push(m.clone());
                false
            }
        });
        if let Some(cb) = &self.on_change {
            cb(ChangeEvent::Retained { role, removed });
        }
    }
}

/// 清空整个 incremental 区。Reply = `()`。
///
/// 触发 [`ChangeEvent::Cleared`]。
pub struct ClearMsg;

impl<B: ContextBackend> Message<ClearMsg> for AgentContext<B> {
    type Reply = ();

    async fn handle(
        &mut self,
        _msg: ClearMsg,
        _ctx: &mut Context<Self, Self::Reply>,
    ) -> Self::Reply {
        if !self.incremental.is_empty() {
            let removed = std::mem::take(&mut self.incremental);
            if let Some(cb) = &self.on_change {
                cb(ChangeEvent::Cleared { removed });
            }
        }
    }
}

/// 触发上下文压缩。Reply = `()`。
///
/// 根据 [`CompressStrategy`] 对 incremental 区执行压缩。
/// 摘要由后端 LLM 生成，通过 `opts` 传递请求选项。
pub struct CompressMsg<O> {
    /// 压缩策略
    pub strategy: CompressStrategy,
    /// 传递给后端 `send()` 的请求选项
    pub opts: O,
}

impl<B: ContextBackend> Message<CompressMsg<B::Opts>> for AgentContext<B> {
    type Reply = ();

    async fn handle(
        &mut self,
        msg: CompressMsg<B::Opts>,
        _ctx: &mut Context<Self, Self::Reply>,
    ) -> Self::Reply {
        match msg.strategy {
            CompressStrategy::Summarize { keep, prompt } => {
                let total = self.incremental.len();
                if total > keep {
                    let split = total - keep;
                    let to_summarize: Vec<B::Message> = self.incremental.drain(..split).collect();
                    if !to_summarize.is_empty() {
                        let summary_prompt = prompt.unwrap_or_else(Self::default_summary_prompt);
                        let mut summary_messages =
                            vec![self.backend.system_message(summary_prompt)];
                        summary_messages.append(&mut self.compressed);
                        summary_messages.extend(to_summarize);
                        let result = self.backend.send(&summary_messages, &msg.opts).await;
                        if let Ok(response) = result {
                            if let Ok(raw_msgs) =
                                self.backend.extract_messages_from_backend_response(
                                    std::slice::from_ref(&response),
                                )
                            {
                                if let Ok(request_msgs) = self.backend.to_request_messages(raw_msgs)
                                {
                                    let summary: Vec<B::Message> = request_msgs
                                        .into_iter()
                                        .map(|msg| self.backend.to_system_message(msg))
                                        .collect();
                                    let kept: Vec<B::Message> =
                                        self.incremental.drain(..).collect();
                                    let (final_summary, final_kept) =
                                        if let Some(cb) = &self.on_compressed {
                                            cb(summary, kept)
                                        } else {
                                            (summary, kept)
                                        };
                                    self.compressed = final_summary;
                                    self.incremental = final_kept;
                                } else {
                                    log::warn!("压缩摘要转换请求格式失败，已跳过");
                                }
                            } else {
                                log::warn!("压缩摘要提取消息失败，已跳过");
                            }
                        }
                    }
                }
            }
        }
    }
}

/// 非流式对话。Reply = `Result<Response, AgentError>`。
///
/// 将三区消息拼接后，如 opts 实现了 [`ScratchOpts`] 且 `scratch()` 返回 `Some`，
/// 追加一条 system 消息作为 Scratch，然后发送给后端 LLM。
/// 成功后自动将响应消息存入 incremental 区，触发 [`ChangeEvent::Appended`]。
pub struct SendMsg<O> {
    /// 传递给后端 `send()` 的请求选项（需实现 [`ScratchOpts`]）
    pub opts: O,
}

impl<B: ContextBackend> Message<SendMsg<B::Opts>> for AgentContext<B> {
    type Reply = Result<B::Response, AgentError>;

    async fn handle(
        &mut self,
        msg: SendMsg<B::Opts>,
        _ctx: &mut Context<Self, Self::Reply>,
    ) -> Self::Reply {
        let scratch = msg.opts.scratch().map(|s| s.to_string());
        let mut all_messages: Vec<B::Message> = self
            .immutable
            .iter()
            .chain(self.compressed.iter())
            .chain(self.incremental.iter())
            .cloned()
            .collect();
        if let Some(content) = scratch {
            all_messages.push(self.backend.system_message(content));
        }
        let response = self.backend.send(&all_messages, &msg.opts).await?;
        let raw_msgs = self
            .backend
            .extract_messages_from_backend_response(std::slice::from_ref(&response))?;
        let request_msgs = self.backend.to_request_messages(raw_msgs)?;
        for msg in &request_msgs {
            self.incremental.push(msg.clone());
            if let Some(cb) = &self.on_change {
                cb(ChangeEvent::Appended(msg.clone()));
            }
        }
        Ok(response)
    }
}

/// 流式对话。Reply = [`AgentSendStream<B>`]。
///
/// 将三区消息拼接后，如 opts 实现了 [`ScratchOpts`] 且 `scratch()` 返回 `Some`，
/// 追加一条 system 消息作为 Scratch，然后发送给后端 LLM。
/// 调用者消费流直到结束，drop 时自动将响应消息存入 incremental 区。
pub struct SendStreamMsg<O> {
    /// 传递给后端 `send_stream()` 的请求选项（需实现 [`ScratchOpts`]）
    pub opts: O,
}

impl<B: ContextBackend + Clone> Message<SendStreamMsg<B::Opts>> for AgentContext<B> {
    type Reply = AgentSendStream<B>;

    async fn handle(
        &mut self,
        msg: SendStreamMsg<B::Opts>,
        ctx: &mut Context<Self, Self::Reply>,
    ) -> Self::Reply {
        let scratch = msg.opts.scratch().map(|s| s.to_string());
        let mut all_messages: Vec<B::Message> = self
            .immutable
            .iter()
            .chain(self.compressed.iter())
            .chain(self.incremental.iter())
            .cloned()
            .collect();
        if let Some(content) = scratch {
            all_messages.push(self.backend.system_message(content));
        }
        let stream = self.backend.send_stream(all_messages, msg.opts);
        AgentSendStream::new(
            self.backend.clone(),
            stream,
            ctx.actor_ref().clone(),
            self.on_change.clone(),
        )
    }
}

/// 估算三区全部消息的 token 数量。Reply = `usize`。
///
/// 委托给后端的 [`estimate_tokens`](ContextBackend::estimate_tokens)。
/// 失败时降级返回 0。
pub struct EstimateTokensMsg;

impl<B: ContextBackend> Message<EstimateTokensMsg> for AgentContext<B> {
    type Reply = usize;

    async fn handle(
        &mut self,
        _msg: EstimateTokensMsg,
        _ctx: &mut Context<Self, Self::Reply>,
    ) -> Self::Reply {
        let all: Vec<B::Message> = self
            .immutable
            .iter()
            .chain(self.compressed.iter())
            .chain(self.incremental.iter())
            .cloned()
            .collect();
        self.backend.estimate_tokens(&all).await.unwrap_or(0)
    }
}

/// 检查上下文是否已满（token 数 >= [`context_window`](ContextBackend::context_window)）。Reply = `bool`。
///
/// 估算失败时降级返回 `true`（安全策略：宁可误判已满，不可溢出）。
pub struct IsFullMsg;

impl<B: ContextBackend> Message<IsFullMsg> for AgentContext<B> {
    type Reply = bool;

    async fn handle(
        &mut self,
        _msg: IsFullMsg,
        _ctx: &mut Context<Self, Self::Reply>,
    ) -> Self::Reply {
        let all: Vec<B::Message> = self
            .immutable
            .iter()
            .chain(self.compressed.iter())
            .chain(self.incremental.iter())
            .cloned()
            .collect();
        let tokens = self
            .backend
            .estimate_tokens(&all)
            .await
            .unwrap_or(usize::MAX);
        tokens >= self.backend.context_window()
    }
}

/// 将三区全部消息导出为 JSONL 字符串，每行一条 JSON。Reply = `Result<String, AgentError>`。
///
/// 消息按 immutable → compressed → incremental 顺序输出。
pub struct ToJsonlMsg;

impl<B: ContextBackend> Message<ToJsonlMsg> for AgentContext<B> {
    type Reply = Result<String, AgentError>;

    async fn handle(
        &mut self,
        _msg: ToJsonlMsg,
        _ctx: &mut Context<Self, Self::Reply>,
    ) -> Self::Reply {
        let lines: Vec<String> = self
            .immutable
            .iter()
            .chain(self.compressed.iter())
            .chain(self.incremental.iter())
            .map(|m| self.backend.message_to_jsonl(m))
            .collect::<Result<_, _>>()?;
        Ok(lines.join("\n"))
    }
}

/// 从 JSONL 字符串加载消息到 incremental 区。Reply = `Result<(), AgentError>`。
///
/// 每行一条 JSON，空行跳过。解析失败或触发 `preserve_reasoning` 时返回错误。
/// 加载的消息逐条触发 [`ChangeEvent::Appended`] 回调。
pub struct FromJsonlMsg {
    /// JSONL 字符串，每行一条消息
    pub jsonl: String,
}

impl<B: ContextBackend> Message<FromJsonlMsg> for AgentContext<B> {
    type Reply = Result<(), AgentError>;

    async fn handle(
        &mut self,
        msg: FromJsonlMsg,
        _ctx: &mut Context<Self, Self::Reply>,
    ) -> Self::Reply {
        for line in msg.jsonl.lines() {
            let line = line.trim();
            if line.is_empty() {
                continue;
            }
            let message: B::Message = self.backend.message_from_jsonl(line)?;
            self.incremental.push(message.clone());
            if let Some(ref cb) = self.on_change {
                cb(ChangeEvent::Appended(message));
            }
        }
        Ok(())
    }
}