use std::collections::HashMap;

use crate::{
    diag, lexer,
    project::{ProjectArtifact, SourceFile, ToolProject},
    sexpr::{Atom, SExpr, SExprKind},
    token::Span,
};

pub fn render_file(path: &str, source: &str) -> Result<String, Vec<diag::Diagnostic>> {
    let source = SourceFile {
        path: path.to_string(),
        source: source.to_string(),
    };
    render_sources(&[source], None)
}

pub fn render_project(project: &ToolProject) -> Result<String, Vec<diag::Diagnostic>> {
    render_sources(&project.sources, Some(&project.artifact))
}

fn render_sources(
    sources: &[SourceFile],
    artifact: Option<&ProjectArtifact>,
) -> Result<String, Vec<diag::Diagnostic>> {
    let package_by_path = package_names_by_path(artifact);
    let mut modules = Vec::new();
    let mut diagnostics = Vec::new();

    let mut sorted_sources = sources.iter().collect::<Vec<_>>();
    sorted_sources.sort_by(|left, right| left.path.cmp(&right.path));

    for source in sorted_sources {
        match module_symbols(
            source,
            package_by_path.get(&source.path).map(String::as_str),
        ) {
            Ok(module) => modules.push(module),
            Err(mut errs) => diagnostics.append(&mut errs),
        }
    }

    if !diagnostics.is_empty() {
        return Err(diagnostics);
    }

    let mut out = String::new();
    out.push_str("(symbols\n");
    out.push_str("  (schema slovo.symbols)\n");
    out.push_str("  (version \"1.0.0-beta.10\")\n");
    out.push_str("  (modules");
    if modules.is_empty() {
        out.push_str(")\n");
    } else {
        out.push('\n');
        for module in &modules {
            render_module(module, &mut out);
        }
        out.push_str("  )\n");
    }
    out.push_str(")\n");
    Ok(out)
}

fn package_names_by_path(artifact: Option<&ProjectArtifact>) -> HashMap<String, String> {
    let mut packages = HashMap::new();
    let Some(workspace) = artifact.and_then(|artifact| artifact.workspace.as_ref()) else {
        return packages;
    };
    for package in &workspace.packages {
        for module in &package.modules {
            packages.insert(module.path.clone(), package.name.clone());
        }
    }
    packages
}

fn module_symbols(
    source: &SourceFile,
    package: Option<&str>,
) -> Result<ModuleSymbols, Vec<diag::Diagnostic>> {
    let tokens = lexer::lex(&source.path, &source.source)?;
    let forms = crate::sexpr::parse(&source.path, &tokens)?;

    let mut module = ModuleSymbols {
        name: "main".to_string(),
        path: source.path.clone(),
        package: package.map(str::to_string),
        module_span: None,
        exports: Vec::new(),
        imports: Vec::new(),
        type_aliases: Vec::new(),
        structs: Vec::new(),
        enums: Vec::new(),
        functions: Vec::new(),
        tests: Vec::new(),
        source: source.source.clone(),
    };

    for form in &forms {
        match list_head(form) {
            Some("module") => extract_module(form, &mut module),
            Some("import") => extract_import(form, &mut module.imports),
            Some("type") => extract_type_alias(form, &mut module.type_aliases),
            Some("struct") => extract_struct(form, &mut module.structs),
            Some("enum") => extract_enum(form, &mut module.enums),
            Some("fn") => extract_function(form, &mut module.functions),
            Some("test") => extract_test(form, &mut module.tests),
            _ => {}
        }
    }

    Ok(module)
}

fn extract_module(form: &SExpr, module: &mut ModuleSymbols) {
    let Some(items) = expect_list(form) else {
        return;
    };
    if let Some(name) = items.get(1).and_then(expect_ident) {
        module.name = name.to_string();
        module.module_span = Some(form.span);
    }
    if let Some(export_items) = items.get(2).and_then(expect_list) {
        if matches!(export_items.first().and_then(expect_ident), Some("export")) {
            module
                .exports
                .extend(export_items[1..].iter().filter_map(named_symbol_from_ident));
        }
    }
}

fn extract_import(form: &SExpr, imports: &mut Vec<ImportSymbol>) {
    let Some(items) = expect_list(form) else {
        return;
    };
    let Some(module_name) = items.get(1).and_then(expect_ident) else {
        return;
    };
    let names = items
        .get(2)
        .and_then(expect_list)
        .map(|items| items.iter().filter_map(named_symbol_from_ident).collect())
        .unwrap_or_default();
    imports.push(ImportSymbol {
        module: module_name.to_string(),
        span: form.span,
        module_span: items[1].span,
        names,
    });
}

