knack-registry 0.1.0

use std::{
    collections::BTreeMap,
    fs::File,
    net::SocketAddr,
    path::{Path, PathBuf},
    process::Command as ProcessCommand,
    sync::Arc,
    time::Duration,
};

use anyhow::{Context, Result, bail};
use axum::{
    Json, Router,
    body::Body,
    extract::{Path as AxumPath, Query, State},
    http::{StatusCode, header},
    response::{IntoResponse, Response},
    routing::get,
};
use clap::Parser;
use flate2::{Compression, write::GzEncoder};
use serde::Deserialize;
use knack_core::{IndexedSkill, RegistryIndex, collect_files, read_skill, validate_skill};
use tar::{Builder, Header};
use tokio::sync::RwLock;

#[derive(Debug, Parser)]
#[command(name = "knack-registry")]
#[command(version, about = "Serve and search a knack registry index")]
struct Cli {
    /// Path to a knack registry index TOML file.
    #[arg(long, default_value = "knack.index.toml")]
    index: PathBuf,

    /// Address to bind.
    #[arg(long, default_value = "127.0.0.1:7349")]
    bind: SocketAddr,

    /// Optional local root containing skill directories to serve as archives.
    #[arg(long)]
    skills_root: Option<PathBuf>,

    /// Optional registry alias to return as install sources, e.g. company.
    #[arg(long)]
    public_alias: Option<String>,

    /// Source alias used to resolve backing sources, e.g. tea=git+ssh://git@gitea.example.com.
    #[arg(long = "source-alias")]
    source_aliases: Vec<String>,

    /// Periodically refresh dynamic sources. Set to 0 to disable background refresh.
    #[arg(long, default_value_t = 300)]
    refresh_interval_seconds: u64,
}

#[derive(Clone)]
struct AppState {
    index: Arc<RwLock<RegistryIndex>>,
    index_path: PathBuf,
    skills_root: Option<PathBuf>,
    public_alias: Option<String>,
    source_aliases: BTreeMap<String, String>,
}

#[derive(Debug, Deserialize)]
struct SearchParams {
    q: String,
}

#[tokio::main]
async fn main() -> Result<()> {
    let cli = Cli::parse();
    let source_aliases = parse_source_aliases(&cli.source_aliases)?;
    let index = refresh_index(&cli.index, &source_aliases)?;
    let state = AppState {
        index: Arc::new(RwLock::new(index)),
        index_path: cli.index,
        skills_root: cli.skills_root,
        public_alias: cli.public_alias,
        source_aliases,
    };

    if cli.refresh_interval_seconds > 0 {
        spawn_refresh_task(
            state.index.clone(),
            state.index_path.clone(),
            state.source_aliases.clone(),
            Duration::from_secs(cli.refresh_interval_seconds),
        );
    }

    let app = Router::new()
        .route("/health", get(health))
        .route("/index", get(get_index))
        .route("/search", get(search))
        .route("/skills/{name}/archive", get(skill_archive))
        .with_state(state);

    let listener = tokio::net::TcpListener::bind(cli.bind)
        .await
        .with_context(|| format!("failed to bind {}", cli.bind))?;
    println!("knack-registry listening on http://{}", cli.bind);
    axum::serve(listener, app)
        .with_graceful_shutdown(shutdown_signal())
        .await
        .context("registry server failed")?;

    Ok(())
}

fn read_index(path: &PathBuf) -> Result<RegistryIndex> {
    let contents = std::fs::read_to_string(path)
        .with_context(|| format!("failed to read {}", path.display()))?;
    let index: RegistryIndex =
        toml::from_str(&contents).with_context(|| format!("failed to parse {}", path.display()))?;
    index.validate()?;
    Ok(index)
}

fn refresh_index(
    path: &PathBuf,
    source_aliases: &BTreeMap<String, String>,
) -> Result<RegistryIndex> {
    let mut index = read_index(path)?;
    materialize_dynamic_sources(&mut index, source_aliases)?;
    Ok(index)
}

fn spawn_refresh_task(
    index: Arc<RwLock<RegistryIndex>>,
    index_path: PathBuf,
    source_aliases: BTreeMap<String, String>,
    interval: Duration,
) {
    tokio::spawn(async move {
        let mut ticker = tokio::time::interval(interval);
        ticker.tick().await;
        loop {
            ticker.tick().await;
            match refresh_index(&index_path, &source_aliases) {
                Ok(refreshed) => {
                    let mut index = index.write().await;
                    *index = refreshed;
                    eprintln!("refreshed knack registry index");
                }
                Err(error) => {
                    eprintln!("failed to refresh knack registry index: {error:#}");
                }
            }
        }
    });
}

