use miette::Diagnostic;
use mlf_codegen::plugin::GeneratorContext;
use std::path::{Path, PathBuf};
use thiserror::Error;

#[derive(Error, Debug, Diagnostic)]
pub enum GenerateError {
    #[error("Failed to parse lexicon: {path}")]
    #[diagnostic(code(mlf::generate::parse_lexicon))]
    ParseLexicon {
        path: String,
        #[help]
        help: Option<String>,
    },

    #[error("Failed to write output: {path}")]
    #[diagnostic(code(mlf::generate::write_output))]
    WriteOutput {
        path: String,
        #[source]
        source: std::io::Error,
    },

    #[error("Generator '{name}' not found")]
    #[diagnostic(code(mlf::generate::generator_not_found))]
    #[help("Available generators: {}", available.join(", "))]
    GeneratorNotFound {
        name: String,
        available: Vec<String>,
    },

    #[error("Code generation failed: {0}")]
    #[diagnostic(code(mlf::generate::generation_failed))]
    #[allow(dead_code)]
    GenerationFailed(String),
}

pub fn run(
    generator_name: Option<String>,
    input_paths: Vec<PathBuf>,
    output_dir: Option<PathBuf>,
    root: Option<PathBuf>,
    flat: bool,
) -> Result<(), GenerateError> {
    let current_dir = std::env::current_dir().map_err(|source| GenerateError::WriteOutput {
        path: "current directory".to_string(),
        source,
    })?;

    // Load mlf.toml if available
    let project_root = crate::config::find_project_root(&current_dir).ok();
    let config = project_root.as_ref().and_then(|root| {
        let config_path = root.join("mlf.toml");
        crate::config::MlfConfig::load(&config_path).ok()
    });

    // Determine generator name
    let generator_name = if let Some(explicit) = generator_name {
        explicit
    } else if let Some(cfg) = &config {
        // Find first non-lexicon, non-mlf output in mlf.toml
        cfg.output
            .iter()
            .find(|o| o.r#type != "lexicon" && o.r#type != "mlf")
            .map(|o| o.r#type.clone())
            .ok_or_else(|| GenerateError::GeneratorNotFound {
                name: "any".to_string(),
                available: vec!["No code generator outputs configured in mlf.toml. Either add an output configuration or provide --generator flag.".to_string()],
            })?
    } else {
        return Err(GenerateError::GeneratorNotFound {
            name: "any".to_string(),
            available: vec!["No mlf.toml found and no --generator flag provided. Either create a mlf.toml or provide --generator flag.".to_string()],
        });
    };

    // Find the generator
    let generators = mlf_codegen::plugin::generators();
    let generator = generators
        .iter()
        .find(|g| g.name() == generator_name)
        .ok_or_else(|| {
            let available: Vec<String> = generators.iter().map(|g| g.name().to_string()).collect();
            GenerateError::GeneratorNotFound {
                name: generator_name.clone(),
                available,
            }
        })?;

    println!(
        "Using generator: {} ({})",
        generator.name(),
        generator.description()
    );
    println!("Output extension: {}\n", generator.file_extension());

    // Determine output directory
    let output_dir = if let Some(explicit) = output_dir {
        explicit
    } else if let Some(cfg) = &config {
        // Find output matching the generator type
        cfg.output
            .iter()
            .find(|o| o.r#type == generator_name)
            .map(|o| PathBuf::from(&o.directory))
            .ok_or_else(|| GenerateError::WriteOutput {
                path: "mlf.toml".to_string(),
                source: std::io::Error::new(
                    std::io::ErrorKind::NotFound,
                    format!(
                        "No output configured for generator '{}' in mlf.toml",
                        generator_name
                    ),
                ),
            })?
    } else {
        return Err(GenerateError::WriteOutput {
            path: "mlf.toml".to_string(),
            source: std::io::Error::new(
                std::io::ErrorKind::NotFound,
                "No mlf.toml found and no --output flag provided",
            ),
        });
    };

    // Determine root directory
    let root_dir = if let Some(explicit) = root {
        explicit
    } else if let Some(cfg) = &config {
        project_root.as_ref().unwrap().join(&cfg.source.directory)
    } else {
        current_dir.clone()
    };

    // Determine input paths
    let input_paths = if input_paths.is_empty() {
        if let Some(cfg) = &config {
            vec![project_root.as_ref().unwrap().join(&cfg.source.directory)]
        } else {
            return Err(GenerateError::WriteOutput {
                path: "input".to_string(),
                source: std::io::Error::new(
                    std::io::ErrorKind::NotFound,
                    "No input files specified and no mlf.toml found",
                ),
            });
        }
    } else {
        input_paths
    };

    // Collect input files
    let mut file_paths = Vec::new();
    for path in input_paths {
        if path.is_dir() {
            file_paths.extend(collect_mlf_files(&path)?);
        } else if path.is_file() && path.extension().and_then(|s| s.to_str()) == Some("mlf") {
            file_paths.push(path);
        }
    }

    std::fs::create_dir_all(&output_dir).map_err(|source| GenerateError::WriteOutput {
        path: output_dir.display().to_string(),
        source,
    })?;

    let mut errors = Vec::new();
    let mut success_count = 0;

    for file_path in file_paths {
        let source = match std::fs::read_to_string(&file_path) {
            Ok(s) => s,
            Err(e) => {
                errors.push((
                    file_path.display().to_string(),
                    format!("Failed to read file: {}", e),
                ));
                continue;
            }
        };

        let lexicon = match mlf_lang::parse_lexicon(&source) {
            Ok(lex) => lex,
            Err(e) => {
                errors.push((file_path.display().to_string(), format!("{:?}", e)));
                continue;
            }
        };

        let namespace = match extract_namespace(&file_path, &root_dir) {
            Ok(ns) => ns,
            Err(e) => {
                errors.push((
                    file_path.display().to_string(),
                    format!("Failed to extract namespace: {}", e),
                ));
                continue;
            }
        };

        // Create workspace with standard library and .mlf cache
        let mlf_cache_dir = crate::config::find_project_root(&std::env::current_dir().unwrap())
            .ok()
            .map(|root| crate::config::get_mlf_cache_dir(&root));

        let mut workspace =
            match crate::workspace_ext::workspace_with_std_and_cache(mlf_cache_dir.as_deref()) {
                Ok(ws) => ws,
                Err(e) => {
                    errors.push((
                        file_path.display().to_string(),
                        format!("Failed to load workspace: {}", e),
                    ));
                    continue;
                }
            };

        // Add the module to the workspace
        if let Err(e) = workspace.add_module(namespace.clone(), lexicon.clone()) {
            errors.push((
                file_path.display().to_string(),
                format!("Failed to add module: {:?}", e),
            ));
            continue;
        }

        // Resolve types
        if let Err(e) = workspace.resolve() {
            errors.push((
                file_path.display().to_string(),
                format!("Type resolution error: {:?}", e),
            ));
            continue;
        }

        // Generate code using the selected generator
        let ctx = GeneratorContext {
            namespace: &namespace,
            lexicon: &lexicon,
            workspace: &workspace,
        };

        let generated_code = match generator.generate(&ctx) {
            Ok(code) => code,
            Err(e) => {
                errors.push((
                    file_path.display().to_string(),
                    format!("Generation error: {}", e),
                ));
                continue;
            }
        };

        // Determine output path
        let output_path = if flat {
            let filename = format!("{}{}", namespace, generator.file_extension());
            output_dir.join(filename)
        } else {
            let mut path = output_dir.clone();
            for segment in namespace.split('.') {
                path.push(segment);
            }
            if let Err(e) = std::fs::create_dir_all(&path.parent().unwrap()) {
                errors.push((
                    file_path.display().to_string(),
                    format!("Failed to create directory: {}", e),
                ));
                continue;
            }
            path.set_extension(generator.file_extension().trim_start_matches('.'));
            path
        };

        // Write generated code
        if let Err(source) = std::fs::write(&output_path, generated_code) {
            errors.push((
                output_path.display().to_string(),
                format!("Failed to write file: {}", source),
            ));
            continue;
        }

        println!("Generated: {}", output_path.display());
        success_count += 1;
    }

    if !errors.is_empty() {
        eprintln!(
            "\n{} file(s) generated successfully, {} error(s) encountered:\n",
            success_count,
            errors.len()
        );
        for (path, error) in &errors {
            eprintln!("  {} - {}", path, error);
        }
        eprintln!();
        return Err(GenerateError::ParseLexicon {
            path: "multiple files".to_string(),
            help: Some(format!("{} errors total", errors.len())),
        });
    }

    println!("\nSuccessfully generated {} file(s)", success_count);
    Ok(())
}

/// Collect all .mlf files recursively from a directory
fn collect_mlf_files(dir: &Path) -> Result<Vec<PathBuf>, GenerateError> {
    let mut files = Vec::new();

    for entry in std::fs::read_dir(dir).map_err(|source| GenerateError::WriteOutput {
        path: dir.display().to_string(),
        source,
    })? {
        let entry = entry.map_err(|source| GenerateError::WriteOutput {
            path: dir.display().to_string(),
            source,
        })?;

        let path = entry.path();

        if path.is_dir() {
            files.extend(collect_mlf_files(&path)?);
        } else if path.extension().and_then(|s| s.to_str()) == Some("mlf") {
            files.push(path);
        }
    }

    Ok(files)
}

fn extract_namespace(file_path: &Path, root_dir: &Path) -> Result<String, std::io::Error> {
    // Canonicalize both paths for comparison
    let file_canonical = file_path.canonicalize()?;
    let root_canonical = root_dir.canonicalize()?;

    // Get the relative path from root to file
    let relative_path = file_canonical.strip_prefix(&root_canonical).map_err(|_| {
        std::io::Error::new(
            std::io::ErrorKind::Other,
            format!(
                "File path {} is not under root directory {}",
                file_path.display(),
                root_dir.display()
            ),
        )
    })?;

    // Convert path components to namespace parts
    let mut namespace_parts = Vec::new();

    for component in relative_path.components() {
        match component {
            std::path::Component::Normal(os_str) => {
                if let Some(s) = os_str.to_str() {
                    namespace_parts.push(s);
                }
            }
            _ => continue,
        }
    }

    // Remove .mlf extension from the last component if present
    if let Some(last) = namespace_parts.last_mut() {
        if let Some(stem) = last.strip_suffix(".mlf") {
            *last = stem;
        }
    }

    if namespace_parts.is_empty() {
        return Err(std::io::Error::new(
            std::io::ErrorKind::Other,
            format!(
                "Could not extract namespace from path: {}",
                file_path.display()
            ),
        ));
    }

    Ok(namespace_parts.join("."))
}