fn extract_type_alias(form: &SExpr, aliases: &mut Vec<TypeAliasSymbol>) {
    let Some(items) = expect_list(form) else {
        return;
    };
    let Some(name) = items.get(1).and_then(named_symbol_from_ident) else {
        return;
    };
    aliases.push(TypeAliasSymbol {
        name,
        span: form.span,
        target_span: items.get(2).map(|item| item.span),
    });
}

fn extract_struct(form: &SExpr, structs: &mut Vec<StructSymbol>) {
    let Some(items) = expect_list(form) else {
        return;
    };
    let Some(name) = items.get(1).and_then(named_symbol_from_ident) else {
        return;
    };
    let fields = items[2..]
        .iter()
        .filter_map(|item| {
            let pair = expect_list(item)?;
            let field = pair.first().and_then(named_symbol_from_ident)?;
            Some(FieldSymbol {
                name: field,
                span: item.span,
                type_span: pair.get(1).map(|ty| ty.span),
            })
        })
        .collect();
    structs.push(StructSymbol {
        name,
        span: form.span,
        fields,
    });
}

fn extract_enum(form: &SExpr, enums: &mut Vec<EnumSymbol>) {
    let Some(items) = expect_list(form) else {
        return;
    };
    let Some(name) = items.get(1).and_then(named_symbol_from_ident) else {
        return;
    };
    let variants = items[2..]
        .iter()
        .filter_map(|item| {
            if let Some(name) = named_symbol_from_ident(item) {
                return Some(VariantSymbol {
                    name,
                    span: item.span,
                    payload_span: None,
                });
            }
            let variant_items = expect_list(item)?;
            let name = variant_items.first().and_then(named_symbol_from_ident)?;
            Some(VariantSymbol {
                name,
                span: item.span,
                payload_span: variant_items.get(1).map(|payload| payload.span),
            })
        })
        .collect();
    enums.push(EnumSymbol {
        name,
        span: form.span,
        variants,
    });
}

fn extract_function(form: &SExpr, functions: &mut Vec<FunctionSymbol>) {
    let Some(items) = expect_list(form) else {
        return;
    };
    let Some(name) = items.get(1).and_then(named_symbol_from_ident) else {
        return;
    };
    let params = items
        .get(2)
        .and_then(expect_list)
        .map(|items| {
            items
                .iter()
                .filter_map(|item| {
                    let pair = expect_list(item)?;
                    let name = pair.first().and_then(named_symbol_from_ident)?;
                    Some(ParamSymbol {
                        name,
                        span: item.span,
                        type_span: pair.get(1).map(|ty| ty.span),
                    })
                })
                .collect()
        })
        .unwrap_or_default();
    functions.push(FunctionSymbol {
        name,
        span: form.span,
        params,
        return_span: items.get(4).map(|item| item.span),
    });
}

fn extract_test(form: &SExpr, tests: &mut Vec<TestSymbol>) {
    let Some(items) = expect_list(form) else {
        return;
    };
    let Some(name_expr) = items.get(1) else {
        return;
    };
    let Some(name) = expect_string(name_expr) else {
        return;
    };
    tests.push(TestSymbol {
        name: NamedSymbol {
            name: name.to_string(),
            span: name_expr.span,
        },
        span: form.span,
    });
}

fn render_module(module: &ModuleSymbols, out: &mut String) {
    out.push_str("    (module\n");
    out.push_str(&format!(
        "      (name {})\n",
        diag::render_string(&module.name)
    ));
    if let Some(package) = module.package.as_deref() {
        out.push_str(&format!(
            "      (package {})\n",
            diag::render_string(package)
        ));
    }
    out.push_str(&format!(
        "      (path {})\n",
        diag::render_string(&module.path)
    ));
    if let Some(span) = module.module_span {
        render_span("      ", "module_span", span, &module.source, out);
    }
    render_named_symbols(
        "exports",
        "export",
        &module.exports,
        "      ",
        &module.source,
        out,
    );
    render_imports(&module.imports, module, out);
    render_type_aliases(&module.type_aliases, module, out);
    render_structs(&module.structs, module, out);
    render_enums(&module.enums, module, out);
    render_functions(&module.functions, module, out);
    render_tests(&module.tests, module, out);
    out.push_str("    )\n");
}