fn materialize_dynamic_sources(
    index: &mut RegistryIndex,
    source_aliases: &BTreeMap<String, String>,
) -> Result<()> {
    let static_skill_names: Vec<String> =
        index.skill.iter().map(|skill| skill.name.clone()).collect();
    let dynamic_sources = index.source.clone();
    for source in dynamic_sources {
        let fetched = fetch_source_root(&source.source, source_aliases)?;
        for skill_dir in collect_skill_dirs(&fetched.path)? {
            let skill = read_skill(&skill_dir)?;
            validate_skill(&skill)?;
            if static_skill_names.iter().any(|name| name == &skill.name)
                || index.skill.iter().any(|indexed| indexed.name == skill.name)
            {
                continue;
            }
            let relative = skill_dir.strip_prefix(&fetched.path).with_context(|| {
                format!(
                    "failed to make {} relative to {}",
                    skill_dir.display(),
                    fetched.path.display()
                )
            })?;
            let relative = relative.to_string_lossy().replace('\\', "/");
            let skill_source = if relative.is_empty() {
                source.source.clone()
            } else {
                format!("{}/{}", source.source.trim_end_matches('/'), relative)
            };
            index.skill.push(IndexedSkill {
                name: skill.name,
                description: skill.description,
                source: skill_source,
                tags: source.tags.clone(),
            });
        }
    }
    index
        .skill
        .sort_by(|left, right| left.name.cmp(&right.name));
    index.validate()?;
    Ok(())
}

fn collect_skill_dirs(root: &Path) -> Result<Vec<PathBuf>> {
    let mut skills = Vec::new();
    collect_skill_dirs_inner(root, &mut skills)?;
    skills.sort();
    Ok(skills)
}

fn collect_skill_dirs_inner(path: &Path, skills: &mut Vec<PathBuf>) -> Result<()> {
    if path.join("SKILL.md").is_file() {
        skills.push(path.to_path_buf());
        return Ok(());
    }

    for entry in
        std::fs::read_dir(path).with_context(|| format!("failed to read {}", path.display()))?
    {
        let entry = entry?;
        let path = entry.path();
        let file_type = entry.file_type()?;
        if file_type.is_dir() && !is_ignored_scan_dir(&path) {
            collect_skill_dirs_inner(&path, skills)?;
        }
    }

    Ok(())
}

fn is_ignored_scan_dir(path: &Path) -> bool {
    path.file_name()
        .and_then(|name| name.to_str())
        .is_some_and(|name| matches!(name, ".git" | "target" | "node_modules"))
}

async fn health() -> &'static str {
    "ok"
}

async fn get_index(State(state): State<AppState>) -> Json<RegistryIndex> {
    Json(state.index.read().await.clone())
}

async fn search(
    State(state): State<AppState>,
    Query(params): Query<SearchParams>,
) -> Json<Vec<IndexedSkill>> {
    let index = state.index.read().await;
    let mut results: Vec<IndexedSkill> = index.search(&params.q).into_iter().cloned().collect();
    drop(index);
    if let Some(alias) = &state.public_alias {
        for skill in &mut results {
            skill.source = format!("{}:{}", alias, skill.name);
        }
    }
    Json(results)
}

async fn skill_archive(
    State(state): State<AppState>,
    AxumPath(name): AxumPath<String>,
) -> Response {
    match create_skill_archive(&state, &name).await {
        Ok(bytes) => (
            [
                (header::CONTENT_TYPE, "application/gzip"),
                (
                    header::CONTENT_DISPOSITION,
                    &format!("attachment; filename=\"{}.skill.tar.gz\"", name),
                ),
            ],
            Body::from(bytes),
        )
            .into_response(),
        Err(error) => (StatusCode::NOT_FOUND, error.to_string()).into_response(),
    }
}

async fn create_skill_archive(state: &AppState, name: &str) -> Result<Vec<u8>> {
    if let Some(skills_root) = &state.skills_root {
        let skill_dir = skills_root.join(name);
        if skill_dir.join("SKILL.md").is_file() {
            return create_skill_archive_from_dir(&skill_dir);
        }
    }

    let index = state.index.read().await;
    let source = index
        .skill
        .iter()
        .find(|skill| skill.name == name)
        .map(|skill| skill.source.clone())
        .with_context(|| format!("skill not found: {name}"))?;
    drop(index);
    let fetched = fetch_backing_source(&source, state)?;
    create_skill_archive_from_dir(&fetched.path)
}

fn create_skill_archive_from_dir(skill_dir: &Path) -> Result<Vec<u8>> {
    let skill = read_skill(skill_dir)?;
    validate_skill(&skill)?;

    let buffer = Vec::new();
    let encoder = GzEncoder::new(buffer, Compression::default());
    let mut archive = Builder::new(encoder);
    for file in collect_files(&skill_dir)? {
        let relative = file.strip_prefix(&skill_dir).with_context(|| {
            format!(
                "failed to make {} relative to {}",
                file.display(),
                skill_dir.display()
            )
        })?;
        let archive_name = Path::new(&skill.name).join(relative);
        append_file(&mut archive, &file, &archive_name)?;
    }
    archive.finish()?;
    let encoder = archive.into_inner()?;
    Ok(encoder.finish()?)
}

