iAs/ias-context/src/routes.rs

use axum::extract::Path;
use axum::http::{HeaderMap, StatusCode};
use axum::routing::{patch, post};
use axum::{Extension, Json, Router};
use serde_json::{Value, json};
use sqlx::PgPool;

use crate::model::{
    ContextScope, MemoryCreateRequest, MemoryDeleteRequest, MemoryDeleteResponse, MemoryResponse,
    MemorySearchRequest, MemorySearchResponse, MemoryUpdateRequest, ToolCallInstructionGetRequest,
    ToolCallInstructionResponse, ToolCallInstructionUpdateRequest,
};
use crate::repository::ContextRepository;

type ApiError = (StatusCode, Json<Value>);

#[derive(Clone)]
pub struct ContextRouteState {
    pub db: PgPool,
}

impl ContextRouteState {
    pub fn new(db: PgPool) -> Self {
        Self { db }
    }
}

pub fn routes<S>(db: PgPool) -> Router<S>
where
    S: Clone + Send + Sync + 'static,
{
    Router::new()
        .route("/v1/context/memories", post(create_memory))
        .route("/v1/context/memories/search", post(search_memories))
        .route(
            "/v1/context/memories/{id}",
            patch(update_memory).delete(delete_memory),
        )
        .route(
            "/v1/context/tool-instructions",
            post(get_tool_call_instruction).patch(update_tool_call_instruction),
        )
        .layer(Extension(ContextRouteState::new(db)))
}

async fn create_memory(
    headers: HeaderMap,
    Extension(state): Extension<ContextRouteState>,
    Json(request): Json<MemoryCreateRequest>,
) -> Result<Json<MemoryResponse>, ApiError> {
    authorize(&headers)?;
    let scope = validate_scope(request.identity.scope())?;
    let content = required_text(&request.content, "content")?;
    let metadata = request.metadata.unwrap_or_else(|| json!({})).to_string();

    let memory = ContextRepository::create_memory(&state.db, &scope, content, metadata)
        .await
        .map_err(api_internal_error)?;

    Ok(Json(MemoryResponse {
        success: true,
        memory,
    }))
}

async fn search_memories(
    headers: HeaderMap,
    Extension(state): Extension<ContextRouteState>,
    Json(request): Json<MemorySearchRequest>,
) -> Result<Json<MemorySearchResponse>, ApiError> {
    authorize(&headers)?;
    let scope = validate_scope(request.identity.scope())?;
    let limit = request.limit.unwrap_or(20).clamp(1, 50);
    let query = request
        .query
        .as_deref()
        .map(str::trim)
        .filter(|query| !query.is_empty());

    let memories = ContextRepository::search_memories(&state.db, &scope.scope_key, query, limit)
        .await
        .map_err(api_internal_error)?;

    Ok(Json(MemorySearchResponse {
        success: true,
        memories,
    }))
}

async fn update_memory(
    headers: HeaderMap,
    Extension(state): Extension<ContextRouteState>,
    Path(id): Path<i64>,
    Json(request): Json<MemoryUpdateRequest>,
) -> Result<Json<MemoryResponse>, ApiError> {
    authorize(&headers)?;
    let scope = validate_scope(request.identity.scope())?;
    let content = required_text(&request.content, "content")?;
    let metadata = request.metadata.map(|value| value.to_string());

    let memory =
        ContextRepository::update_memory(&state.db, &scope.scope_key, id, content, metadata)
            .await
            .map_err(api_internal_error)?
            .ok_or_else(|| api_not_found("长期记忆不存在"))?;

    Ok(Json(MemoryResponse {
        success: true,
        memory,
    }))
}

async fn delete_memory(
    headers: HeaderMap,
    Extension(state): Extension<ContextRouteState>,
    Path(id): Path<i64>,
    Json(request): Json<MemoryDeleteRequest>,
) -> Result<Json<MemoryDeleteResponse>, ApiError> {
    authorize(&headers)?;
    let scope = validate_scope(request.identity.scope())?;
    let deleted = ContextRepository::delete_memory(&state.db, &scope.scope_key, id)
        .await
        .map_err(api_internal_error)?;

    Ok(Json(MemoryDeleteResponse {
        success: true,
        deleted,
    }))
}

