use std::{
    collections::{BTreeMap, BTreeSet},
    fs,
    path::Path,
};

use crate::{
    ast::Program,
    diag::Diagnostic,
    lexer, lower,
    project::{self, ProjectArtifact, SourceFile, ToolFailure, WorkspaceArtifact},
    sexpr::{Atom, SExpr, SExprKind},
    types::Type,
};

pub fn generate(input: &str, output_dir: &str) -> Result<(), ToolFailure> {
    let docs = if project::is_project_input(input) {
        let loaded = project::load_project_sources_for_tools(input)?;
        render_project(&loaded.artifact, &loaded.sources)?
    } else {
        let source = fs::read_to_string(input).map_err(|err| ToolFailure {
            diagnostics: vec![Diagnostic::new(
                input,
                "InputReadFailed",
                format!("cannot read `{}`: {}", input, err),
            )],
            sources: Vec::new(),
            artifact: None,
        })?;
        let sources = vec![SourceFile {
            path: input.to_string(),
            source,
        }];
        render_file(input, &sources)?
    };

    fs::create_dir_all(output_dir).map_err(|err| ToolFailure {
        diagnostics: vec![Diagnostic::new(
            output_dir,
            "OutputWriteFailed",
            format!(
                "cannot create documentation directory `{}`: {}",
                output_dir, err
            ),
        )],
        sources: Vec::new(),
        artifact: None,
    })?;
    fs::write(Path::new(output_dir).join("index.md"), docs).map_err(|err| ToolFailure {
        diagnostics: vec![Diagnostic::new(
            output_dir,
            "OutputWriteFailed",
            format!("cannot write documentation into `{}`: {}", output_dir, err),
        )],
        sources: Vec::new(),
        artifact: None,
    })
}

fn render_project(
    artifact: &ProjectArtifact,
    sources: &[SourceFile],
) -> Result<String, ToolFailure> {
    let modules = sources
        .iter()
        .map(document_source)
        .collect::<Result<Vec<_>, _>>()?;

    let mut out = String::new();
    out.push_str("# Project ");
    out.push_str(&artifact.project_name);
    out.push_str("\n\n");
    if let Some(workspace) = &artifact.workspace {
        render_workspace(&mut out, workspace);
        render_workspace_package_public_api(&mut out, workspace, &modules);
    } else {
        render_project_package_public_api(&mut out, artifact, &modules);
    }
    for module in &modules {
        render_module(&mut out, module);
    }
    Ok(out)
}

fn render_workspace(out: &mut String, workspace: &WorkspaceArtifact) {
    out.push_str("## Workspace\n\n");
    render_list(out, "Members", &workspace.members);

    let packages = workspace
        .packages
        .iter()
        .map(|package| {
            format!(
                "{} {} (entry {})",
                package.name, package.version, package.entry
            )
        })
        .collect::<Vec<_>>();
    render_list(out, "Packages", &packages);

    let dependencies = workspace
        .dependencies
        .iter()
        .map(|dependency| {
            format!(
                "{} -> {} ({}, {})",
                dependency.from, dependency.to, dependency.kind, dependency.path
            )
        })
        .collect::<Vec<_>>();
    render_list(out, "Package Dependencies", &dependencies);
}

fn render_file(title: &str, sources: &[SourceFile]) -> Result<String, ToolFailure> {
    let mut out = String::new();
    out.push_str("# File ");
    out.push_str(title);
    out.push_str("\n\n");
    let module = document_source(&sources[0])?;
    render_module(&mut out, &module);
    Ok(out)
}

fn document_source(source: &SourceFile) -> Result<DocModule, ToolFailure> {
    let tokens = lexer::lex(&source.path, &source.source).map_err(|diagnostics| ToolFailure {
        diagnostics,
        sources: vec![source.clone()],
        artifact: None,
    })?;
    let forms = crate::sexpr::parse(&source.path, &tokens).map_err(|diagnostics| ToolFailure {
        diagnostics,
        sources: vec![source.clone()],
        artifact: None,
    })?;
    let lowerable_forms = forms
        .iter()
        .filter(|form| !matches!(list_head(form), Some("import")))
        .cloned()
        .collect::<Vec<_>>();
    let program = lower::lower_program(&source.path, &lowerable_forms).ok();
    Ok(module_from_forms(&source.path, &forms, program.as_ref()))
}

