BlueArchiveToolkit/infrastructure/src/bin/bat/translation_query.rs

use super::report_output::print_json_value;
use super::*;
use bat_core::domain::{validate_glossary_override, GlossaryOverride, TranslationMemoryContext};
use bat_core::repositories::TranslationMemoryRepository;
use std::collections::BTreeMap;

#[derive(Debug, Deserialize)]
struct TranslationTaskResultUpdateParam {
    unit_id: String,
    source_text: String,
    translated_text: String,
    #[serde(default)]
    glossary_override: Option<GlossaryOverride>,
}

pub(super) fn build_translation_tasks_report(
    state_dir: &Path,
    query: OfficialTextUnitTaskQuery,
    offset: usize,
    limit: usize,
) -> anyhow::Result<serde_json::Value> {
    let (_, version_state) = read_daemon_resource_state(state_dir)?;
    let current = version_state
        .as_ref()
        .and_then(|state| state.current_completed_version.as_ref());
    let Some(record) = current else {
        return Ok(serde_json::json!({ "available": false }));
    };
    let task_queue_path = record
        .resource_root
        .join(bat_infrastructure::OFFICIAL_TEXTUNIT_TASK_QUEUE_FILE);
    let task_repository_path =
        SqliteTranslationTaskRepository::repository_path(&record.resource_root);
    let Some(queue) = bat_infrastructure::read_textunit_task_queue_at(&record.resource_root)
        .map_err(anyhow::Error::msg)?
    else {
        return Ok(serde_json::json!({
            "available": false,
            "current_version_id": record.id,
            "resource_root": record.resource_root,
            "textunit_task_queue_path": task_queue_path,
            "task_repository_path": task_repository_path,
            "task_repository_available": false,
        }));
    };
    let task_repository_available =
        sqlite_file_exists_no_symlink(&task_repository_path, "翻译任务状态数据库")?;
    let (total_entries, entries) = if task_repository_available {
        query_translation_task_repository(&task_repository_path, &query, offset, limit)?
    } else {
        let mut queue_query = query.clone();
        queue_query.task_status = None;
        queue_query.has_failure_reason = None;
        let matches = bat_infrastructure::query_textunit_tasks(&queue, &queue_query);
        let persisted = matches
            .into_iter()
            .cloned()
            .map(|task| {
                bat_infrastructure::PersistedTranslationTask::from_queued_task(
                    task,
                    queue.generated_unix_seconds,
                )
            })
            .filter(|task| {
                query
                    .task_status
                    .as_ref()
                    .is_none_or(|status| task.task_status.as_str() == status)
            })
            .filter(|task| {
                query
                    .has_failure_reason
                    .is_none_or(|has_reason| task.failure_reason.is_some() == has_reason)
            })
            .collect::<Vec<_>>();
        let total_entries = persisted.len();
        let entries = persisted
            .into_iter()
            .skip(offset)
            .take(limit)
            .collect::<Vec<_>>();
        (total_entries as u64, entries)
    };
    Ok(serde_json::json!({
        "available": true,
        "current_version_id": record.id,
        "resource_root": record.resource_root,
        "textunit_task_queue_path": task_queue_path,
        "task_repository_path": task_repository_path,
        "task_repository_available": task_repository_available,
        "task_repository_schema_version": bat_infrastructure::TRANSLATION_TASK_SCHEMA_VERSION,
        "summary": queue.summary,
        "total_entries": total_entries,
        "offset": offset,
        "limit": limit,
        "query": translation_task_query_json(&query),
        "entries": entries,
    }))
}

