dial9-viewer 0.3.4

use clap::{Parser, Subcommand};
use std::path::PathBuf;

mod skills {
    // Generated by build.rs from skills/ and toolkit/ directories
    include!(concat!(env!("OUT_DIR"), "/skill_files.rs"));
    include!(concat!(env!("OUT_DIR"), "/toolkit_files.rs"));

    /// Look up a skill segment by name. Returns the content if found.
    pub fn get(name: &str) -> Option<&'static str> {
        SKILL_FILES
            .iter()
            .find(|(n, _, _)| *n == name)
            .map(|(_, _, content)| *content)
    }
}

#[derive(Parser, Debug)]
#[command(
    name = "dial9-viewer",
    about = "Trace browser and viewer for dial9-tokio-telemetry"
)]
pub struct Cli {
    #[command(subcommand)]
    command: Option<Commands>,

    /// Port to listen on (when running without a subcommand)
    #[arg(long, default_value = "3000", global = true)]
    port: u16,

    /// S3 bucket name (when running without a subcommand)
    #[arg(long, global = true)]
    bucket: Option<String>,

    /// S3 key prefix (when running without a subcommand)
    #[arg(long, global = true)]
    prefix: Option<String>,

    /// Serve traces from a local directory instead of S3
    #[arg(long, global = true, conflicts_with = "bucket")]
    local_dir: Option<PathBuf>,

    /// Dev mode: serve UI files from disk for faster iteration
    #[arg(long, global = true)]
    dev: bool,
}

#[derive(Subcommand, Debug)]
enum Commands {
    /// Agent skill documentation and analysis toolkit
    Agents {
        #[command(subcommand)]
        action: Option<AgentsAction>,
    },
    /// Start the web server
    Serve {},
}

#[derive(Subcommand, Debug)]
enum AgentsAction {
    /// Copy the analysis toolkit (JS modules + analyze.js) to a directory
    Toolkit {
        /// Directory to write toolkit files into (created if missing)
        path: PathBuf,
    },
    /// Print a specific skill segment
    Skill {
        /// Segment name: loading, analysis, recipes, red-flags
        name: String,
    },
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let cli = Cli::parse();

    match cli.command {
        Some(Commands::Agents { action }) => match action {
            None => print!("{}", skills::HEADER),
            Some(AgentsAction::Toolkit { path }) => {
                std::fs::create_dir_all(&path)?;
                for (name, content) in skills::TOOLKIT_FILES {
                    std::fs::write(path.join(name), content)?;
                }
                let abs = std::fs::canonicalize(&path)?;
                eprintln!("Toolkit written to {}", abs.display());
                eprintln!("Run: node {}/analyze.js <trace.bin>", abs.display());
            }
            Some(AgentsAction::Skill { name }) => match skills::get(&name) {
                Some(content) => print!("{}", content),
                None => {
                    eprintln!("Unknown segment: {name}");
                    eprintln!("Available segments:");
                    for (name, title, _) in skills::SKILL_FILES {
                        eprintln!("  {name:12} — {title}");
                    }
                    std::process::exit(1);
                }
            },
        },
        Some(Commands::Serve {}) | None => {
            return serve(cli.port, cli.bucket, cli.prefix, cli.local_dir, cli.dev).await;
        }
    }
    Ok(())
}

async fn detect_bucket_region(bucket: &str) -> Option<String> {
    let config = aws_config::load_defaults(aws_config::BehaviorVersion::latest()).await;
    let client = aws_sdk_s3::Client::new(&config);
    match client.head_bucket().bucket(bucket).send().await {
        Ok(resp) => resp.bucket_region().map(|r| r.to_string()),
        Err(err) => {
            // HeadBucket errors include the x-amz-bucket-region header
            let raw = err.raw_response();
            raw.and_then(|r| {
                r.headers()
                    .get("x-amz-bucket-region")
                    .map(|v| v.to_string())
            })
        }
    }
}