struct DocModule {
    path: String,
    title: String,
    imports: Vec<String>,
    exports: Vec<String>,
    structs: Vec<String>,
    functions: Vec<String>,
    tests: Vec<String>,
    public_api: PublicApi,
}

#[derive(Default)]
struct PublicApi {
    functions: Vec<DocFunction>,
    structs: Vec<DocStruct>,
    enums: Vec<DocEnum>,
}

struct DocFunction {
    name: String,
    signature: String,
}

struct DocStruct {
    name: String,
    fields: Vec<String>,
}

struct DocEnum {
    name: String,
    variants: Vec<String>,
}

fn module_from_forms(file: &str, forms: &[SExpr], program: Option<&Program>) -> DocModule {
    let mut module = DocModule {
        path: file.to_string(),
        title: file.to_string(),
        imports: Vec::new(),
        exports: Vec::new(),
        structs: Vec::new(),
        functions: Vec::new(),
        tests: Vec::new(),
        public_api: PublicApi::default(),
    };

    for form in forms {
        let Some(items) = list(form) else {
            continue;
        };
        match items.first().and_then(ident) {
            Some("module") => {
                if let Some(name) = items.get(1).and_then(ident) {
                    module.title = name.to_string();
                }
                if let Some(exports) = items.get(2).and_then(list) {
                    if matches!(exports.first().and_then(ident), Some("export")) {
                        module
                            .exports
                            .extend(exports[1..].iter().filter_map(ident).map(str::to_string));
                    }
                }
            }
            Some("import") => {
                if let Some(name) = items.get(1).and_then(ident) {
                    module.imports.push(name.to_string());
                }
            }
            Some("struct") => {
                if let Some(name) = items.get(1).and_then(ident) {
                    module.structs.push(name.to_string());
                }
            }
            Some("fn") => {
                if let Some(name) = items.get(1).and_then(ident) {
                    module.functions.push(name.to_string());
                }
            }
            Some("test") => {
                if let Some(name) = items.get(1).and_then(string_atom) {
                    module.tests.push(name.to_string());
                }
            }
            _ => {}
        }
    }

    if let Some(program) = program {
        module.title = program.module.clone();
        module.structs = program
            .structs
            .iter()
            .map(|decl| decl.name.clone())
            .collect();
        module.functions = program
            .functions
            .iter()
            .map(|function| {
                let params = function
                    .params
                    .iter()
                    .map(|param| format!("{} {}", param.name, param.ty))
                    .collect::<Vec<_>>()
                    .join(", ");
                format!("{}({}) -> {}", function.name, params, function.return_type)
            })
            .collect();
        module.tests = program.tests.iter().map(|test| test.name.clone()).collect();
        module.public_api = public_api_from_program(program, &module.exports);
    }

    module.imports.sort();
    module.exports.sort();
    module.structs.sort();
    module.functions.sort();
    module.tests.sort();
    module
}

fn render_module(out: &mut String, module: &DocModule) {
    out.push_str("## Module ");
    out.push_str(&module.title);
    out.push_str("\n\n");
    render_list(out, "Imports", &module.imports);
    render_list(out, "Exports", &module.exports);
    render_public_api_section(out, "###", &module.public_api);
    render_list(out, "Structs", &module.structs);
    render_list(out, "Functions", &module.functions);
    render_list(out, "Tests", &module.tests);
}

fn render_project_package_public_api(
    out: &mut String,
    artifact: &ProjectArtifact,
    modules: &[DocModule],
) {
    let local_root = Path::new(&artifact.project_root).join(&artifact.source_root);
    let local_modules = modules
        .iter()
        .filter(|module| Path::new(&module.path).starts_with(&local_root))
        .collect::<Vec<_>>();
    render_package_public_api(out, &artifact.project_name, &local_modules);
}

fn render_workspace_package_public_api(
    out: &mut String,
    workspace: &WorkspaceArtifact,
    modules: &[DocModule],
) {
    let by_path = modules
        .iter()
        .map(|module| (module.path.as_str(), module))
        .collect::<BTreeMap<_, _>>();

    for package in &workspace.packages {
        let local_root = Path::new(&package.root).join(&package.source_root);
        let package_modules = package
            .modules
            .iter()
            .filter(|module| Path::new(&module.path).starts_with(&local_root))
            .filter_map(|module| by_path.get(module.path.as_str()).copied())
            .collect::<Vec<_>>();
        render_package_public_api(
            out,
            &format!("{} {}", package.name, package.version),
            &package_modules,
        );
    }
}

