depbank 0.2.7

use anyhow::{Context, Result};
use depbank::{
    DependencyCollection, calculate_directory_tokens, calculate_file_tokens, collect_dependencies,
    extract_dependency_info, find_cargo_lock, find_cargo_toml_files, generate_all_code_banks,
    is_dependency_available, resolve_dependency_versions, resolve_registry_path,
};
use std::collections::{HashMap, HashSet};
use std::fmt::Write as _;
use std::fs;
use std::path::{Path, PathBuf};

// Constants for formatting strings
const README_HEADER: &str = "# Code Bank Summary\n\n";
const README_TABLE_HEADER: &str = "| Dependency | Version | Tokens | Size (bytes) |\n";
const README_TABLE_SEPARATOR: &str = "|------------|---------|--------|-------------|\n";
const README_TOKEN_SUMMARY_HEADER: &str = "\n## Token Summary\n\n";
const README_ABOUT_HEADER: &str = "\n## About Code Banks\n\n";
const README_ABOUT_P1: &str = "Code banks are generated summaries of your project's dependencies. ";
const README_ABOUT_P2: &str = "They provide an overview of the code structure and key components ";
const README_ABOUT_P3: &str =
    "of each dependency, helping you understand the libraries your project uses ";
const README_ABOUT_P4: &str = "without having to browse through the original source code.\n\n";
const README_ABOUT_P5: &str =
    "Each file contains a summary of the corresponding dependency's source code, ";
const README_ABOUT_P6: &str = "including important types, functions, and structures.\n\n";
const README_ABOUT_P7: &str = "Generated by [DepBank](https://github.com/tyrchen/depbank).\n";

pub fn generate_command(project_path: &Path, output_dir: &Path, dry_run: bool) -> Result<()> {
    println!("Analyzing project...");

    // Find and analyze dependencies
    let available_deps = analyze_dependencies(project_path, dry_run)?;

    if dry_run {
        println!("Dry run enabled, skipping generation");
        return Ok(());
    }

    // Generate code banks
    println!("Generating code banks...");
    let registry_path = resolve_registry_path()?;
    let code_bank_files = generate_all_code_banks(&available_deps, &registry_path, output_dir)?;
    println!("Generated {} code bank files", code_bank_files.len());

    // Calculate tokens and generate README
    generate_code_bank_readme(
        output_dir,
        project_path,
        &available_deps,
        code_bank_files.len(),
    )?;

    Ok(())
}

fn analyze_dependencies(project_path: &Path, _dry_run: bool) -> Result<DependencyCollection> {
    // Find all Cargo.toml files
    let cargo_toml_files = find_cargo_toml_files(project_path)?;
    println!("Found {} Cargo.toml files", cargo_toml_files.len());

    if cargo_toml_files.is_empty() {
        return Err(anyhow::anyhow!("No Cargo.toml files found"));
    }

    // Extract dependencies from *all* found Cargo.toml files
    let mut dependency_info = DependencyCollection::new();
    let mut unique_deps_for_reporting = HashSet::new(); // Track unique dep names for reporting

    for cargo_toml_path in &cargo_toml_files {
        let file_deps = extract_dependency_info(cargo_toml_path)?;
        for dep in file_deps.iter() {
            // Add to the main collection
            dependency_info.add(dep.clone());
            // Track unique names for reporting count
            unique_deps_for_reporting.insert(dep.name.clone());
        }
    }

    println!(
        "Found {} unique dependencies",
        unique_deps_for_reporting.len()
    );

    // Find Cargo.lock - should still be at the workspace root
    let cargo_lock_path = find_cargo_lock(project_path)?;
    println!("Found Cargo.lock");

    // Resolve exact versions from Cargo.lock using the aggregated dependency info
    let resolved_versions = resolve_dependency_versions(cargo_lock_path, &dependency_info)?;
    println!("Resolved {} versions", resolved_versions.len());

    // Resolve registry path
    let registry_path = resolve_registry_path()?;

    // Check which dependencies are available locally
    let mut available_deps = DependencyCollection::new();
    // Keep track of names we've already added to available_deps to avoid duplicates if
    // the same dep name appears in multiple Cargo.toml files after version resolution.
    let mut added_names = HashSet::new();
    for dependency in resolved_versions.iter() {
        if !added_names.contains(&dependency.name)
            && is_dependency_available(&registry_path, dependency)
        {
            available_deps.add(dependency.clone());
            added_names.insert(dependency.name.clone());
        }
    }

    println!(
        "{}/{} unique dependencies available locally",
        available_deps.len(),
        unique_deps_for_reporting.len() // Report against unique names found initially
    );

    if available_deps.is_empty() {
        // Adjust the error message to be more accurate in a workspace context
        return Err(anyhow::anyhow!(
            "No resolved dependencies available locally in the Cargo registry"
        ));
    }

    // Return the available dependencies
    Ok(available_deps)
}

