tori/src/template.rs

use std::path::Path;

use serde::Deserialize;

use crate::llm::{ChatMessage, LlmClient};
use crate::tools::ExternalToolManager;

#[derive(Debug, Deserialize)]
pub struct TemplateInfo {
    pub name: String,
    pub description: String,
    pub match_hint: String,
}

#[allow(dead_code)]
pub struct LoadedTemplate {
    pub id: String,
    pub info: TemplateInfo,
    pub instructions: String,
    pub external_tools: ExternalToolManager,
    pub kb_files: Vec<(String, String)>,
}

pub fn templates_dir() -> &'static str {
    if Path::new("/app/templates").is_dir() {
        "/app/templates"
    } else {
        "app-templates"
    }
}

/// Scan available templates and ask LLM to pick one (or none).
pub async fn select_template(llm: &LlmClient, requirement: &str) -> Option<String> {
    let base = Path::new(templates_dir());
    let mut entries = match tokio::fs::read_dir(base).await {
        Ok(e) => e,
        Err(_) => return None,
    };

    let mut templates: Vec<(String, TemplateInfo)> = Vec::new();
    while let Ok(Some(entry)) = entries.next_entry().await {
        // Use metadata() instead of file_type() to follow symlinks
        let is_dir = tokio::fs::metadata(entry.path())
            .await
            .map(|m| m.is_dir())
            .unwrap_or(false);
        if !is_dir {
            continue;
        }
        let id = entry.file_name().to_string_lossy().to_string();
        let meta_path = entry.path().join("template.json");
        if let Ok(data) = tokio::fs::read_to_string(&meta_path).await {
            if let Ok(info) = serde_json::from_str::<TemplateInfo>(&data) {
                templates.push((id, info));
            }
        }
    }

    if templates.is_empty() {
        return None;
    }

    let listing: String = templates
        .iter()
        .map(|(id, info)| {
            format!(
                "- id: {}\n  名称: {}\n  描述: {}\n  适用场景: {}",
                id, info.name, info.description, info.match_hint
            )
        })
        .collect::<Vec<_>>()
        .join("\n");

    let prompt = format!(
        "以下是可用的项目模板：\n{}\n\n用户需求：{}\n\n选择最匹配的模板 ID，如果都不合适则回复 none。只回复模板 ID 或 none，不要其他内容。",
        listing, requirement
    );

    let response = llm
        .chat(vec![
            ChatMessage::system("你是一个模板选择助手。根据用户需求选择最合适的项目模板。只回复模板 ID 或 none。"),
            ChatMessage::user(&prompt),
        ])
        .await
        .ok()?;

    let answer = response.trim().to_lowercase();
    tracing::info!("Template selection LLM response: '{}' (available: {:?})",
        answer, templates.iter().map(|(id, _)| id.as_str()).collect::<Vec<_>>());
    if answer == "none" {
        return None;
    }

    // Verify the answer matches an actual template ID
    let result = templates
        .iter()
        .find(|(id, _)| id == &answer)
        .map(|(id, _)| id.clone());
    if result.is_none() {
        tracing::warn!("Template selection: LLM returned '{}' which doesn't match any template ID", answer);
    }
    result
}

/// Copy template contents to workdir (excluding template.json, tools/, kb/).
pub async fn apply_template(template_id: &str, workdir: &str) -> anyhow::Result<()> {
    let src = Path::new(templates_dir()).join(template_id);
    if !src.is_dir() {
        anyhow::bail!("Template directory not found: {}", template_id);
    }
    copy_dir_recursive(&src, Path::new(workdir)).await
}

/// Recursively copy directory contents, skipping template.json/tools/kb at the top level.
async fn copy_dir_recursive(src: &Path, dst: &Path) -> anyhow::Result<()> {
    let mut stack: Vec<(std::path::PathBuf, std::path::PathBuf, bool)> =
        vec![(src.to_path_buf(), dst.to_path_buf(), true)];

    while let Some((src_dir, dst_dir, top_level)) = stack.pop() {
        tokio::fs::create_dir_all(&dst_dir).await?;
        let mut entries = tokio::fs::read_dir(&src_dir).await?;
        while let Some(entry) = entries.next_entry().await? {
            let name = entry.file_name();
            let name_str = name.to_string_lossy();

            // Skip metadata/convention dirs at top level
            if top_level && (name_str == "template.json" || name_str == "tools" || name_str == "kb")
            {
                continue;
            }

            let src_path = entry.path();
            let dst_path = dst_dir.join(&name);

            if entry.file_type().await?.is_dir() {
                stack.push((src_path, dst_path, false));
            } else {
                tokio::fs::copy(&src_path, &dst_path).await?;
            }
        }
    }
    Ok(())
}

impl LoadedTemplate {
    /// Load a template: discover tools, read KB files, read instructions.
    pub async fn load(template_id: &str) -> anyhow::Result<Self> {
        let base = Path::new(templates_dir()).join(template_id);
        if !base.is_dir() {
            anyhow::bail!("Template directory not found: {}", template_id);
        }

        // Read template.json
        let meta_path = base.join("template.json");
        let meta_data = tokio::fs::read_to_string(&meta_path).await?;
        let info: TemplateInfo = serde_json::from_str(&meta_data)?;

        // Read INSTRUCTIONS.md
        let instructions_path = base.join("INSTRUCTIONS.md");
        let instructions = tokio::fs::read_to_string(&instructions_path)
            .await
            .unwrap_or_default();

        // Discover external tools
        let tools_dir = base.join("tools");
        let external_tools = ExternalToolManager::discover(&tools_dir).await;
        tracing::info!(
            "Template '{}': discovered {} external tools",
            template_id,
            external_tools.len()
        );

        // Scan KB files
        let kb_dir = base.join("kb");
        let kb_files = scan_kb_files(&kb_dir).await;
        tracing::info!(
            "Template '{}': found {} KB files",
            template_id,
            kb_files.len()
        );

        Ok(Self {
            id: template_id.to_string(),
            info,
            instructions,
            external_tools,
            kb_files,
        })
    }
}

/// Scan kb/ directory for .md files. Returns (title, content) pairs.
/// Title is extracted from the first `# heading` line, or falls back to the filename.
async fn scan_kb_files(kb_dir: &Path) -> Vec<(String, String)> {
    let mut results = Vec::new();

    let mut entries = match tokio::fs::read_dir(kb_dir).await {
        Ok(e) => e,
        Err(_) => return results,
    };

    while let Ok(Some(entry)) = entries.next_entry().await {
        let path = entry.path();

        // Resolve symlinks
        let real_path = match tokio::fs::canonicalize(&path).await {
            Ok(p) => p,
            Err(_) => path.clone(),
        };

        // Only process .md files
        let ext = real_path
            .extension()
            .and_then(|e| e.to_str())
            .unwrap_or("");
        if ext != "md" {
            continue;
        }

        let content = match tokio::fs::read_to_string(&real_path).await {
            Ok(c) => c,
            Err(e) => {
                tracing::warn!("Failed to read KB file {}: {}", real_path.display(), e);
                continue;
            }
        };

        // Extract title: first `# heading` line, or filename without extension
        let title = content
            .lines()
            .find(|l| l.starts_with("# "))
            .map(|l| l.trim_start_matches("# ").trim().to_string())
            .unwrap_or_else(|| {
                path.file_stem()
                    .and_then(|s| s.to_str())
                    .unwrap_or("untitled")
                    .to_string()
            });

        results.push((title, content));
    }

    results
}