use anyhow::{Context, Result};
use serde::Deserialize;
use std::collections::{HashMap, HashSet};
use std::fs;
use std::path::PathBuf;

#[derive(Debug, Deserialize)]
struct WorkspaceToml {
	workspace: WorkspaceConfig,
}

#[derive(Debug, Deserialize)]
struct WorkspaceConfig {
	members: Vec<String>,
}

#[derive(Debug, Deserialize)]
struct CrateToml {
	package: PackageConfig,
	dependencies: Option<HashMap<String, CrateDependency>>,
}

#[derive(Debug, Deserialize)]
struct PackageConfig {
	name: String,
}

/// Represents a crate-level dependency in Cargo.toml
/// The Simple variant's String is needed for serde deserialization but never read directly
#[derive(Debug, Deserialize)]
#[serde(untagged)]
#[allow(dead_code)]
enum CrateDependency {
	Simple(String),
	Detailed {
		path: Option<String>,
		#[serde(flatten)]
		other: HashMap<String, toml::Value>,
	},
}

struct CrateInfo {
	name: String,
	path: PathBuf,
	dependencies: Vec<String>,
}

/// Remove transitive dependencies from the crate list.
/// If A depends on B and C, and B depends on C, then A->C is removed.
fn remove_transitive_dependencies(crates: &mut [CrateInfo]) {
	// Build a map from crate name to its dependencies for quick lookup
	let dep_map: HashMap<String, HashSet<String>> = crates.iter().map(|c| (c.name.clone(), c.dependencies.iter().cloned().collect())).collect();

	// For each crate, compute which dependencies are reachable through other dependencies
	for crate_info in crates.iter_mut() {
		let mut transitive_deps = HashSet::new();

		// For each direct dependency, find all its transitive dependencies
		for direct_dep in &crate_info.dependencies {
			// Recursively collect all transitive dependencies of this direct dependency
			let mut visited = HashSet::new();
			collect_all_dependencies(direct_dep, &dep_map, &mut visited);
			// Remove the direct dependency itself from the visited set
			visited.remove(direct_dep);
			transitive_deps.extend(visited);
		}

		// Remove dependencies that are transitive
		crate_info.dependencies.retain(|dep| !transitive_deps.contains(dep));
	}
}

/// Recursively collect all dependencies of a crate
fn collect_all_dependencies(crate_name: &str, dep_map: &HashMap<String, HashSet<String>>, visited: &mut HashSet<String>) {
	if !visited.insert(crate_name.to_string()) {
		return; // Already visited, avoid cycles
	}
	if let Some(deps) = dep_map.get(crate_name) {
		for dep in deps {
			collect_all_dependencies(dep, dep_map, visited);
		}
	}
}