fn generate_code_bank_readme(
    output_dir: &Path,
    project_path: &Path,
    dependencies: &DependencyCollection,
    code_bank_files_count: usize,
) -> Result<()> {
    println!("Calculating tokens for generated code banks (may take a while)...");
    // Calculate tokens for generated code banks
    let file_stats = calculate_directory_tokens(output_dir, Some("md"))?;

    // Sort stats by token count
    let mut stats_vec: Vec<_> = file_stats.iter().collect();
    stats_vec.sort_by(|a, b| b.1.token_count.cmp(&a.1.token_count)); // Sort by token count (descending)

    // Generate README content
    let (readme_content, total_tokens) = create_readme_content(
        project_path,
        code_bank_files_count,
        &stats_vec,
        dependencies,
        &file_stats,
    );

    // Write README.md to the output directory
    let readme_path = output_dir.join("README.md");
    fs::write(&readme_path, readme_content).with_context(|| {
        format!(
            "Failed to write README.md to file: {}",
            readme_path.display()
        )
    })?;

    // Print concise summary to console
    println!("\nSummary:");
    println!("- Generated {} code bank files", code_bank_files_count);
    println!("- Total tokens: {}", total_tokens);
    println!("- Added README.md with summary and token information");
    println!("- Output directory: {}", output_dir.display());

    Ok(())
}

fn create_readme_content(
    project_path: &Path,
    code_bank_files_count: usize,
    stats_vec: &[(&String, &depbank::FileStats)],
    dependencies: &DependencyCollection,
    file_stats: &HashMap<String, depbank::FileStats>,
) -> (String, usize) {
    let mut total_tokens = 0;
    let mut total_size = 0;
    let mut readme_content = String::new();

    // Add header
    readme_content.push_str(README_HEADER);
    write!(
        readme_content,
        "Generated for project: {}\n\n",
        project_path.display()
    )
    .unwrap();

    // Add dependency section
    write!(
        readme_content,
        "## Dependencies ({} total)\n\n",
        code_bank_files_count
    )
    .unwrap();

    // Add dependency list with versions
    readme_content.push_str(README_TABLE_HEADER);
    readme_content.push_str(README_TABLE_SEPARATOR);

    for (name, stats) in stats_vec {
        let name_without_md = name.trim_end_matches(".md");
        let version = dependencies
            .get_version(name_without_md)
            .map(|v| v.as_str())
            .unwrap_or("unknown");

        writeln!(
            readme_content,
            "| {} | {} | {} | {} |",
            name_without_md, version, stats.token_count, stats.size_bytes
        )
        .unwrap();

        total_tokens += stats.token_count;
        total_size += stats.size_bytes;
    }

    // Add token summary
    readme_content.push_str(README_TOKEN_SUMMARY_HEADER);
    writeln!(readme_content, "- **Total tokens:** {}", total_tokens).unwrap();
    writeln!(readme_content, "- **Total size:** {} bytes", total_size).unwrap();
    writeln!(
        readme_content,
        "- **Number of files:** {}",
        file_stats.len()
    )
    .unwrap();

    // Add explanation of what code banks are
    readme_content.push_str(README_ABOUT_HEADER);
    readme_content.push_str(README_ABOUT_P1);
    readme_content.push_str(README_ABOUT_P2);
    readme_content.push_str(README_ABOUT_P3);
    readme_content.push_str(README_ABOUT_P4);
    readme_content.push_str(README_ABOUT_P5);
    readme_content.push_str(README_ABOUT_P6);
    readme_content.push_str(README_ABOUT_P7);

    (readme_content, total_tokens)
}

pub fn tokens_command(path: &Path, extension: Option<&str>) -> Result<()> {
    if path.is_file() {
        analyze_file_tokens(path)?;
    } else if path.is_dir() {
        analyze_directory_tokens(path, extension)?;
    } else {
        return Err(anyhow::anyhow!(
            "Path does not exist or is not accessible: {}",
            path.display()
        ));
    }

    Ok(())
}

fn analyze_file_tokens(path: &Path) -> Result<()> {
    // Calculate tokens for a single file
    let token_count = calculate_file_tokens(path)?;
    let file_size = std::fs::metadata(path)?.len();
    println!(
        "{}: {} tokens, {} bytes",
        path.display(),
        token_count,
        file_size
    );
    Ok(())
}