pub(super) fn build_translation_handoff_report(
    state_dir: &Path,
) -> anyhow::Result<serde_json::Value> {
    let (_, version_state) = read_daemon_resource_state(state_dir)?;
    let current = version_state
        .as_ref()
        .and_then(|state| state.current_completed_version.as_ref());
    let Some(record) = current else {
        return Ok(serde_json::json!({ "available": false }));
    };
    let resource_root = &record.resource_root;
    let task_queue_path = resource_root.join(bat_infrastructure::OFFICIAL_TEXTUNIT_TASK_QUEUE_FILE);
    let handoff_path = resource_root.join(bat_infrastructure::TRANSLATION_HANDOFF_FILE);
    let repository_path = SqliteTranslationTaskRepository::repository_path(resource_root);
    let Some(queue) = bat_infrastructure::read_textunit_task_queue_at(resource_root)
        .map_err(anyhow::Error::msg)?
    else {
        return Ok(serde_json::json!({
            "available": false,
            "current_version_id": record.id,
            "resource_root": resource_root,
            "textunit_task_queue_path": task_queue_path,
            "translation_handoff_path": handoff_path,
            "task_repository_path": repository_path,
        }));
    };
    let task_repository_available =
        sqlite_file_exists_no_symlink(&repository_path, "翻译任务状态数据库")?;
    let tasks = if task_repository_available {
        let runtime = tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()?;
        runtime.block_on(async {
            let repository = SqliteTranslationTaskRepository::open(&repository_path)
                .await
                .map_err(|error| anyhow::anyhow!("{error}"))?;
            repository
                .list(&OfficialTextUnitTaskQuery::default())
                .await
                .map_err(|error| anyhow::anyhow!("{error}"))
        })?
    } else {
        queue
            .tasks
            .iter()
            .cloned()
            .map(|task| {
                bat_infrastructure::PersistedTranslationTask::from_queued_task(
                    task,
                    queue.generated_unix_seconds,
                )
            })
            .collect::<Vec<_>>()
    };
    let handoff = bat_infrastructure::build_translation_handoff(&queue, &tasks);
    let handoff_file_available = sqlite_file_exists_no_symlink(&handoff_path, "翻译 handoff")?;
    Ok(serde_json::json!({
        "available": true,
        "current_version_id": record.id,
        "resource_root": resource_root,
        "textunit_task_queue_path": task_queue_path,
        "translation_handoff_path": handoff_path,
        "translation_handoff_file_available": handoff_file_available,
        "task_repository_path": repository_path,
        "task_repository_available": task_repository_available,
        "task_repository_schema_version": bat_infrastructure::TRANSLATION_TASK_SCHEMA_VERSION,
        "handoff_schema_version": bat_infrastructure::TRANSLATION_HANDOFF_SCHEMA_VERSION,
        "handoff": handoff,
    }))
}

