iAs/ias-service/src/context_simulator_data.rs

use chrono::{DateTime, Utc};
use serde::Serialize;
use serde_json::Value;
use sqlx::{PgPool, Row};

// ── 数据库行 ────────────────────────────────────────────

#[derive(Debug, Clone, Serialize)]
pub struct SessionRow {
    pub id: i64,
    pub started_at: String,
    pub last_message_at: String,
    pub ended_at: Option<DateTime<Utc>>,
    pub end_reason: Option<String>,
    pub summary: Option<String>,
}

#[derive(Debug, Clone, Serialize)]
pub struct MessageRow {
    pub id: i64,
    pub role: String,
    pub content: String,
    pub tool_call_id: Option<String>,
    pub tool_calls: Option<String>,
    pub created_at: String,
}

#[derive(Debug, Clone, Serialize)]
pub struct MemoryRow {
    pub id: i64,
    pub content: String,
    pub metadata: String,
    pub created_at: DateTime<Utc>,
    pub updated_at: DateTime<Utc>,
    pub deleted_at: Option<DateTime<Utc>>,
}

#[derive(Debug, Clone, Serialize)]
pub struct SummaryRow {
    pub id: i64,
    pub summary_type: String,
    pub content: String,
    pub started_at: Option<DateTime<Utc>>,
    pub ended_at: Option<DateTime<Utc>>,
    pub created_at: DateTime<Utc>,
}

#[derive(Debug, Clone, Serialize)]
pub struct ToolRow {
    pub name: String,
    pub description: String,
    pub source: String,
    pub path: Option<String>,
    pub params: Vec<ToolParamRow>,
}

#[derive(Debug, Clone, Serialize)]
pub struct ToolParamRow {
    pub name: String,
    pub param_type: String,
    pub required: bool,
    pub description: String,
}

// ── 数据库查询 ──────────────────────────────────────────

pub async fn list_sessions(pool: &PgPool, scope_key: &str) -> anyhow::Result<Vec<SessionRow>> {
    let rows = sqlx::query(
        r#"
        SELECT id, started_at, last_message_at, ended_at, end_reason, summary
        FROM ias_context_sessions
        WHERE scope_key = $1
        ORDER BY started_at DESC, id DESC
        LIMIT 50
        "#,
    )
    .bind(scope_key)
    .fetch_all(pool)
    .await?;

    let sessions = rows
        .into_iter()
        .map(|row| SessionRow {
            id: row.try_get("id").unwrap_or_default(),
            started_at: row
                .try_get::<DateTime<Utc>, _>("started_at")
                .map(|t| t.format("%Y-%m-%d %H:%M:%S").to_string())
                .unwrap_or_default(),
            last_message_at: row
                .try_get::<DateTime<Utc>, _>("last_message_at")
                .map(|t| t.format("%Y-%m-%d %H:%M:%S").to_string())
                .unwrap_or_default(),
            ended_at: row.try_get("ended_at").ok(),
            end_reason: row.try_get("end_reason").ok(),
            summary: row.try_get("summary").ok(),
        })
        .collect();

    Ok(sessions)
}

pub async fn list_session_messages(
    pool: &PgPool,
    session_id: i64,
) -> anyhow::Result<Vec<MessageRow>> {
    let rows = sqlx::query(
        r#"
        SELECT id, role, content, tool_call_id, tool_calls, created_at
        FROM ias_context_messages
        WHERE session_id = $1
        ORDER BY id ASC
        LIMIT 200
        "#,
    )
    .bind(session_id)
    .fetch_all(pool)
    .await?;

    let messages = rows
        .into_iter()
        .map(|row| MessageRow {
            id: row.try_get("id").unwrap_or_default(),
            role: row.try_get("role").unwrap_or_default(),
            content: row.try_get("content").unwrap_or_default(),
            tool_call_id: row.try_get("tool_call_id").ok(),
            tool_calls: row.try_get("tool_calls").ok(),
            created_at: row
                .try_get::<DateTime<Utc>, _>("created_at")
                .map(|t| t.format("%Y-%m-%d %H:%M:%S").to_string())
                .unwrap_or_default(),
        })
        .collect();

    Ok(messages)
}

pub async fn list_memories(pool: &PgPool, scope_key: &str) -> anyhow::Result<Vec<MemoryRow>> {
    let rows = sqlx::query(
        r#"
        SELECT id, content, metadata, created_at, updated_at, deleted_at
        FROM ias_long_term_memories
        WHERE scope_key = $1
        ORDER BY updated_at DESC, id DESC
        LIMIT 30
        "#,
    )
    .bind(scope_key)
    .fetch_all(pool)
    .await?;

    let memories = rows
        .into_iter()
        .map(|row| MemoryRow {
            id: row.try_get("id").unwrap_or_default(),
            content: row.try_get("content").unwrap_or_default(),
            metadata: row.try_get("metadata").unwrap_or_default(),
            created_at: row.try_get("created_at").unwrap_or_default(),
            updated_at: row.try_get("updated_at").unwrap_or_default(),
            deleted_at: row.try_get("deleted_at").ok(),
        })
        .collect();

    Ok(memories)
}