fn main() -> Result<()> {
	let output_path = std::env::args_os()
		.nth(1)
		.map(PathBuf::from)
		.ok_or_else(|| anyhow::anyhow!("Usage: crate-hierarchy-viz <output-file>"))?;

	let workspace_root = std::env::current_dir().unwrap();
	let workspace_toml_path = workspace_root.join("Cargo.toml");

	// Parse workspace Cargo.toml
	let workspace_content = fs::read_to_string(&workspace_toml_path).with_context(|| format!("Failed to read {:?}", workspace_toml_path))?;
	let workspace_toml: WorkspaceToml = toml::from_str(&workspace_content).with_context(|| "Failed to parse workspace Cargo.toml")?;

	// Expand glob patterns in workspace members (e.g., "node-graph/libraries/*")
	let mut resolved_members = Vec::new();
	let mut seen_members = HashSet::new();
	let abs_root = workspace_root.canonicalize().unwrap_or_else(|_| workspace_root.clone());
	for member in &workspace_toml.workspace.members {
		if member.contains('*') {
			let pattern = abs_root.join(member).to_string_lossy().to_string();
			let matched: Vec<_> = glob::glob(&pattern)
				.with_context(|| format!("Failed to expand glob pattern: {member}"))?
				.filter_map(|entry| entry.ok())
				.filter_map(|path| path.strip_prefix(&abs_root).ok().map(|p| p.to_string_lossy().to_string()))
				.collect();
			if matched.is_empty() {
				eprintln!("Warning: No matches for glob pattern: {member}");
			}
			for m in matched {
				let normalized = m.replace('\\', "/");
				if seen_members.insert(normalized.clone()) {
					resolved_members.push(normalized);
				}
			}
		} else {
			let normalized = member.replace('\\', "/");
			if seen_members.insert(normalized.clone()) {
				resolved_members.push(normalized);
			}
		}
	}

	// Parse each member crate's Cargo.toml
	let mut parsed_crates = Vec::new();
	for member in &resolved_members {
		let crate_path = workspace_root.join(member);
		let cargo_toml_path = crate_path.join("Cargo.toml");

		if !cargo_toml_path.exists() {
			eprintln!("Warning: Cargo.toml not found for member: {}", member);
			continue;
		}

		let crate_content = fs::read_to_string(&cargo_toml_path).with_context(|| format!("Failed to read {:?}", cargo_toml_path))?;
		let crate_toml: CrateToml = toml::from_str(&crate_content).with_context(|| format!("Failed to parse Cargo.toml for {}", member))?;

		parsed_crates.push((crate_path, crate_toml));
	}

	// Collect all workspace crate names
	let workspace_crate_names: HashSet<String> = parsed_crates.iter().map(|(_, toml)| toml.package.name.clone()).collect();

	// Build dependency graph, keeping only workspace-internal dependencies
	let mut crates: Vec<CrateInfo> = parsed_crates
		.into_iter()
		.map(|(path, crate_toml)| {
			let dependencies = crate_toml
				.dependencies
				.unwrap_or_default()
				.into_iter()
				.filter_map(|(dep_name, dep_config)| {
					// Resolve the actual package name (handles renamed dependencies)
					let actual_name = match &dep_config {
						CrateDependency::Detailed { other, .. } => other.get("package").and_then(|v| v.as_str()).map(|s| s.to_string()).unwrap_or(dep_name),
						CrateDependency::Simple(_) => dep_name,
					};
					// Only keep dependencies that are workspace crates
					workspace_crate_names.contains(&actual_name).then_some(actual_name)
				})
				.collect();

			CrateInfo {
				name: crate_toml.package.name,
				path,
				dependencies,
			}
		})
		.collect();

	remove_transitive_dependencies(&mut crates);

	// Generate DOT format and write to output file
	let dot_content = generate_dot(&crates);

	if let Some(parent) = output_path.parent() {
		fs::create_dir_all(parent).with_context(|| format!("Failed to create directory {:?}", parent))?;
	}
	fs::write(&output_path, &dot_content).with_context(|| format!("Failed to write to {:?}", output_path))?;

	Ok(())
}

fn generate_dot(crates: &[CrateInfo]) -> String {
	let mut out = String::new();
	out.push_str("digraph CrateHierarchy {\n");
	out.push_str("    rankdir=LR;\n");
	out.push_str("    node [shape=box, style=\"rounded,filled\", fillcolor=lightblue];\n");
	out.push_str("    edge [color=gray];\n\n");

	// Define subgraph clusters
	let clusters: &[(&str, &str, &str, Box<dyn Fn(&CrateInfo) -> bool>)] = &[
		(
			"cluster_core",
			"Core Components",
			"lightgray",
			Box::new(|c| (c.name.starts_with("graphite-") || c.name == "editor" || c.name == "graphene-cli") && !c.name.contains("desktop")),
		),
		(
			"cluster_nodegraph",
			"Node Graph System",
			"lightyellow",
			Box::new(|c| c.name == "graph-craft" || c.name == "interpreted-executor" || c.name == "node-macro" || c.name == "preprocessor" || c.name == "graphene-cli"),
		),
		(
			"cluster_node_libraries",
			"Node Graph Libraries",
			"lightcyan",
			Box::new(|c| c.path.to_string_lossy().replace('\\', "/").contains("node-graph/libraries")),
		),
		(
			"cluster_nodes",
			"Nodes",
			"lightblue",
			Box::new(|c| c.path.to_string_lossy().replace('\\', "/").contains("node-graph/nodes")),
		),
		(
			"cluster_desktop",
			"Desktop",
			"lightgreen",
			Box::new(|c| c.path.to_string_lossy().replace('\\', "/").contains("desktop")),
		),
	];

	for (id, label, color, filter) in clusters {
		out.push_str(&format!("    subgraph {id} {{\n"));
		out.push_str(&format!("        label=\"{label}\";\n"));
		out.push_str("        style=filled;\n");
		out.push_str(&format!("        fillcolor={color};\n"));
		for c in crates.iter().filter(|c| filter(c)) {
			out.push_str(&format!("        \"{}\";\n", c.name));
		}
		out.push_str("    }\n\n");
	}

	// Add dependency edges
	for crate_info in crates {
		for dep in &crate_info.dependencies {
			if dep == "dyn-any" || dep == "node-macro" {
				continue;
			}
			out.push_str(&format!("    \"{}\" -> \"{}\";\n", crate_info.name, dep));
		}
	}

	out.push_str("}\n");
	out
}