pub(super) fn update_translation_task_status_report(
    state_dir: &Path,
    params: Option<&serde_json::Value>,
    configured_glossary_path: Option<&Path>,
) -> anyhow::Result<serde_json::Value> {
    let task_id = rpc_string_param(params, "task_id")
        .ok_or_else(|| anyhow::anyhow!("translation.task.update 缺少 task_id"))?;
    let status_label = rpc_string_param(params, "status")
        .ok_or_else(|| anyhow::anyhow!("translation.task.update 缺少 status"))?;
    let status = TranslationTaskStatus::parse(status_label)
        .ok_or_else(|| anyhow::anyhow!("不支持的翻译任务 worker 状态：{status_label}"))?;
    let failure_reason = rpc_string_param(params, "failure_reason")
        .or_else(|| rpc_string_param(params, "reason"))
        .map(str::to_string);
    let provider_run_id = rpc_string_param(params, "provider_run_id").map(str::to_string);
    let provider = rpc_string_param(params, "provider").map(str::to_string);
    let result_params = translation_task_result_params(params)?;
    if !result_params.is_empty() && status != TranslationTaskStatus::Completed {
        return Err(anyhow::anyhow!(
            "translation_results 只能随 completed 状态写入"
        ));
    }
    let (_, version_state) = read_daemon_resource_state(state_dir)?;
    let current = version_state
        .as_ref()
        .and_then(|state| state.current_completed_version.as_ref())
        .ok_or_else(|| anyhow::anyhow!("没有可更新翻译任务的当前官方 release"))?;
    let repository_path = SqliteTranslationTaskRepository::repository_path(&current.resource_root);
    if !sqlite_file_exists_no_symlink(&repository_path, "翻译任务状态数据库")? {
        return Err(anyhow::anyhow!(
            "翻译任务状态数据库不存在：{}",
            repository_path.display()
        ));
    }
    let textunit_index = if result_params.is_empty() {
        None
    } else {
        Some(
            read_textunit_index_at(&current.resource_root)
                .map_err(anyhow::Error::msg)?
                .ok_or_else(|| {
                    anyhow::anyhow!("当前 release 缺少 TextUnit 明细索引，无法校验人工校对结果")
                })?,
        )
    };
    let result_provider = provider.clone().unwrap_or_else(|| "manual".to_string());
    let result_timestamp = unix_seconds_now();
    let result_provider_run_id = provider_run_id
        .clone()
        .unwrap_or_else(|| format!("manual-{result_timestamp}"));

    let runtime = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()?;
    let task = runtime.block_on(async {
        let repository = SqliteTranslationTaskRepository::open(&repository_path)
            .await
            .map_err(|error| anyhow::anyhow!("{error}"))?;
        if let Some(index) = textunit_index.as_ref() {
            let current_task = repository
                .find(task_id)
                .await
                .map_err(|error| anyhow::anyhow!("{error}"))?;
            let glossary_path = configured_glossary_path
                .map(PathBuf::from)
                .unwrap_or_else(|| {
                    bat_infrastructure::SqliteGlossaryRepository::repository_path(
                        &current.resource_root,
                    )
                });
            let glossary = if std::fs::symlink_metadata(&glossary_path).is_ok() {
                Some(
                    bat_infrastructure::SqliteGlossaryRepository::open(&glossary_path)
                        .await
                        .map_err(|error| anyhow::anyhow!("打开 Glossary 数据库失败：{error}"))?,
                )
            } else {
                None
            };
            let results = build_manual_translation_results(
                &current_task,
                index,
                &result_params,
                &result_provider,
                &result_provider_run_id,
                result_timestamp,
                glossary.as_ref(),
            )
            .await?;
            repository
                .update_status_with_results(
                    task_id,
                    status,
                    failure_reason,
                    Some(result_provider_run_id),
                    Some(result_provider),
                    Some(&results),
                )
                .await
                .map_err(|error| anyhow::anyhow!("{error}"))
        } else {
            repository
                .update_status(task_id, status, failure_reason, provider_run_id)
                .await
                .map_err(|error| anyhow::anyhow!("{error}"))
        }
    })?;
    Ok(serde_json::json!({
        "available": true,
        "current_version_id": current.id,
        "task_repository_path": repository_path,
        "entry": task,
    }))
}

pub(super) fn textunit_query_json(query: &OfficialTextUnitQuery) -> serde_json::Value {
    serde_json::json!({
        "destination": query.destination.clone(),
        "path_pattern": query.path_pattern.clone(),
        "archive_entry": query.archive_entry.clone(),
        "path_id": query.path_id,
        "class_id": query.class_id,
        "field_path": query.field_path.clone(),
        "format": query.format.clone(),
    })
}

fn translation_task_result_params(
    params: Option<&serde_json::Value>,
) -> anyhow::Result<Vec<TranslationTaskResultUpdateParam>> {
    let Some(value) = params
        .and_then(|params| params.get("translation_results"))
        .or_else(|| params.and_then(|params| params.get("results")))
    else {
        return Ok(Vec::new());
    };
    if value.is_null() {
        return Ok(Vec::new());
    }
    serde_json::from_value(value.clone())
        .map_err(|error| anyhow::anyhow!("translation_results 必须是结果数组：{error}"))
}

