mnote/rust/crates/mnote-web/src/routes/page_ai_workflow.rs

use crate::app::AppState;
use crate::context::RequestContext;
use crate::error::WebError;
use crate::hermes_tools::{block, ToolCallInput};
use axum::extract::{Extension, State};
use axum::http::{HeaderMap, StatusCode};
use axum::Json;
use serde_json::{json, Value};
use std::fs;
use std::path::PathBuf;
use std::time::Instant;
use tracing::info;

pub async fn block_edit_workflow(
    State(state): State<AppState>,
    Extension(context): Extension<RequestContext>,
    Json(payload): Json<Value>,
) -> Result<(StatusCode, HeaderMap, Json<Value>), WebError> {
    let started = Instant::now();
    let workspace_id = string_field(&payload, "workspaceId")
        .or_else(|| context.workspace.workspace_id.clone())
        .ok_or_else(|| {
            WebError::bad_request_code("page_ai_workflow_bad_request", "缺少 workspaceId")
                .with_context(&context)
        })?;
    let document_id = string_field(&payload, "documentId").ok_or_else(|| {
        WebError::bad_request_code("page_ai_workflow_bad_request", "缺少 documentId")
            .with_context(&context)
    })?;
    let message = string_field(&payload, "message").ok_or_else(|| {
        WebError::bad_request_code("page_ai_workflow_bad_request", "缺少 message")
            .with_context(&context)
    })?;
    let trace_id =
        string_field(&payload, "traceId").unwrap_or_else(|| context.trace.trace_id.clone());
    let session_id = string_field(&payload, "sessionId")
        .unwrap_or_else(|| format!("page_ai_fast_{}", context.trace.request_id));
    let run_id = string_field(&payload, "runId").unwrap_or_else(|| session_id.clone());
    info!(
        trace_id = %trace_id,
        run_id = %run_id,
        workspace_id = %workspace_id,
        document_id = %document_id,
        "mnote page AI block workflow started"
    );
    if !looks_like_block_edit(&message) {
        return Err(WebError::bad_request_code(
            "page_ai_workflow_not_block_edit",
            "当前请求不像块编辑任务，交给通用页面 AI",
        )
        .with_context(&context));
    }
    let page_context = payload.get("pageContext").cloned().unwrap_or(Value::Null);
    let ai_context = page_context
        .get("aiContext")
        .cloned()
        .or_else(|| payload.get("aiContext").cloned())
        .ok_or_else(|| {
            WebError::bad_request_code(
                "page_ai_workflow_missing_context",
                "块编辑快路径缺少 mnote.page_ai_context.v1",
            )
            .with_context(&context)
        })?;
    let profile = string_field(&payload, "profile").unwrap_or_else(|| "mnoteai".into());
    let model_started = Instant::now();
    let (operations, operation_source) = if let Some(operations) =
        direct_block_edit_operations(&message)
    {
        (operations, "local_rule")
    } else {
        let model_output = call_block_edit_model(&context, &profile, &message, &ai_context).await?;
        (extract_operations_from_model_text(&model_output)?, "model")
    };
    info!(
        trace_id = %trace_id,
        run_id = %run_id,
        operations = operations.len(),
        operation_source = operation_source,
        model_ms = model_started.elapsed().as_millis(),
        "mnote page AI block workflow model completed"
    );
    if operations.is_empty() {
        return Err(WebError::bad_request_code(
            "page_ai_workflow_empty_operations",
            "模型未返回块操作",
        )
        .with_context(&context));
    }
    let allowed_target_block_ids = ai_context
        .get("allowedTargetBlockIds")
        .cloned()
        .unwrap_or_else(|| json!([]));
    let actor_id =
        if context.auth.actor_id.trim().is_empty() || context.auth.actor_id == "anonymous" {
            state.config().dev_user_id.clone()
        } else {
            context.auth.actor_id.clone()
        };
    let apply_input = ToolCallInput {
        tool_name: "mnote.doc.apply_block_ops".into(),
        workspace_id: Some(workspace_id.clone()),
        document_id: Some(document_id.clone()),
        actor_id: Some(actor_id),
        profile: Some(profile),
        session_id: Some(session_id),
        run_id: Some(run_id.clone()),
        tool_call_id: Some(format!("fast_apply_{}", context.trace.request_id)),
        trace_id: Some(trace_id.clone()),
        idempotency_key: Some(format!("page_ai_fast_apply_{}", context.trace.request_id)),
        dry_run: Some(false),
        capability_scope: Some(vec!["block.write".into(), "page.write".into()]),
        args: Some(json!({
            "operations": operations,
            "allowedTargetBlockIds": allowed_target_block_ids
        })),
    };
    let apply_started = Instant::now();
    let apply_result = block::doc_apply_block_ops(&state, &context, &apply_input).await?;
    let apply_ms = apply_started.elapsed().as_millis();
    info!(
        trace_id = %trace_id,
        run_id = %run_id,
        apply_ms = apply_ms,
        total_ms = started.elapsed().as_millis(),
        "mnote page AI block workflow completed"
    );
    Ok((
        StatusCode::OK,
        HeaderMap::new(),
        Json(json!({
            "ok": true,
            "schema": "mnote.page_ai_block_edit_workflow.v1",
            "fastPath": true,
            "documentId": document_id,
            "workspaceId": workspace_id,
            "runId": run_id,
            "traceId": trace_id,
            "operationSource": operation_source,
            "operations": apply_input.arg_value("operations").unwrap_or_else(|| json!([])),
            "applyResult": apply_result,
            "message": "已通过页面块编辑快路径完成写入。",
            "timingsMs": {
                "total": started.elapsed().as_millis(),
                "apply": apply_ms
            }
        })),
    ))
}

