htknow/src/weknora_sync.rs

use std::{collections::HashMap, time::Duration};

use anyhow::{Context, Result, anyhow};
use log::{info, warn};
use once_cell::sync::Lazy;
use reqwest::Client;
use serde::Serialize;
use sha2::{Digest, Sha256};
use sqlx::{FromRow, SqlitePool};

use crate::api::{FILE_COLS_NO_CONTENT, File};

static CLIENT: Lazy<Client> =
    Lazy::new(|| Client::builder().timeout(Duration::from_secs(120)).build().expect("build WeKnora sync HTTP client"));

#[derive(Debug, Serialize)]
struct ExternalChunk {
    external_slice_id: String,
    index: usize,
    content: String,
}

#[derive(Debug, Serialize)]
struct ExternalImport {
    source: &'static str,
    external_file_id: String,
    source_version: String,
    title: String,
    filename: String,
    file_type: String,
    file_size: i64,
    file_hash: String,
    chunks: Vec<ExternalChunk>,
}

#[derive(Debug, FromRow)]
struct SyncEvent {
    event_id: String,
    event_type: String,
    file_id: i64,
    htknow_kb_id: i64,
    attempts: i64,
}

pub async fn enqueue_file_upsert(pool: &SqlitePool, file_id: i64) -> Result<()> {
    if !sync_enabled() {
        return Ok(());
    }
    let kb_id: i64 = sqlx::query_scalar("SELECT kb_id FROM files WHERE id = ?")
        .bind(file_id)
        .fetch_optional(pool)
        .await?
        .flatten()
        .ok_or_else(|| anyhow!("file {} has no knowledge base", file_id))?;
    enqueue(pool, "upsert", file_id, kb_id).await
}

pub async fn enqueue_file_delete(pool: &SqlitePool, file_id: i64, htknow_kb_id: i64) -> Result<()> {
    if !sync_enabled() {
        return Ok(());
    }
    enqueue(pool, "delete", file_id, htknow_kb_id).await
}

async fn enqueue(pool: &SqlitePool, event_type: &str, file_id: i64, htknow_kb_id: i64) -> Result<()> {
    let event_id = uuid::Uuid::new_v4().to_string();
    sqlx::query(
        "INSERT INTO weknora_sync_outbox
         (event_id, event_type, file_id, htknow_kb_id, attempts, next_attempt_at, last_error, created_at, updated_at)
         VALUES (?, ?, ?, ?, 0, strftime('%s','now'), '', strftime('%s','now'), strftime('%s','now'))
         ON CONFLICT(file_id) DO UPDATE SET
           event_id = excluded.event_id,
           event_type = excluded.event_type,
           htknow_kb_id = excluded.htknow_kb_id,
           attempts = 0,
           next_attempt_at = excluded.next_attempt_at,
           last_error = '',
           updated_at = excluded.updated_at",
    )
    .bind(event_id)
    .bind(event_type)
    .bind(file_id)
    .bind(htknow_kb_id)
    .execute(pool)
    .await?;
    Ok(())
}

pub fn start_worker(pool: SqlitePool) {
    if !sync_enabled() {
        info!("WeKnora synchronization is disabled");
        return;
    }
    tokio::spawn(async move {
        loop {
            if let Err(err) = process_due_events(&pool).await {
                warn!("WeKnora outbox worker failed: {}", err);
            }
            tokio::time::sleep(Duration::from_secs(2)).await;
        }
    });
}

async fn process_due_events(pool: &SqlitePool) -> Result<()> {
    let events: Vec<SyncEvent> = sqlx::query_as(
        "SELECT event_id, event_type, file_id, htknow_kb_id, attempts
         FROM weknora_sync_outbox
         WHERE next_attempt_at <= strftime('%s','now')
         ORDER BY created_at, file_id
         LIMIT 10",
    )
    .fetch_all(pool)
    .await?;
    for event in events {
        let result = match event.event_type.as_str() {
            "upsert" => sync_file(pool, event.file_id).await,
            "delete" => delete_file(event.file_id, event.htknow_kb_id).await,
            other => Err(anyhow!("unsupported outbox event type {}", other)),
        };
        match result {
            Ok(()) => {
                sqlx::query("DELETE FROM weknora_sync_outbox WHERE file_id = ? AND event_id = ?")
                    .bind(event.file_id)
                    .bind(&event.event_id)
                    .execute(pool)
                    .await?;
                info!("WeKnora {} accepted for file {}", event.event_type, event.file_id);
            }
            Err(err) => {
                let attempts = event.attempts + 1;
                let delay = (2_i64.saturating_pow(attempts.min(10) as u32)).min(3600);
                sqlx::query(
                    "UPDATE weknora_sync_outbox
                     SET attempts = ?, next_attempt_at = strftime('%s','now') + ?, last_error = ?, updated_at = strftime('%s','now')
                     WHERE file_id = ? AND event_id = ?",
                )
                .bind(attempts)
                .bind(delay)
                .bind(err.to_string())
                .bind(event.file_id)
                .bind(&event.event_id)
                .execute(pool)
                .await?;
                warn!(
                    "WeKnora {} attempt {} failed for file {}; retry in {}s: {}",
                    event.event_type, attempts, event.file_id, delay, err
                );
            }
        }
    }
    Ok(())
}