async fn build_manual_translation_results(
    task: &bat_infrastructure::PersistedTranslationTask,
    index: &bat_infrastructure::OfficialTextUnitIndex,
    params: &[TranslationTaskResultUpdateParam],
    provider: &str,
    provider_run_id: &str,
    translated_unix_seconds: u64,
    glossary: Option<&bat_infrastructure::SqliteGlossaryRepository>,
) -> anyhow::Result<Vec<bat_infrastructure::TranslationTaskUnitResult>> {
    let index_by_id = index
        .units
        .iter()
        .map(|unit| (unit.id.as_str(), unit))
        .collect::<std::collections::BTreeMap<_, _>>();
    let mut seen = std::collections::BTreeSet::new();
    let mut results = Vec::with_capacity(params.len());
    for param in params {
        let unit_id = param.unit_id.trim();
        if unit_id.is_empty() {
            return Err(anyhow::anyhow!("translation_results[].unit_id 不能为空"));
        }
        if !seen.insert(unit_id.to_string()) {
            return Err(anyhow::anyhow!(
                "translation_results 包含重复 TextUnit：{unit_id}"
            ));
        }
        let unit = index_by_id
            .get(unit_id)
            .ok_or_else(|| anyhow::anyhow!("translation_results 引用了未知 TextUnit：{unit_id}"))?;
        if unit.destination != task.task.destination
            || unit.archive_entry != task.task.archive_entry
        {
            return Err(anyhow::anyhow!(
                "TextUnit {unit_id} 不属于翻译任务 {}",
                task.task.task_id
            ));
        }
        if param.source_text != unit.source_text {
            return Err(anyhow::anyhow!(
                "TextUnit {unit_id} 的 source_text 与当前索引不一致"
            ));
        }
        let glossary_qa = if let Some(glossary) = glossary {
            let context = bat_infrastructure::translation_memory_context(
                &unit.destination,
                unit.archive_entry.as_deref(),
                unit.serialized_file.as_deref(),
                unit.path_id,
                unit.class_id,
                unit.field_path.as_deref(),
                unit.format.as_deref(),
                unit.asset_name.as_deref(),
                unit.text_source_kind.as_deref(),
                &unit.context,
            );
            Some(
                glossary
                    .diagnose(&unit.source_text, &context)
                    .await
                    .map_err(|error| anyhow::anyhow!("Glossary QA 失败：{error}"))?
                    .check_translation(&param.translated_text),
            )
        } else {
            None
        };
        if let Some(qa) = glossary_qa.as_ref() {
            if qa.status.is_blocked() {
                validate_glossary_override(qa, param.glossary_override.as_ref()).map_err(
                    |error| {
                        anyhow::anyhow!("TextUnit {} 的 glossary_override 无效：{error}", unit_id)
                    },
                )?;
            } else if param.glossary_override.is_some() {
                return Err(anyhow::anyhow!(
                    "TextUnit {} 不能为非 blocking Glossary QA 指定 override",
                    unit_id
                ));
            }
        } else if param.glossary_override.is_some() {
            return Err(anyhow::anyhow!(
                "TextUnit {} 不能为非 blocking Glossary QA 指定 override",
                unit_id
            ));
        }
        results.push(bat_infrastructure::TranslationTaskUnitResult {
            unit_id: unit_id.to_string(),
            source_text: param.source_text.clone(),
            translated_text: param.translated_text.clone(),
            source_kind: bat_infrastructure::TranslationTaskResultSourceKind::Manual,
            translation_memory_record_id: None,
            provider: provider.to_string(),
            provider_run_id: provider_run_id.to_string(),
            translated_unix_seconds,
            glossary_qa,
            glossary_override: param.glossary_override.clone(),
        });
    }
    Ok(results)
}

pub(super) fn translation_task_query_json(query: &OfficialTextUnitTaskQuery) -> serde_json::Value {
    serde_json::json!({
        "task_id": query.task_id.clone(),
        "official_release_id": query.official_release_id.clone(),
        "destination": query.destination.clone(),
        "path_pattern": query.path_pattern.clone(),
        "archive_entry": query.archive_entry.clone(),
        "status": query.status.clone(),
        "task_status": query.task_status.clone(),
        "parse_status": query.parse_status.clone(),
        "text_unit_format": query.text_unit_format.clone(),
        "has_reason": query.has_reason,
        "has_failure_reason": query.has_failure_reason,
    })
}