fn extract_operations_from_model_text(text: &str) -> Result<Vec<Value>, WebError> {
    let parsed = parse_model_json(text)?;
    if let Some(content) = parsed
        .pointer("/choices/0/message/content")
        .and_then(Value::as_str)
    {
        return extract_operations_from_model_text(content);
    }
    if let Some(operations) = parsed.get("operations").and_then(Value::as_array) {
        return Ok(operations.clone());
    }
    if let Some(operations) = parsed
        .get("arguments")
        .and_then(|value| value.get("operations"))
        .and_then(Value::as_array)
    {
        return Ok(operations.clone());
    }
    if let Some(operations) = parsed.as_array() {
        return Ok(operations.clone());
    }
    Err(WebError::bad_request_code(
        "page_ai_workflow_bad_model_output",
        "模型输出未包含 operations",
    ))
}

fn parse_model_json(text: &str) -> Result<Value, WebError> {
    let trimmed = strip_code_fence(text.trim());
    if let Ok(value) = serde_json::from_str::<Value>(&trimmed) {
        return Ok(value);
    }
    if let Some(slice) = first_json_slice(&trimmed) {
        if let Ok(value) = serde_json::from_str::<Value>(slice) {
            return Ok(value);
        }
    }
    Err(WebError::bad_request_code(
        "page_ai_workflow_bad_model_json",
        "模型输出不是可解析 JSON",
    ))
}

fn strip_code_fence(text: &str) -> String {
    let trimmed = text.trim();
    if !trimmed.starts_with("```") {
        return trimmed.to_string();
    }
    let without_open = trimmed.lines().skip(1).collect::<Vec<_>>().join("\n");
    without_open
        .trim()
        .strip_suffix("```")
        .unwrap_or(without_open.trim())
        .trim()
        .to_string()
}

fn first_json_slice(text: &str) -> Option<&str> {
    let start = text.find('{').or_else(|| text.find('['))?;
    let open = text.as_bytes()[start] as char;
    let close = if open == '{' { '}' } else { ']' };
    let mut depth = 0usize;
    let mut in_string = false;
    let mut escaped = false;
    for (offset, ch) in text[start..].char_indices() {
        if in_string {
            if escaped {
                escaped = false;
            } else if ch == '\\' {
                escaped = true;
            } else if ch == '"' {
                in_string = false;
            }
            continue;
        }
        if ch == '"' {
            in_string = true;
        } else if ch == open {
            depth += 1;
        } else if ch == close {
            depth = depth.saturating_sub(1);
            if depth == 0 {
                return Some(&text[start..start + offset + ch.len_utf8()]);
            }
        }
    }
    None
}