fn render_named_symbols(
    section: &str,
    item_name: &str,
    symbols: &[NamedSymbol],
    indent: &str,
    source: &str,
    out: &mut String,
) {
    out.push_str(indent);
    out.push('(');
    out.push_str(section);
    if symbols.is_empty() {
        out.push_str(")\n");
        return;
    }
    out.push('\n');
    for symbol in symbols {
        out.push_str(indent);
        out.push_str("  (");
        out.push_str(item_name);
        out.push('\n');
        out.push_str(indent);
        out.push_str("    (name ");
        out.push_str(&diag::render_string(&symbol.name));
        out.push_str(")\n");
        let symbol_indent = format!("{}    ", indent);
        render_span(&symbol_indent, "span", symbol.span, source, out);
        out.push_str(indent);
        out.push_str("  )\n");
    }
    out.push_str(indent);
    out.push_str(")\n");
}

fn render_imports(imports: &[ImportSymbol], module: &ModuleSymbols, out: &mut String) {
    out.push_str("      (imports");
    if imports.is_empty() {
        out.push_str(")\n");
        return;
    }
    out.push('\n');
    for import in imports {
        out.push_str("        (import\n");
        out.push_str(&format!(
            "          (module {})\n",
            diag::render_string(&import.module)
        ));
        render_span("          ", "span", import.span, &module.source, out);
        render_span(
            "          ",
            "module_span",
            import.module_span,
            &module.source,
            out,
        );
        render_named_symbols(
            "names",
            "name",
            &import.names,
            "          ",
            &module.source,
            out,
        );
        out.push_str("        )\n");
    }
    out.push_str("      )\n");
}

fn render_type_aliases(aliases: &[TypeAliasSymbol], module: &ModuleSymbols, out: &mut String) {
    out.push_str("      (type_aliases");
    if aliases.is_empty() {
        out.push_str(")\n");
        return;
    }
    out.push('\n');
    for alias in aliases {
        out.push_str("        (type_alias\n");
        render_decl_name(&alias.name, "          ", &module.source, out);
        render_span("          ", "span", alias.span, &module.source, out);
        if let Some(span) = alias.target_span {
            render_span("          ", "target_span", span, &module.source, out);
        }
        out.push_str("        )\n");
    }
    out.push_str("      )\n");
}

fn render_structs(structs: &[StructSymbol], module: &ModuleSymbols, out: &mut String) {
    out.push_str("      (structs");
    if structs.is_empty() {
        out.push_str(")\n");
        return;
    }
    out.push('\n');
    for struct_symbol in structs {
        out.push_str("        (struct\n");
        render_decl_name(&struct_symbol.name, "          ", &module.source, out);
        render_span(
            "          ",
            "span",
            struct_symbol.span,
            &module.source,
            out,
        );
        out.push_str("          (fields");
        if struct_symbol.fields.is_empty() {
            out.push_str(")\n");
        } else {
            out.push('\n');
            for field in &struct_symbol.fields {
                out.push_str("            (field\n");
                render_decl_name(&field.name, "              ", &module.source, out);
                render_span("              ", "span", field.span, &module.source, out);
                if let Some(span) = field.type_span {
                    render_span("              ", "type_span", span, &module.source, out);
                }
                out.push_str("            )\n");
            }
            out.push_str("          )\n");
        }
        out.push_str("        )\n");
    }
    out.push_str("      )\n");
}

fn render_enums(enums: &[EnumSymbol], module: &ModuleSymbols, out: &mut String) {
    out.push_str("      (enums");
    if enums.is_empty() {
        out.push_str(")\n");
        return;
    }
    out.push('\n');
    for enum_symbol in enums {
        out.push_str("        (enum\n");
        render_decl_name(&enum_symbol.name, "          ", &module.source, out);
        render_span("          ", "span", enum_symbol.span, &module.source, out);
        out.push_str("          (variants");
        if enum_symbol.variants.is_empty() {
            out.push_str(")\n");
        } else {
            out.push('\n');
            for variant in &enum_symbol.variants {
                out.push_str("            (variant\n");
                render_decl_name(&variant.name, "              ", &module.source, out);
                render_span("              ", "span", variant.span, &module.source, out);
                if let Some(span) = variant.payload_span {
                    render_span("              ", "payload_span", span, &module.source, out);
                }
                out.push_str("            )\n");
            }
            out.push_str("          )\n");
        }
        out.push_str("        )\n");
    }
    out.push_str("      )\n");
}