pub(super) fn run_translation_memory_command(options: &CliOptions) -> anyhow::Result<()> {
    let method = match options.command {
        CliCommand::TranslationMemorySummary => RPC_METHOD_TRANSLATION_MEMORY_SUMMARY,
        CliCommand::TranslationMemoryQuery => RPC_METHOD_TRANSLATION_MEMORY_QUERY,
        CliCommand::TranslationMemoryConfirm => RPC_METHOD_TRANSLATION_MEMORY_CONFIRM,
        _ => return Err(anyhow::anyhow!("不是 Translation Memory 命令")),
    };
    if daemon_rpc_available(&options.state_dir)
        && options.resource_root.is_none()
        && !options.output_explicit
    {
        let _control_lock = DaemonControlLock::acquire(&options.state_dir)?;
        let report = daemon_rpc_call(
            &options.state_dir,
            method,
            translation_memory_cli_params(options)?,
        )?;
        print_json_value(options.output_format, &report)?;
        return Ok(());
    }
    let path = translation_memory_cli_path(options)?;
    let report = match options.command {
        CliCommand::TranslationMemorySummary => build_translation_memory_summary_report(&path)?,
        CliCommand::TranslationMemoryQuery => {
            let source_text = options
                .translation_memory_source_text
                .as_deref()
                .ok_or_else(|| anyhow::anyhow!("TM query 必须指定 --tm-source-text"))?;
            let context = parse_translation_memory_context(
                options.translation_memory_context_json.as_deref(),
            )?;
            build_translation_memory_query_report(
                &path,
                source_text,
                &context,
                options.query_limit,
            )?
        }
        CliCommand::TranslationMemoryConfirm => {
            let record_id = options
                .translation_memory_record_id
                .as_deref()
                .ok_or_else(|| anyhow::anyhow!("TM confirm 必须指定 --tm-record-id"))?;
            let reviewer = options
                .translation_memory_reviewer
                .as_deref()
                .ok_or_else(|| anyhow::anyhow!("TM confirm 必须指定 --tm-reviewer"))?;
            build_translation_memory_confirm_report(
                &path,
                record_id,
                reviewer,
                options.translation_memory_reason.clone(),
            )?
        }
        _ => unreachable!(),
    };
    print_json_value(options.output_format, &report)
}

fn translation_memory_cli_path(options: &CliOptions) -> anyhow::Result<std::path::PathBuf> {
    if let Some(path) = options.translation_memory_path.as_ref() {
        return lexical_absolute(path).map_err(anyhow::Error::msg);
    }
    let resource_root = options
        .resource_root
        .as_deref()
        .map(lexical_absolute)
        .transpose()
        .map_err(anyhow::Error::msg)?
        .unwrap_or(active_official_resource_root(&options.config.output_root)?);
    Ok(bat_infrastructure::translation_memory_repository_path(
        &resource_root,
    ))
}

fn translation_memory_cli_params(
    options: &CliOptions,
) -> anyhow::Result<Option<serde_json::Value>> {
    let mut params = serde_json::Map::new();
    if let Some(path) = options.translation_memory_path.as_ref() {
        params.insert(
            "translation_memory_path".to_string(),
            serde_json::json!(path),
        );
    }
    match options.command {
        CliCommand::TranslationMemorySummary => {}
        CliCommand::TranslationMemoryQuery => {
            let source_text = options
                .translation_memory_source_text
                .as_deref()
                .ok_or_else(|| anyhow::anyhow!("TM query 必须指定 --tm-source-text"))?;
            let context = parse_translation_memory_context(
                options.translation_memory_context_json.as_deref(),
            )?;
            params.insert("source_text".to_string(), serde_json::json!(source_text));
            params.insert("source_context".to_string(), serde_json::json!(context));
            params.insert("limit".to_string(), serde_json::json!(options.query_limit));
        }
        CliCommand::TranslationMemoryConfirm => {
            let record_id = options
                .translation_memory_record_id
                .as_deref()
                .ok_or_else(|| anyhow::anyhow!("TM confirm 必须指定 --tm-record-id"))?;
            let reviewer = options
                .translation_memory_reviewer
                .as_deref()
                .ok_or_else(|| anyhow::anyhow!("TM confirm 必须指定 --tm-reviewer"))?;
            params.insert("record_id".to_string(), serde_json::json!(record_id));
            params.insert("reviewer".to_string(), serde_json::json!(reviewer));
            if let Some(reason) = options.translation_memory_reason.as_deref() {
                params.insert("reason".to_string(), serde_json::json!(reason));
            }
        }
        _ => unreachable!(),
    }
    Ok(Some(serde_json::Value::Object(params)))
}