async fn call_block_edit_model(
    context: &RequestContext,
    profile: &str,
    message: &str,
    ai_context: &Value,
) -> Result<String, WebError> {
    let model = workflow_model_config(profile);
    let page_xml = ai_context
        .get("pageXml")
        .and_then(Value::as_str)
        .unwrap_or_default();
    let page_text = ai_context
        .get("pageText")
        .and_then(Value::as_str)
        .unwrap_or_default();
    let allowed = ai_context
        .get("allowedTargetBlockIds")
        .cloned()
        .unwrap_or_else(|| json!([]));
    let body = json!({
        "model": model.model,
        "temperature": 0,
        "max_tokens": 900,
        "response_format": {"type": "json_object"},
        "messages": [
            {
                "role": "system",
                "content": "你是 mnote 页面块编辑 workflow。只输出 JSON：{\"operations\":[...] ,\"summary\":\"...\"}。operations 的 op 只能是 replace、insert_after、delete、move_after。优先使用 page_xml 中的 block id；禁止输出解释文字。"
            },
            {
                "role": "user",
                "content": format!(
                    "用户指令：{}\n\nallowedTargetBlockIds：{}\n\npage_xml：\n{}\n\npage_text：\n{}",
                    message,
                    allowed,
                    page_xml,
                    page_text
                )
            }
        ]
    });
    let url = format!("{}/chat/completions", model.base_url.trim_end_matches('/'));
    let response = reqwest::Client::builder()
        .timeout(std::time::Duration::from_secs(20))
        .build()
        .map_err(|error| {
            WebError::internal(format!("页面 AI workflow HTTP client 构造失败: {error}"))
        })?
        .post(url)
        .bearer_auth(model.api_key)
        .json(&body)
        .send()
        .await
        .map_err(|error| {
            WebError::bad_gateway_code(
                "page_ai_workflow_model_unavailable",
                format!("页面 AI workflow 模型请求失败: {error}"),
            )
            .with_context(context)
        })?;
    let status = response.status();
    let text = response.text().await.unwrap_or_default();
    if !status.is_success() {
        return Err(WebError::bad_gateway_code(
            "page_ai_workflow_model_failed",
            format!("页面 AI workflow 模型返回失败: {status}"),
        )
        .with_context(context));
    }
    let payload = parse_model_json(&text)?;
    payload
        .pointer("/choices/0/message/content")
        .and_then(Value::as_str)
        .map(ToOwned::to_owned)
        .ok_or_else(|| {
            WebError::bad_gateway_code(
                "page_ai_workflow_model_no_content",
                "页面 AI workflow 模型响应缺少 message.content",
            )
            .with_context(context)
        })
}

struct WorkflowModelConfig {
    model: String,
    base_url: String,
    api_key: String,
}

fn workflow_model_config(profile: &str) -> WorkflowModelConfig {
    let config = fs::read_to_string(profile_config_path(profile)).unwrap_or_default();
    let provider =
        yaml_path_value(&config, &["model", "provider"]).unwrap_or_else(|| "deepseek".into());
    let model = yaml_path_value(&config, &["model", "default"])
        .or_else(|| yaml_path_value(&config, &["providers", &provider, "model"]))
        .unwrap_or_else(|| "deepseek-v4-flash".into());
    let base_url = yaml_path_value(&config, &["model", "base_url"])
        .or_else(|| yaml_path_value(&config, &["providers", &provider, "base_url"]))
        .unwrap_or_else(|| "https://api.deepseek.com/v1".into());
    let api_key = yaml_path_value(&config, &["model", "api_key"])
        .or_else(|| yaml_path_value(&config, &["providers", &provider, "api_key"]))
        .or_else(|| {
            yaml_path_value(&config, &["model", "key_env"])
                .or_else(|| yaml_path_value(&config, &["providers", &provider, "key_env"]))
                .and_then(|env_key| std::env::var(env_key).ok())
        })
        .unwrap_or_default();
    WorkflowModelConfig {
        model,
        base_url,
        api_key,
    }
}

fn string_field(payload: &Value, key: &str) -> Option<String> {
    payload
        .get(key)
        .and_then(Value::as_str)
        .map(str::trim)
        .filter(|value| !value.is_empty())
        .map(ToOwned::to_owned)
}

fn looks_like_block_edit(message: &str) -> bool {
    [
        "新增", "添加", "插入", "删除", "删掉", "修改", "替换", "改成", "移动", "移到", "move",
        "replace", "delete", "insert",
    ]
    .iter()
    .any(|needle| message.contains(needle))
}