fn render_package_public_api(out: &mut String, package_name: &str, modules: &[&DocModule]) {
    out.push_str("## Package API ");
    out.push_str(package_name);
    out.push_str("\n\n");

    let public_modules = modules
        .iter()
        .copied()
        .filter(|module| !module.public_api.is_empty())
        .collect::<Vec<_>>();
    if public_modules.is_empty() {
        out.push_str("None.\n\n");
        return;
    }

    for module in public_modules {
        out.push_str("### Module ");
        out.push_str(&module.title);
        out.push_str("\n\n");
        render_public_api_body(out, "####", &module.public_api);
    }
}

fn render_public_api_section(out: &mut String, heading: &str, public_api: &PublicApi) {
    out.push_str(heading);
    out.push_str(" Public API\n\n");
    render_public_api_body(out, "####", public_api);
}

fn render_public_api_body(out: &mut String, heading: &str, public_api: &PublicApi) {
    if public_api.is_empty() {
        out.push_str("None.\n\n");
        return;
    }

    if !public_api.functions.is_empty() {
        out.push_str(heading);
        out.push_str(" Functions\n");
        for function in &public_api.functions {
            render_code_bullet(out, &function.signature);
        }
        out.push('\n');
    }

    if !public_api.structs.is_empty() {
        out.push_str(heading);
        out.push_str(" Structs\n");
        for struct_decl in &public_api.structs {
            render_code_bullet(out, &format!("struct {}", struct_decl.name));
            for field in &struct_decl.fields {
                render_indented_code_bullet(out, field);
            }
        }
        out.push('\n');
    }

    if !public_api.enums.is_empty() {
        out.push_str(heading);
        out.push_str(" Enums\n");
        for enum_decl in &public_api.enums {
            render_code_bullet(out, &format!("enum {}", enum_decl.name));
            for variant in &enum_decl.variants {
                render_indented_code_bullet(out, variant);
            }
        }
        out.push('\n');
    }
}

fn render_list(out: &mut String, title: &str, values: &[String]) {
    out.push_str("### ");
    out.push_str(title);
    out.push('\n');
    if values.is_empty() {
        out.push_str("\nNone.\n\n");
        return;
    }
    for value in values {
        out.push_str("- `");
        out.push_str(&value.replace('`', "\\`"));
        out.push_str("`\n");
    }
    out.push('\n');
}

fn list_head(expr: &SExpr) -> Option<&str> {
    list(expr).and_then(|items| items.first()).and_then(ident)
}

fn render_code_bullet(out: &mut String, value: &str) {
    out.push_str("- `");
    out.push_str(&value.replace('`', "\\`"));
    out.push_str("`\n");
}

fn render_indented_code_bullet(out: &mut String, value: &str) {
    out.push_str("  - `");
    out.push_str(&value.replace('`', "\\`"));
    out.push_str("`\n");
}

impl PublicApi {
    fn is_empty(&self) -> bool {
        self.functions.is_empty() && self.structs.is_empty() && self.enums.is_empty()
    }
}