pub(super) fn build_translation_memory_summary_report(
    path: &std::path::Path,
) -> anyhow::Result<serde_json::Value> {
    if !sqlite_file_exists_no_symlink(path, "Translation Memory 数据库")? {
        return Ok(serde_json::json!({
            "available": false,
            "path": path,
            "reason": "database_missing",
        }));
    }
    let runtime = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()?;
    let summary = runtime.block_on(async {
        let repository = bat_infrastructure::SqliteTranslationMemoryRepository::open(path)
            .await
            .map_err(|error| anyhow::anyhow!("{error}"))?;
        repository
            .summary()
            .await
            .map_err(|error| anyhow::anyhow!("{error}"))
    })?;
    Ok(serde_json::json!({
        "available": true,
        "path": path,
        "schema_version": summary.schema_version,
        "summary": summary,
    }))
}

pub(super) fn build_translation_memory_query_report(
    path: &std::path::Path,
    source_text: &str,
    source_context: &TranslationMemoryContext,
    limit: usize,
) -> anyhow::Result<serde_json::Value> {
    if source_text.trim().is_empty() {
        return Err(anyhow::anyhow!("TM query 的 source_text 不能为空"));
    }
    if !(1..=1000).contains(&limit) {
        return Err(anyhow::anyhow!("TM query 的 limit 必须在 1..=1000 范围内"));
    }
    if !sqlite_file_exists_no_symlink(path, "Translation Memory 数据库")? {
        return Ok(serde_json::json!({
            "available": false,
            "path": path,
            "source_text": source_text,
            "source_context": source_context,
            "matches": [],
            "reason": "database_missing",
        }));
    }
    let runtime = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()?;
    let matches = runtime.block_on(async {
        let repository = bat_infrastructure::SqliteTranslationMemoryRepository::open(path)
            .await
            .map_err(|error| anyhow::anyhow!("{error}"))?;
        repository
            .find_matches(source_text, source_context, limit)
            .await
            .map_err(|error| anyhow::anyhow!("{error}"))
    })?;
    Ok(serde_json::json!({
        "available": true,
        "path": path,
        "source_text": source_text,
        "source_context": source_context,
        "matches": matches,
    }))
}

pub(super) fn build_translation_memory_confirm_report(
    path: &std::path::Path,
    record_id: &str,
    reviewer: &str,
    reason: Option<String>,
) -> anyhow::Result<serde_json::Value> {
    if record_id.trim().is_empty() || reviewer.trim().is_empty() {
        return Err(anyhow::anyhow!(
            "TM confirm 必须指定非空 record_id 和 reviewer"
        ));
    }
    if !sqlite_file_exists_no_symlink(path, "Translation Memory 数据库")? {
        return Err(anyhow::anyhow!(
            "Translation Memory 数据库不存在：{}",
            path.display()
        ));
    }
    let runtime = tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()?;
    let entry = runtime.block_on(async {
        let repository = bat_infrastructure::SqliteTranslationMemoryRepository::open(path)
            .await
            .map_err(|error| anyhow::anyhow!("{error}"))?;
        repository
            .confirm(record_id, reviewer, reason)
            .await
            .map_err(|error| anyhow::anyhow!("{error}"))
    })?;
    Ok(serde_json::json!({
        "available": true,
        "path": path,
        "entry": entry,
    }))
}

fn parse_translation_memory_context(
    value: Option<&str>,
) -> anyhow::Result<TranslationMemoryContext> {
    let Some(value) = value else {
        return Ok(BTreeMap::new());
    };
    serde_json::from_str::<TranslationMemoryContext>(value)
        .map_err(|error| anyhow::anyhow!("--tm-context-json 必须是 JSON object：{error}"))
}