fn direct_block_edit_operations(message: &str) -> Option<Vec<Value>> {
    let mut operations = Vec::new();
    for clause in message
        .split(|ch| matches!(ch, '；' | ';' | '\n'))
        .map(str::trim)
        .filter(|value| !value.is_empty())
    {
        let quoted = quoted_segments(clause);
        if (clause.contains("替换") || clause.contains("改成")) && quoted.len() >= 2 {
            operations.push(json!({
                "op": "replace",
                "matchText": quoted[0],
                "content": quoted[1]
            }));
        } else if (clause.contains("插入") || clause.contains("新增") || clause.contains("添加"))
            && quoted.len() >= 2
        {
            operations.push(json!({
                "op": "insert_after",
                "matchText": quoted[0],
                "content": quoted[1]
            }));
        } else if (clause.contains("删除") || clause.contains("删掉")) && !quoted.is_empty() {
            operations.push(json!({
                "op": "delete",
                "matchText": quoted[0]
            }));
        }
    }
    if operations.is_empty() {
        None
    } else {
        Some(operations)
    }
}

fn quoted_segments(value: &str) -> Vec<String> {
    let mut segments = Vec::new();
    let mut start: Option<char> = None;
    let mut current = String::new();
    for ch in value.chars() {
        match (start, ch) {
            (None, '「' | '“' | '"') => {
                start = Some(ch);
                current.clear();
            }
            (Some('「'), '」') | (Some('“'), '”') | (Some('"'), '"') => {
                if !current.trim().is_empty() {
                    segments.push(current.trim().to_string());
                }
                current.clear();
                start = None;
            }
            (Some(_), _) => current.push(ch),
            (None, _) => {}
        }
    }
    segments
}

fn hermes_home() -> PathBuf {
    std::env::var("HERMES_HOME")
        .ok()
        .filter(|value| !value.trim().is_empty())
        .map(PathBuf::from)
        .or_else(|| {
            std::env::var("HOME")
                .ok()
                .map(|home| PathBuf::from(home).join(".hermes"))
        })
        .unwrap_or_else(|| PathBuf::from(".hermes"))
}

fn profile_config_path(profile: &str) -> PathBuf {
    let home = hermes_home();
    let profile = profile.trim();
    if profile.is_empty() || profile == "default" {
        return home.join("config.yaml");
    }
    let candidate = home.join("profiles").join(profile);
    if candidate.exists() {
        candidate.join("config.yaml")
    } else {
        home.join("config.yaml")
    }
}

fn yaml_path_value(content: &str, path: &[&str]) -> Option<String> {
    let mut stack: Vec<(usize, String)> = Vec::new();
    for raw_line in content.lines() {
        let line = raw_line.trim_end_matches('\r');
        let trimmed = line.trim();
        if trimmed.is_empty() || trimmed.starts_with('#') || !line.contains(':') {
            continue;
        }
        let indent = line.chars().take_while(|ch| ch.is_whitespace()).count();
        while stack
            .last()
            .map(|(level, _)| *level >= indent)
            .unwrap_or(false)
        {
            stack.pop();
        }
        let Some((key, value)) = trimmed.split_once(':') else {
            continue;
        };
        let key = key.trim().trim_matches('"').trim_matches('\'').to_string();
        let value = value
            .trim()
            .trim_matches('"')
            .trim_matches('\'')
            .to_string();
        stack.push((indent, key));
        if stack.len() == path.len()
            && stack
                .iter()
                .zip(path.iter())
                .all(|((_, key), expected)| key == expected)
            && !value.is_empty()
        {
            return Some(value);
        }
    }
    None
}

#[cfg(test)]
mod tests {
    use super::{direct_block_edit_operations, extract_operations_from_model_text};

    #[test]
    fn extracts_operations_from_fenced_model_json() {
        let operations = extract_operations_from_model_text(
            r#"```json
{"operations":[{"op":"replace","matchText":"旧文本","content":"新文本"}],"summary":"ok"}
```"#,
        )
        .expect("operations");
        assert_eq!(operations.len(), 1);
        assert_eq!(operations[0]["op"], "replace");
        assert_eq!(operations[0]["matchText"], "旧文本");
    }

    #[test]
    fn parses_direct_chinese_block_operations() {
        let operations = direct_block_edit_operations(
            "把「第二段」替换为「第二段已修改」；在「第一段」后插入「插入段」；删除「第三段」。只简短回复结果。",
        )
        .expect("operations");
        assert_eq!(operations.len(), 3);
        assert_eq!(operations[0]["op"], "replace");
        assert_eq!(operations[0]["matchText"], "第二段");
        assert_eq!(operations[0]["content"], "第二段已修改");
        assert_eq!(operations[1]["op"], "insert_after");
        assert_eq!(operations[1]["matchText"], "第一段");
        assert_eq!(operations[1]["content"], "插入段");
        assert_eq!(operations[2]["op"], "delete");
        assert_eq!(operations[2]["matchText"], "第三段");
    }
}