fn public_api_from_program(program: &Program, exports: &[String]) -> PublicApi {
    let export_names = exports.iter().cloned().collect::<BTreeSet<_>>();
    let aliases = alias_targets(program);

    let mut functions = program
        .functions
        .iter()
        .filter(|function| export_names.contains(&function.name))
        .map(|function| DocFunction {
            name: function.name.clone(),
            signature: function_signature(
                &function.name,
                function
                    .params
                    .iter()
                    .map(|param| (param.name.as_str(), &param.ty)),
                &function.return_type,
                &aliases,
            ),
        })
        .collect::<Vec<_>>();
    functions.sort_by(|left, right| left.name.cmp(&right.name));

    let mut structs = program
        .structs
        .iter()
        .filter(|struct_decl| export_names.contains(&struct_decl.name))
        .map(|struct_decl| DocStruct {
            name: struct_decl.name.clone(),
            fields: struct_decl
                .fields
                .iter()
                .map(|field| {
                    format!(
                        "{}: {}",
                        field.name,
                        display_public_type(&field.ty, &aliases)
                    )
                })
                .collect(),
        })
        .collect::<Vec<_>>();
    structs.sort_by(|left, right| left.name.cmp(&right.name));

    let mut enums = program
        .enums
        .iter()
        .filter(|enum_decl| export_names.contains(&enum_decl.name))
        .map(|enum_decl| DocEnum {
            name: enum_decl.name.clone(),
            variants: enum_decl
                .variants
                .iter()
                .map(|variant| match &variant.payload_ty {
                    Some(payload_ty) => {
                        format!(
                            "{}({})",
                            variant.name,
                            display_public_type(payload_ty, &aliases)
                        )
                    }
                    None => variant.name.clone(),
                })
                .collect(),
        })
        .collect::<Vec<_>>();
    enums.sort_by(|left, right| left.name.cmp(&right.name));

    PublicApi {
        functions,
        structs,
        enums,
    }
}

fn function_signature<'a>(
    name: &str,
    params: impl Iterator<Item = (&'a str, &'a Type)>,
    return_type: &Type,
    aliases: &BTreeMap<String, Type>,
) -> String {
    let params = params
        .map(|(name, ty)| format!("{}: {}", name, display_public_type(ty, aliases)))
        .collect::<Vec<_>>()
        .join(", ");
    format!(
        "fn {}({}) -> {}",
        name,
        params,
        display_public_type(return_type, aliases)
    )
}

fn alias_targets(program: &Program) -> BTreeMap<String, Type> {
    let raw = program
        .type_aliases
        .iter()
        .map(|alias| (alias.name.clone(), alias.target.clone()))
        .collect::<BTreeMap<_, _>>();
    raw.keys()
        .map(|name| {
            let mut visiting = BTreeSet::new();
            (
                name.clone(),
                resolve_alias_type(&Type::Named(name.clone()), &raw, &mut visiting),
            )
        })
        .collect()
}

fn display_public_type(ty: &Type, aliases: &BTreeMap<String, Type>) -> String {
    let mut visiting = BTreeSet::new();
    resolve_alias_type(ty, aliases, &mut visiting).to_string()
}

fn resolve_alias_type(
    ty: &Type,
    aliases: &BTreeMap<String, Type>,
    visiting: &mut BTreeSet<String>,
) -> Type {
    match ty {
        Type::Named(name) => {
            let Some(target) = aliases.get(name) else {
                return ty.clone();
            };
            if !visiting.insert(name.clone()) {
                return ty.clone();
            }
            let resolved = resolve_alias_type(target, aliases, visiting);
            visiting.remove(name);
            resolved
        }
        Type::Ptr(inner) => Type::Ptr(Box::new(resolve_alias_type(inner, aliases, visiting))),
        Type::Array(inner, len) => {
            Type::Array(Box::new(resolve_alias_type(inner, aliases, visiting)), *len)
        }
        Type::Vec(inner) => Type::Vec(Box::new(resolve_alias_type(inner, aliases, visiting))),
        Type::Slice(inner) => Type::Slice(Box::new(resolve_alias_type(inner, aliases, visiting))),
        Type::Option(inner) => Type::Option(Box::new(resolve_alias_type(inner, aliases, visiting))),
        Type::Result(ok, err) => Type::Result(
            Box::new(resolve_alias_type(ok, aliases, visiting)),
            Box::new(resolve_alias_type(err, aliases, visiting)),
        ),
        Type::I32
        | Type::I64
        | Type::U32
        | Type::U64
        | Type::F64
        | Type::Bool
        | Type::Unit
        | Type::String => ty.clone(),
    }
}

fn list(expr: &SExpr) -> Option<&[SExpr]> {
    match &expr.kind {
        SExprKind::List(items) => Some(items),
        _ => None,
    }
}

fn ident(expr: &SExpr) -> Option<&str> {
    match &expr.kind {
        SExprKind::Atom(Atom::Ident(value)) => Some(value),
        _ => None,
    }
}

fn string_atom(expr: &SExpr) -> Option<&str> {
    match &expr.kind {
        SExprKind::Atom(Atom::String(value)) => Some(value),
        _ => None,
    }
}