pub async fn list_summaries(pool: &PgPool, scope_key: &str) -> anyhow::Result<Vec<SummaryRow>> {
    let rows = sqlx::query(
        r#"
        SELECT id, summary_type, content, started_at, ended_at, created_at
        FROM ias_context_summaries
        WHERE scope_key = $1
        ORDER BY created_at DESC, id DESC
        LIMIT 10
        "#,
    )
    .bind(scope_key)
    .fetch_all(pool)
    .await?;

    let summaries = rows
        .into_iter()
        .map(|row| SummaryRow {
            id: row.try_get("id").unwrap_or_default(),
            summary_type: row.try_get("summary_type").unwrap_or_default(),
            content: row.try_get("content").unwrap_or_default(),
            started_at: row.try_get("started_at").ok(),
            ended_at: row.try_get("ended_at").ok(),
            created_at: row.try_get("created_at").unwrap_or_default(),
        })
        .collect();

    Ok(summaries)
}

pub async fn load_tools() -> anyhow::Result<Vec<ToolRow>> {
    let builtin = load_builtin_tools();
    let capability = load_capability_tools().await;
    let mut all = builtin;
    all.extend(capability);
    Ok(all)
}

fn load_builtin_tools() -> Vec<ToolRow> {
    ias_ai::toolkit::get_tool()
        .into_iter()
        .filter_map(extract_builtin_tool)
        .collect()
}

fn extract_builtin_tool(value: Value) -> Option<ToolRow> {
    let function = value.get("function")?;
    let name = function.get("name")?.as_str()?.to_string();
    let description = function
        .get("description")
        .and_then(Value::as_str)
        .unwrap_or_default()
        .to_string();
    let parameters = function.get("parameters");
    let required = parameters
        .and_then(|v| v.get("required"))
        .and_then(Value::as_array)
        .map(|vals| {
            vals.iter()
                .filter_map(Value::as_str)
                .map(ToString::to_string)
                .collect::<Vec<_>>()
        })
        .unwrap_or_default();

    let mut params = parameters
        .and_then(|v| v.get("properties"))
        .and_then(Value::as_object)
        .map(|props| {
            let mut params: Vec<ToolParamRow> = props
                .iter()
                .map(|(param_name, spec)| ToolParamRow {
                    name: param_name.clone(),
                    param_type: spec
                        .get("type")
                        .and_then(Value::as_str)
                        .unwrap_or("object")
                        .to_string(),
                    required: required.iter().any(|r| r == param_name),
                    description: spec
                        .get("description")
                        .and_then(Value::as_str)
                        .unwrap_or_default()
                        .to_string(),
                })
                .collect();
            params.sort_by(|a, b| a.name.cmp(&b.name));
            params
        })
        .unwrap_or_default();
    params.sort_by_key(|p| (!p.required, p.name.clone()));

    Some(ToolRow {
        name,
        description,
        source: "内置".to_string(),
        path: None,
        params,
    })
}

async fn load_capability_tools() -> Vec<ToolRow> {
    let tools_path = configured_tools_path();
    let Ok(mut dir) = tokio::fs::read_dir(&tools_path).await else {
        return Vec::new();
    };

    let mut files = Vec::new();
    loop {
        match dir.next_entry().await {
            Ok(Some(entry)) => {
                let path = entry.path();
                if is_yaml_file(&path) {
                    files.push(path);
                }
            }
            Ok(None) => break,
            Err(_) => break,
        }
    }
    files.sort_by_key(|p| p.file_name().map(|n| n.to_os_string()));

    let mut tools = Vec::new();
    for file in files {
        let file_label = file
            .file_name()
            .and_then(|n| n.to_str())
            .unwrap_or_default()
            .to_string();

        let Ok(content) = tokio::fs::read_to_string(&file).await else {
            continue;
        };
        let Ok(config) = serde_yaml::from_str::<ias_ai::workflow::types::FlowConfig>(&content)
        else {
            continue;
        };

        tools.push(ToolRow {
            name: config.name,
            description: config.desc,
            source: "YAML".to_string(),
            path: Some(file_label),
            params: config
                .params
                .into_iter()
                .map(|p| ToolParamRow {
                    name: p.name,
                    param_type: p.param_type,
                    required: p.required,
                    description: p.description,
                })
                .collect(),
        });
    }

    tools
}

fn configured_tools_path() -> std::path::PathBuf {
    std::env::var("IA_TOOLS_PATH")
        .ok()
        .map(|v| v.trim().to_string())
        .filter(|v| !v.is_empty())
        .map(std::path::PathBuf::from)
        .unwrap_or_else(|| {
            std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
                .parent()
                .unwrap_or_else(|| std::path::Path::new("."))
                .join("tools")
        })
}

fn is_yaml_file(path: &std::path::Path) -> bool {
    path.extension()
        .and_then(|e| e.to_str())
        .map(|e| matches!(e.to_ascii_lowercase().as_str(), "yaml" | "yml"))
        .unwrap_or(false)
}

pub fn pretty_metadata(raw: &str) -> String {
    match serde_json::from_str::<Value>(raw) {
        Ok(Value::Object(obj)) => {
            let parts: Vec<String> = obj
                .iter()
                .filter_map(|(k, v)| {
                    if v.is_null() {
                        None
                    } else {
                        Some(format!("{}: {}", k, v))
                    }
                })
                .collect();
            if parts.is_empty() {
                String::new()
            } else {
                format!("{{ {} }}", parts.join(", "))
            }
        }
        _ => {
            let trimmed = raw.trim();
            if trimmed.len() > 200 {
                format!("{}...", trimmed.chars().take(200).collect::<String>())
            } else {
                trimmed.to_string()
            }
        }
    }
}