fn analyze_directory_tokens(dir_path: &Path, extension: Option<&str>) -> Result<()> {
    // Calculate tokens for all files in the directory
    let file_stats = calculate_directory_tokens(dir_path, extension)?;

    // Print token counts in a sorted manner
    let mut stats_vec: Vec<_> = file_stats.iter().collect();
    stats_vec.sort_by(|a, b| b.1.token_count.cmp(&a.1.token_count)); // Sort by token count (descending)

    let (total_tokens, total_size) = print_token_stats(dir_path, &stats_vec, file_stats.len());

    println!(
        "\nTotal: {} tokens, {} bytes across {} files",
        total_tokens,
        total_size,
        file_stats.len()
    );

    Ok(())
}

fn print_token_stats(
    dir_path: &Path,
    stats_vec: &[(&String, &depbank::FileStats)],
    _file_count: usize,
) -> (usize, usize) {
    let mut total_tokens = 0;
    let mut total_size = 0;

    println!("Token counts for files in {}:", dir_path.display());
    for (name, stats) in stats_vec {
        println!(
            "{}: {} tokens, {} bytes",
            name, stats.token_count, stats.size_bytes
        );

        total_tokens += stats.token_count;
        total_size += stats.size_bytes;
    }

    (total_tokens, total_size)
}

pub fn list_command(project_path: &Path, detailed: bool) -> Result<()> {
    // Find all Cargo.toml files
    let cargo_toml_files = find_cargo_toml_files(project_path)?;
    println!("Found {} Cargo.toml files", cargo_toml_files.len());

    if cargo_toml_files.is_empty() {
        return Err(anyhow::anyhow!("No Cargo.toml files found"));
    }

    // Collect all dependencies
    let dependencies = collect_dependencies(&cargo_toml_files)?;
    println!("\nFound {} unique dependencies:", dependencies.len());

    if detailed {
        display_detailed_dependency_info(project_path, &cargo_toml_files)?;
    } else {
        display_simple_dependency_list(&dependencies);
    }

    Ok(())
}

fn display_simple_dependency_list(dependencies: &HashSet<String>) {
    // Sort dependencies for consistent output
    let mut sorted_deps: Vec<_> = dependencies.iter().collect();
    sorted_deps.sort();

    // For simple view, just list dependencies
    for dep in sorted_deps {
        println!("- {}", dep);
    }
}

fn display_detailed_dependency_info(
    project_path: &Path,
    cargo_toml_files: &[PathBuf],
) -> Result<()> {
    // For detailed view, show dependency info from each Cargo.toml
    display_dependency_specs_by_file(cargo_toml_files)?;

    // Try to resolve versions from Cargo.lock if available
    display_cargo_lock_versions(project_path, cargo_toml_files)?;

    Ok(())
}

fn display_dependency_specs_by_file(cargo_toml_files: &[PathBuf]) -> Result<()> {
    for (index, cargo_toml) in cargo_toml_files.iter().enumerate() {
        println!("\nDependency specifications from {}:", cargo_toml.display());

        let dependency_info = extract_dependency_info(cargo_toml)?;

        // Sort dependencies for consistent output
        let mut sorted_info: Vec<_> = dependency_info.iter().collect();
        sorted_info.sort_by(|a, b| a.version.cmp(&b.version));

        for dep in sorted_info {
            println!("{}: {}", dep.name, dep.version);
        }

        // If not the last Cargo.toml, add a separator
        if index < cargo_toml_files.len() - 1 {
            println!("\n---");
        }
    }

    Ok(())
}

fn display_cargo_lock_versions(project_path: &Path, cargo_toml_files: &[PathBuf]) -> Result<()> {
    if let Ok(cargo_lock_path) = find_cargo_lock(project_path) {
        println!("\nFound Cargo.lock at: {}", cargo_lock_path.display());

        // Extract dependencies from the first Cargo.toml for resolution
        let first_cargo_toml = &cargo_toml_files[0];
        let dependency_info = extract_dependency_info(first_cargo_toml)?;

        // Resolve exact versions from Cargo.lock
        if let Ok(resolved_versions) =
            resolve_dependency_versions(cargo_lock_path, &dependency_info)
        {
            println!("\nResolved dependency versions from Cargo.lock:");

            // Sort dependencies for consistent output
            let mut sorted_resolved: Vec<_> = resolved_versions.iter().collect();
            sorted_resolved.sort_by(|a, b| a.version.cmp(&b.version));

            for dep in sorted_resolved {
                println!("{}: {}", dep.name, dep.version);
            }
        }
    }

    Ok(())
}