async fn update_tool_call_instruction(
    headers: HeaderMap,
    Extension(state): Extension<ContextRouteState>,
    Json(request): Json<ToolCallInstructionUpdateRequest>,
) -> Result<Json<ToolCallInstructionResponse>, ApiError> {
    authorize(&headers)?;
    let scope = validate_scope(request.identity.scope())?;
    let content = normalize_prompt_content(request.content)?;
    let reason = optional_limited_text(request.reason, 500);

    let instruction =
        ContextRepository::upsert_tool_call_instruction(&state.db, &scope, content, reason, "ai")
            .await
            .map_err(api_internal_error)?;

    Ok(Json(ToolCallInstructionResponse {
        success: true,
        instruction: Some(instruction),
    }))
}

async fn get_tool_call_instruction(
    headers: HeaderMap,
    Extension(state): Extension<ContextRouteState>,
    Json(request): Json<ToolCallInstructionGetRequest>,
) -> Result<Json<ToolCallInstructionResponse>, ApiError> {
    authorize(&headers)?;
    let scope = validate_scope(request.identity.scope())?;
    let instruction = ContextRepository::get_tool_call_instruction(&state.db, &scope.scope_key)
        .await
        .map_err(api_internal_error)?;

    Ok(Json(ToolCallInstructionResponse {
        success: true,
        instruction,
    }))
}

fn authorize(headers: &HeaderMap) -> Result<(), ApiError> {
    let expected = context_api_key()?;
    let token = headers
        .get("x-ias-context-token")
        .and_then(|value| value.to_str().ok())
        .map(str::trim)
        .filter(|value| !value.is_empty())
        .or_else(|| bearer_token(headers));

    if token.is_some_and(|token| token == expected) {
        Ok(())
    } else {
        Err(api_unauthorized("context 内部接口鉴权失败"))
    }
}

fn bearer_token(headers: &HeaderMap) -> Option<&str> {
    let value = headers.get("authorization")?.to_str().ok()?.trim();
    value
        .strip_prefix("Bearer ")
        .or_else(|| value.strip_prefix("bearer "))
        .map(str::trim)
        .filter(|value| !value.is_empty())
}

fn context_api_key() -> Result<String, ApiError> {
    std::env::var("IA_CONTEXT_API_KEY")
        .or_else(|_| std::env::var("IA_CONTEXT_TOKEN"))
        .or_else(|_| std::env::var("ROUTER_API_KEY"))
        .map(|value| value.trim().to_string())
        .ok()
        .filter(|value| !value.is_empty())
        .ok_or_else(|| api_unavailable("未配置 context 内部访问密钥"))
}

fn validate_scope(scope: ContextScope) -> Result<ContextScope, ApiError> {
    if scope.account_id.trim().is_empty() {
        return Err(api_bad_request("account_id 不能为空"));
    }
    if scope.from_user_id.trim().is_empty() {
        return Err(api_bad_request("from_user_id 不能为空"));
    }
    Ok(scope)
}

fn required_text(value: &str, field: &str) -> Result<String, ApiError> {
    let value = value.trim();
    if value.is_empty() {
        return Err(api_bad_request(format!("{field} 不能为空")));
    }
    Ok(value.to_string())
}

fn normalize_prompt_content(value: String) -> Result<String, ApiError> {
    let value = value.trim().to_string();
    if value.chars().count() > 8_000 {
        return Err(api_bad_request("content 不能超过 8000 个字符"));
    }
    Ok(value)
}

fn optional_limited_text(value: Option<String>, max_chars: usize) -> Option<String> {
    value
        .map(|value| value.trim().chars().take(max_chars).collect::<String>())
        .filter(|value| !value.is_empty())
}

fn api_bad_request(message: impl Into<String>) -> ApiError {
    (
        StatusCode::BAD_REQUEST,
        Json(json!({
            "success": false,
            "message": message.into(),
        })),
    )
}

fn api_not_found(message: impl Into<String>) -> ApiError {
    (
        StatusCode::NOT_FOUND,
        Json(json!({
            "success": false,
            "message": message.into(),
        })),
    )
}

fn api_unauthorized(message: impl Into<String>) -> ApiError {
    (
        StatusCode::UNAUTHORIZED,
        Json(json!({
            "success": false,
            "message": message.into(),
        })),
    )
}

fn api_unavailable(message: impl Into<String>) -> ApiError {
    (
        StatusCode::SERVICE_UNAVAILABLE,
        Json(json!({
            "success": false,
            "message": message.into(),
        })),
    )
}

fn api_internal_error(error: anyhow::Error) -> ApiError {
    eprintln!("context 接口失败: {error}");
    (
        StatusCode::INTERNAL_SERVER_ERROR,
        Json(json!({
            "success": false,
            "message": "context 接口失败",
        })),
    )
}