async fn serve(
    port: u16,
    bucket: Option<String>,
    prefix: Option<String>,
    local_dir: Option<PathBuf>,
    dev: bool,
) -> anyhow::Result<()> {
    tracing_subscriber::fmt()
        .with_env_filter(
            tracing_subscriber::EnvFilter::try_from_default_env()
                .unwrap_or_else(|_| "dial9_viewer=info".parse().unwrap()),
        )
        .init();

    let dev_ui_dir = if dev {
        // In dev mode, find the ui/ directory relative to the manifest or CWD
        let candidates = [PathBuf::from("ui"), PathBuf::from("dial9-viewer/ui")];
        let dir = candidates.into_iter().find(|p| p.exists());
        match dir {
            Some(d) => {
                tracing::info!(path = %d.display(), "dev mode: serving UI from disk");
                Some(d)
            }
            None => {
                anyhow::bail!(
                    "--dev: could not find ui/ directory. Run from the dial9-viewer/ or repo root directory."
                );
            }
        }
    } else {
        None
    };

    let app_state = if let Some(dir) = &local_dir {
        let dir = std::fs::canonicalize(dir)?;
        tracing::info!(path = %dir.display(), "serving traces from local directory");
        let backend = dial9_viewer::storage::LocalBackend::new(&dir);
        let mut state = dial9_viewer::server::AppState::new(
            std::sync::Arc::new(backend),
            Some("local".into()),
            prefix.clone(),
        );
        if let Some(d) = dev_ui_dir {
            state = state.with_dev_ui_dir(d);
        }
        state
    } else {
        // Detect bucket region if a bucket is provided
        if let Some(bucket_name) = &bucket {
            if let Some(region) = detect_bucket_region(bucket_name).await {
                tracing::info!(%region, bucket = %bucket_name, "detected bucket region");
                let config = aws_config::defaults(aws_config::BehaviorVersion::latest())
                    .region(aws_sdk_s3::config::Region::new(region))
                    .load()
                    .await;
                let client = aws_sdk_s3::Client::new(&config);
                let backend = dial9_viewer::storage::S3Backend::from_client(client);
                let mut state = dial9_viewer::server::AppState::new(
                    std::sync::Arc::new(backend),
                    bucket.clone(),
                    prefix.clone(),
                );
                if let Some(d) = dev_ui_dir {
                    state = state.with_dev_ui_dir(d);
                }
                state
            } else {
                tracing::warn!(bucket = %bucket_name, "could not detect bucket region, using default");
                let backend = dial9_viewer::storage::S3Backend::from_env().await;
                let mut state = dial9_viewer::server::AppState::new(
                    std::sync::Arc::new(backend),
                    bucket.clone(),
                    prefix.clone(),
                );
                if let Some(d) = dev_ui_dir {
                    state = state.with_dev_ui_dir(d);
                }
                state
            }
        } else {
            let backend = dial9_viewer::storage::S3Backend::from_env().await;
            let mut state = dial9_viewer::server::AppState::new(
                std::sync::Arc::new(backend),
                bucket.clone(),
                prefix.clone(),
            );
            if let Some(d) = dev_ui_dir {
                state = state.with_dev_ui_dir(d);
            }
            state
        }
    };

    let app = dial9_viewer::server::router(app_state);

    let listener = tokio::net::TcpListener::bind(("0.0.0.0", port)).await?;
    tracing::info!(port, dev, "dial9-viewer listening");
    println!("\n  → http://localhost:{}\n", port);
    if let Some(dir) = &local_dir {
        tracing::info!(path = %dir.display(), "local directory mode");
    } else if let Some(bucket) = &bucket {
        tracing::info!(%bucket, "default bucket");
    }

    axum::serve(listener, app)
        .with_graceful_shutdown(shutdown_signal())
        .await?;

    Ok(())
}

async fn shutdown_signal() {
    tokio::signal::ctrl_c()
        .await
        .expect("failed to install CTRL+C handler");
    tracing::info!("shutting down");
}