fn sync_enabled() -> bool {
    std::env::var("HTKNOW_WEKNORA_URL").ok().is_some_and(|v| !v.trim().is_empty())
}

fn mapped_weknora_kb_id(htknow_kb_id: i64) -> Result<String> {
    let kb_map_raw = std::env::var("HTKNOW_WEKNORA_KB_MAP").context("HTKNOW_WEKNORA_KB_MAP is required")?;
    let kb_map: HashMap<String, String> =
        serde_json::from_str(&kb_map_raw).context("parse HTKNOW_WEKNORA_KB_MAP JSON")?;
    kb_map
        .get(&htknow_kb_id.to_string())
        .or_else(|| kb_map.get("*"))
        .cloned()
        .ok_or_else(|| anyhow!("no WeKnora KB mapping for htknow KB {} and no '*' default mapping", htknow_kb_id))
}

async fn sync_file(pool: &SqlitePool, file_id: i64) -> Result<()> {
    let file: File = sqlx::query_as(&format!("SELECT {} FROM files WHERE id = ?", FILE_COLS_NO_CONTENT))
        .bind(file_id)
        .fetch_one(pool)
        .await
        .context("load parsed file")?;
    if file.status != 1 {
        return Err(anyhow!("file status is {}, expected completed", file.status));
    }
    let htknow_kb_id = file.kb_id.ok_or_else(|| anyhow!("file has no knowledge base"))?;
    let weknora_kb_id = mapped_weknora_kb_id(htknow_kb_id)?;

    let source_id = crate::api::effective_parse_file_id(pool, file_id).await?;
    let ids: Vec<i64> = sqlx::query_scalar("SELECT id FROM slices WHERE file_id = ? ORDER BY ordinal, id")
        .bind(source_id)
        .fetch_all(pool)
        .await?;
    let contents = crate::slice_content::read_all(source_id).await?;
    let chunks = ids
        .into_iter()
        .enumerate()
        .filter_map(|(index, id)| {
            let content = contents.get(&id)?.trim();
            if content.is_empty() {
                return None;
            }
            Some(ExternalChunk { external_slice_id: id.to_string(), index, content: content.to_string() })
        })
        .collect::<Vec<_>>();
    if chunks.is_empty() {
        return Err(anyhow!("parsed file contains no text chunks"));
    }

    let mut version_hasher = Sha256::new();
    version_hasher.update(file.filename.as_bytes());
    version_hasher.update(file.hash.as_bytes());
    for chunk in &chunks {
        version_hasher.update(chunk.external_slice_id.as_bytes());
        version_hasher.update([0]);
        version_hasher.update(chunk.content.as_bytes());
        version_hasher.update([0xff]);
    }
    let source_version = hex::encode(version_hasher.finalize());
    let file_type =
        std::path::Path::new(&file.filename).extension().and_then(|v| v.to_str()).unwrap_or("markdown").to_lowercase();
    let payload = ExternalImport {
        source: "htknow",
        external_file_id: file.id.to_string(),
        source_version,
        title: file.filename.clone(),
        filename: file.filename,
        file_type,
        file_size: file.size,
        file_hash: file.hash,
        chunks,
    };
    let base = std::env::var("HTKNOW_WEKNORA_URL")?;
    let url =
        format!("{}/api/v1/knowledge-bases/{}/knowledge/external-chunks", base.trim_end_matches('/'), weknora_kb_id);
    let mut request = CLIENT.post(url).json(&payload);
    if let Ok(key) = std::env::var("HTKNOW_WEKNORA_API_KEY")
        && !key.trim().is_empty()
    {
        request = request.header("X-API-Key", key);
    }
    let response = request.send().await.context("call WeKnora external chunk import")?;
    let status = response.status();
    if !status.is_success() {
        let body = response.text().await.unwrap_or_default();
        return Err(anyhow!("WeKnora returned {}: {}", status, body));
    }
    Ok(())
}

async fn delete_file(file_id: i64, htknow_kb_id: i64) -> Result<()> {
    let weknora_kb_id = mapped_weknora_kb_id(htknow_kb_id)?;
    let base = std::env::var("HTKNOW_WEKNORA_URL")?;
    let url = format!(
        "{}/api/v1/knowledge-bases/{}/knowledge/external-chunks/htknow/{}",
        base.trim_end_matches('/'),
        weknora_kb_id,
        file_id
    );
    let mut request = CLIENT.delete(url);
    if let Ok(key) = std::env::var("HTKNOW_WEKNORA_API_KEY")
        && !key.trim().is_empty()
    {
        request = request.header("X-API-Key", key);
    }
    let response = request.send().await.context("call WeKnora external knowledge delete")?;
    let status = response.status();
    if !status.is_success() {
        let body = response.text().await.unwrap_or_default();
        return Err(anyhow!("WeKnora returned {}: {}", status, body));
    }
    Ok(())
}