fn render_functions(functions: &[FunctionSymbol], module: &ModuleSymbols, out: &mut String) {
    out.push_str("      (functions");
    if functions.is_empty() {
        out.push_str(")\n");
        return;
    }
    out.push('\n');
    for function in functions {
        out.push_str("        (function\n");
        render_decl_name(&function.name, "          ", &module.source, out);
        render_span("          ", "span", function.span, &module.source, out);
        if let Some(span) = function.return_span {
            render_span("          ", "return_span", span, &module.source, out);
        }
        out.push_str("          (params");
        if function.params.is_empty() {
            out.push_str(")\n");
        } else {
            out.push('\n');
            for param in &function.params {
                out.push_str("            (param\n");
                render_decl_name(&param.name, "              ", &module.source, out);
                render_span("              ", "span", param.span, &module.source, out);
                if let Some(span) = param.type_span {
                    render_span("              ", "type_span", span, &module.source, out);
                }
                out.push_str("            )\n");
            }
            out.push_str("          )\n");
        }
        out.push_str("        )\n");
    }
    out.push_str("      )\n");
}

fn render_tests(tests: &[TestSymbol], module: &ModuleSymbols, out: &mut String) {
    out.push_str("      (tests");
    if tests.is_empty() {
        out.push_str(")\n");
        return;
    }
    out.push('\n');
    for test in tests {
        out.push_str("        (test\n");
        render_decl_name(&test.name, "          ", &module.source, out);
        render_span("          ", "span", test.span, &module.source, out);
        out.push_str("        )\n");
    }
    out.push_str("      )\n");
}

fn render_decl_name(symbol: &NamedSymbol, indent: &str, source: &str, out: &mut String) {
    out.push_str(indent);
    out.push_str(&format!("(name {})\n", diag::render_string(&symbol.name)));
    render_span(indent, "name_span", symbol.span, source, out);
}

fn render_span(indent: &str, label: &str, span: Span, source: &str, out: &mut String) {
    let start = position(source, span.start);
    let end = position(source, span.end);
    out.push_str(indent);
    out.push('(');
    out.push_str(label);
    out.push_str(&format!(
        " (span (start {}) (end {})) (range (start (line {}) (column {})) (end (line {}) (column {}))))\n",
        span.start, span.end, start.line, start.column, end.line, end.column
    ));
}

fn position(source: &str, offset: usize) -> Position {
    let mut line = 1;
    let mut column = 1;
    for (index, byte) in source.as_bytes().iter().enumerate() {
        if index >= offset {
            break;
        }
        if *byte == b'\n' {
            line += 1;
            column = 1;
        } else {
            column += 1;
        }
    }
    Position { line, column }
}

fn list_head(form: &SExpr) -> Option<&str> {
    expect_list(form)?.first().and_then(expect_ident)
}

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

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

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

fn named_symbol_from_ident(form: &SExpr) -> Option<NamedSymbol> {
    Some(NamedSymbol {
        name: expect_ident(form)?.to_string(),
        span: form.span,
    })
}

struct ModuleSymbols {
    name: String,
    path: String,
    package: Option<String>,
    module_span: Option<Span>,
    exports: Vec<NamedSymbol>,
    imports: Vec<ImportSymbol>,
    type_aliases: Vec<TypeAliasSymbol>,
    structs: Vec<StructSymbol>,
    enums: Vec<EnumSymbol>,
    functions: Vec<FunctionSymbol>,
    tests: Vec<TestSymbol>,
    source: String,
}

struct NamedSymbol {
    name: String,
    span: Span,
}

struct ImportSymbol {
    module: String,
    span: Span,
    module_span: Span,
    names: Vec<NamedSymbol>,
}

struct TypeAliasSymbol {
    name: NamedSymbol,
    span: Span,
    target_span: Option<Span>,
}

struct StructSymbol {
    name: NamedSymbol,
    span: Span,
    fields: Vec<FieldSymbol>,
}

struct FieldSymbol {
    name: NamedSymbol,
    span: Span,
    type_span: Option<Span>,
}

struct EnumSymbol {
    name: NamedSymbol,
    span: Span,
    variants: Vec<VariantSymbol>,
}

struct VariantSymbol {
    name: NamedSymbol,
    span: Span,
    payload_span: Option<Span>,
}

struct FunctionSymbol {
    name: NamedSymbol,
    span: Span,
    params: Vec<ParamSymbol>,
    return_span: Option<Span>,
}

struct ParamSymbol {
    name: NamedSymbol,
    span: Span,
    type_span: Option<Span>,
}

struct TestSymbol {
    name: NamedSymbol,
    span: Span,
}

struct Position {
    line: usize,
    column: usize,
}