#[derive(Debug)]
struct FetchedBackingSource {
    path: PathBuf,
    _temp_dir: tempfile::TempDir,
}

fn fetch_backing_source(source: &str, state: &AppState) -> Result<FetchedBackingSource> {
    fetch_source_root(source, &state.source_aliases).and_then(|fetched| {
        let skill = read_skill(&fetched.path)?;
        validate_skill(&skill)?;
        Ok(fetched)
    })
}

fn fetch_source_root(
    source: &str,
    source_aliases: &BTreeMap<String, String>,
) -> Result<FetchedBackingSource> {
    let (alias, rest) = source
        .split_once(':')
        .ok_or_else(|| anyhow::anyhow!("backing source must be alias:owner/repo[@ref]/path"))?;
    let base_url = source_aliases
        .get(alias)
        .with_context(|| format!("source alias not configured on registry: {alias}"))?;
    let git = parse_git_host_source(base_url, rest)?;

    let temp_dir = tempfile::tempdir().context("failed to create temporary directory")?;
    let repo_dir = temp_dir.path().join("repo");
    let status = ProcessCommand::new("git")
        .arg("clone")
        .arg("--depth")
        .arg("1")
        .arg("--branch")
        .arg(&git.reference)
        .arg(&git.repo_url)
        .arg(&repo_dir)
        .status()
        .with_context(|| "failed to run git clone; is git installed?")?;

    if !status.success() {
        bail!(
            "git clone failed for backing source {} at ref {}",
            git.repo_url,
            git.reference
        );
    }

    let skill_dir = repo_dir.join(git.skill_path);
    Ok(FetchedBackingSource {
        path: skill_dir,
        _temp_dir: temp_dir,
    })
}

#[derive(Debug)]
struct GitBackingSource {
    repo_url: String,
    reference: String,
    skill_path: PathBuf,
}

fn parse_git_host_source(base_url: &str, rest: &str) -> Result<GitBackingSource> {
    let mut parts = rest.splitn(3, '/');
    let owner = parts
        .next()
        .filter(|part| !part.is_empty())
        .context("backing source must include owner")?;
    let repo_with_ref = parts
        .next()
        .filter(|part| !part.is_empty())
        .context("backing source must include repository")?;
    let skill_path = parts.next().unwrap_or("");
    let (repo, reference) = split_repo_ref(repo_with_ref, "main")?;
    let base_url = base_url
        .trim_end_matches('/')
        .strip_prefix("git+")
        .unwrap_or(base_url.trim_end_matches('/'));

    Ok(GitBackingSource {
        repo_url: format!("{base_url}/{owner}/{repo}.git"),
        reference: reference.to_string(),
        skill_path: PathBuf::from(skill_path),
    })
}

fn split_repo_ref<'a>(repo_with_ref: &'a str, default_ref: &'a str) -> Result<(&'a str, &'a str)> {
    let Some(position) = repo_with_ref.rfind('@') else {
        return Ok((repo_with_ref, default_ref));
    };
    let (repo, reference_with_at) = repo_with_ref.split_at(position);
    let reference = &reference_with_at[1..];
    if repo.is_empty() || reference.is_empty() {
        bail!("repository and ref must not be empty");
    }
    Ok((repo, reference))
}

fn parse_source_aliases(values: &[String]) -> Result<BTreeMap<String, String>> {
    let mut aliases = BTreeMap::new();
    for value in values {
        let (name, url) = value
            .split_once('=')
            .with_context(|| format!("source alias must be name=url: {value}"))?;
        if name.is_empty() || url.is_empty() {
            bail!("source alias name and url must not be empty: {value}");
        }
        aliases.insert(name.to_string(), url.to_string());
    }
    Ok(aliases)
}

fn append_file(
    archive: &mut Builder<GzEncoder<Vec<u8>>>,
    source: &Path,
    archive_name: &Path,
) -> Result<()> {
    let mut file =
        File::open(source).with_context(|| format!("failed to open {}", source.display()))?;
    let metadata = file
        .metadata()
        .with_context(|| format!("failed to stat {}", source.display()))?;
    if !metadata.is_file() {
        bail!("not a file: {}", source.display());
    }

    let mut header = Header::new_gnu();
    header.set_size(metadata.len());
    header.set_mode(0o644);
    header.set_mtime(0);
    header.set_uid(0);
    header.set_gid(0);
    header.set_cksum();

    archive
        .append_data(&mut header, archive_name, &mut file)
        .with_context(|| format!("failed to archive {}", source.display()))?;
    Ok(())
}

async fn shutdown_signal() {
    let _ = tokio::signal::ctrl_c().await;
}