pub mod ast;
pub mod bytecode;
pub mod cache;
pub mod diagnostic;
pub mod identity;
pub mod implementation_artifact;
pub mod interface;
pub mod known_enum;
pub mod lexer;
pub mod mir;
pub mod package;
pub mod package_signature;
pub mod parser;
pub mod prelude;
pub mod resolver;
pub mod runtime;
pub mod span;
pub(crate) mod std_package;
pub mod symbol;
pub mod token;
pub mod typed_hir;
pub mod types;
pub mod typing;

use ast::Program;
use bytecode::Program as BytecodeProgram;
use diagnostic::Diagnostic;
use interface::PackageInterfaceGraph;
use mir::Program as MirProgram;
use runtime::RunOutcome;
use std::path::{Path, PathBuf};
use typed_hir::Program as TypedHirProgram;

#[derive(Clone, Debug)]
pub struct PackageAwareCheck {
    pub packages: package::LoadedPackageGraph,
    pub signatures: package_signature::PackageSignatureEnvironment,
    pub module_checks: Vec<PackageModuleCheck>,
    pub typed_program: TypedHirProgram,
}

#[derive(Clone, Debug)]
pub struct PackageModuleCheck {
    pub package: identity::PackageId,
    pub module: identity::ModuleId,
    pub module_path: String,
    pub resolve_output: resolver::ResolveOutput,
    pub type_output: typing::TypeCheckOutput,
    pub typed_program: TypedHirProgram,
}

pub fn check_source(source: &str) -> Result<Program, Vec<Diagnostic>> {
    let tokens = lexer::lex(source)?;
    let program = parser::parse(tokens)?;
    if program.package.is_some() {
        return Err(vec![Diagnostic::new(
            "PK001",
            "package mode requires a file-based entrypoint",
            Default::default(),
        )]);
    }

    let mut diagnostics = resolver::resolve(&program);
    diagnostics.extend(typing::typecheck(&program));

    if diagnostics.is_empty() {
        Ok(program)
    } else {
        Err(diagnostics)
    }
}

pub fn check_path(path: &Path) -> Result<Program, Vec<Diagnostic>> {
    if package::entry_package_path_from_entry(path)?.is_some() {
        let check = check_package_aware_path(path)?;
        return check.packages.entry_program().cloned().ok_or_else(|| {
            vec![Diagnostic::new(
                "PK018",
                "entry module was not loaded in package-aware checking",
                Default::default(),
            )]
        });
    }
    let program = package::load_flattened_program_from_entry(path)?;
    let mut diagnostics = resolver::resolve(&program);
    diagnostics.extend(typing::typecheck(&program));

    if diagnostics.is_empty() {
        Ok(program)
    } else {
        Err(diagnostics)
    }
}

pub fn compile_source(source: &str) -> Result<MirProgram, Vec<Diagnostic>> {
    compile_mir_source(source)
}

pub fn compile_path(path: &Path) -> Result<MirProgram, Vec<Diagnostic>> {
    compile_mir_path(path)
}

pub fn compile_mir_source(source: &str) -> Result<MirProgram, Vec<Diagnostic>> {
    let program = compile_typed_source(source)?;
    Ok(mir::lower_typed(&program))
}

pub fn compile_mir_path(path: &Path) -> Result<MirProgram, Vec<Diagnostic>> {
    let program = compile_typed_path(path)?;
    Ok(mir::lower_typed(&program))
}

pub fn compile_typed_source(source: &str) -> Result<TypedHirProgram, Vec<Diagnostic>> {
    let tokens = lexer::lex(source)?;
    let program = parser::parse(tokens)?;
    if program.package.is_some() {
        return Err(vec![Diagnostic::new(
            "PK001",
            "package mode requires a file-based entrypoint",
            Default::default(),
        )]);
    }

    let resolve_output = resolver::resolve_program(&program);
    let type_output = typing::typecheck_program(&program);
    let mut diagnostics = resolve_output.diagnostics;
    diagnostics.extend(type_output.diagnostics.clone());
    if diagnostics.is_empty() {
        Ok(typed_hir::lower(
            &program,
            &type_output,
            package::PackageSymbolGraph::default(),
        ))
    } else {
        Err(diagnostics)
    }
}

pub fn compile_typed_path(path: &Path) -> Result<TypedHirProgram, Vec<Diagnostic>> {
    if package::entry_package_path_from_entry(path)?.is_some() {
        return check_package_aware_path(path).map(|check| check.typed_program);
    }
    let loaded = package::load_flattened_from_entry(path)?;
    compile_flattened_typed_program(loaded)
}

pub fn check_package_aware_path(path: &Path) -> Result<PackageAwareCheck, Vec<Diagnostic>> {
    let packages = package::load_package_graph_from_entry(path)?;
    let diagnostics = package::validate_loaded_package_graph(&packages);
    if !diagnostics.is_empty() {
        return Err(diagnostics);
    }
    let signatures = package_signature::PackageSignatureEnvironment::from_loaded_graph(&packages)?;
    let module_checks = typecheck_loaded_package_modules(&packages, &signatures)?;
    let typed_program = package_typed_program(&packages, &module_checks)?;
    Ok(PackageAwareCheck {
        packages,
        signatures,
        module_checks,
        typed_program,
    })
}

pub fn check_package_aware_path_against_loaded_interfaces(
    path: &Path,
    interfaces: &PackageInterfaceGraph,
    interface_symbols: &symbol::SymbolTable,
) -> Result<PackageAwareCheck, Vec<Diagnostic>> {
    let packages = package::load_package_graph_from_entry_against_interfaces(
        path,
        interfaces,
        interface_symbols,
    )?;
    let diagnostics = package::validate_loaded_package_graph(&packages);
    if !diagnostics.is_empty() {
        return Err(diagnostics);
    }
    let signatures = package_signature::PackageSignatureEnvironment::from_loaded_graph(&packages)?;
    let module_checks = typecheck_loaded_package_modules(&packages, &signatures)?;
    let typed_program = package_typed_program(&packages, &module_checks)?;
    Ok(PackageAwareCheck {
        packages,
        signatures,
        module_checks,
        typed_program,
    })
}

fn typecheck_loaded_package_modules(
    packages: &package::LoadedPackageGraph,
    signatures: &package_signature::PackageSignatureEnvironment,
) -> Result<Vec<PackageModuleCheck>, Vec<Diagnostic>> {
    let mut module_checks = Vec::new();
    let mut diagnostics = Vec::new();
    for package in &packages.packages {
        if packages.is_loaded_interface_package_path(&package.path) {
            continue;
        }
        let Some(package_id) = packages.package_graph.package_id(&package.path) else {
            continue;
        };
        for file in &package.files {
            let Some(module_id) = packages
                .package_graph
                .module_id(package_id, &file.module_path)
            else {
                continue;
            };
            let resolve_output =
                resolver::resolve_package_module(&file.program, signatures, module_id);
            let has_resolve_diagnostics = !resolve_output.diagnostics.is_empty();
            diagnostics.extend(resolve_output.diagnostics.clone());
            if has_resolve_diagnostics {
                continue;
            }
            let type_output =
                typing::typecheck_package_module(&file.program, signatures, module_id);
            let has_diagnostics = !type_output.diagnostics.is_empty();
            diagnostics.extend(type_output.diagnostics.clone());
            if has_diagnostics {
                continue;
            }
            let mut typed_program =
                typed_hir::lower(&file.program, &type_output, packages.package_graph.clone());
            attach_package_items_to_module_typed_program(
                &mut typed_program,
                &packages.package_graph,
                module_id,
            );
            module_checks.push(PackageModuleCheck {
                package: package_id,
                module: module_id,
                module_path: file.module_path.clone(),
                resolve_output,
                type_output,
                typed_program,
            });
        }
    }
    if diagnostics.is_empty() {
        Ok(module_checks)
    } else {
        Err(diagnostics)
    }
}

fn package_typed_program(
    packages: &package::LoadedPackageGraph,
    module_checks: &[PackageModuleCheck],
) -> Result<TypedHirProgram, Vec<Diagnostic>> {
    let entry_exists = module_checks.iter().any(|check| {
        check.package == packages.entry_package && check.module == packages.entry_module
    });
    if !entry_exists {
        return Err(vec![Diagnostic::new(
            "PK018",
            "entry module was not typechecked in package-aware checking",
            Default::default(),
        )]);
    }
    let modules = module_checks
        .iter()
        .map(|check| check.typed_program.clone())
        .collect::<Vec<_>>();
    Ok(typed_hir::merge_modules(
        &modules,
        packages.package_graph.clone(),
    ))
}

fn attach_package_items_to_module_typed_program(
    program: &mut TypedHirProgram,
    package_graph: &package::PackageSymbolGraph,
    module: identity::ModuleId,
) {
    for statement in &mut program.statements {
        match statement {
            typed_hir::Stmt::Record(record) => {
                record.package_item = record.package_item.or_else(|| {
                    package_graph.item_id_in_module(
                        module,
                        &record.name,
                        package::PackageItemKind::Record,
                    )
                });
            }
            typed_hir::Stmt::Enum(enumeration) => {
                enumeration.package_item = enumeration.package_item.or_else(|| {
                    package_graph.item_id_in_module(
                        module,
                        &enumeration.name,
                        package::PackageItemKind::Enum,
                    )
                });
            }
            typed_hir::Stmt::Function(function) => {
                function.package_item = function.package_item.or_else(|| {
                    package_graph.item_id_in_module(
                        module,
                        &function.name,
                        package::PackageItemKind::Function,
                    )
                });
            }
            _ => {}
        }
    }
}

pub fn compile_typed_path_against_loaded_interfaces(
    path: &Path,
    interfaces: &PackageInterfaceGraph,
    interface_symbols: &symbol::SymbolTable,
) -> Result<TypedHirProgram, Vec<Diagnostic>> {
    check_package_aware_path_against_loaded_interfaces(path, interfaces, interface_symbols)
        .map(|check| check.typed_program)
}

pub fn compile_typed_path_against_interface_artifacts(
    path: &Path,
    interface_root: &Path,
) -> Result<TypedHirProgram, Vec<Diagnostic>> {
    let package_paths = package::import_paths_from_entry(path)?;
    let mut symbols = symbol::SymbolTable::default();
    let interfaces = PackageInterfaceGraph::read_persisted_artifacts(
        interface_root,
        &package_paths,
        &mut symbols,
    )?;
    compile_typed_path_against_loaded_interfaces(path, &interfaces, &symbols)
}

pub fn check_package_aware_path_against_interface_artifacts(
    path: &Path,
    interface_root: &Path,
) -> Result<PackageAwareCheck, Vec<Diagnostic>> {
    let package_paths = package::import_paths_from_entry(path)?;
    let mut symbols = symbol::SymbolTable::default();
    let interfaces = PackageInterfaceGraph::read_persisted_artifacts(
        interface_root,
        &package_paths,
        &mut symbols,
    )?;
    check_package_aware_path_against_loaded_interfaces(path, &interfaces, &symbols)
}

pub fn package_check_cache_key(
    path: &Path,
    interface_root: &Path,
) -> Result<cache::PackageCheckCacheKey, Vec<Diagnostic>> {
    cache::compute_package_check_cache_key(path, interface_root)
}

pub fn write_package_check_cache_artifact(
    path: &Path,
    key: &cache::PackageCheckCacheKey,
) -> Result<(), Diagnostic> {
    cache::write_package_check_artifact(path, key)
}

pub fn write_package_check_cache_artifact_for_root(
    path: &Path,
    artifact_root: &Path,
) -> Result<PathBuf, Vec<Diagnostic>> {
    check_package_aware_path_against_interface_artifacts(path, artifact_root)?;
    let key = cache::compute_package_check_cache_key(path, artifact_root)?;
    let artifact_path = cache::package_check_artifact_path_from_entry(artifact_root, path)?;
    cache::write_package_check_artifact(&artifact_path, &key)
        .map_err(|diagnostic| vec![diagnostic])?;
    Ok(artifact_path)
}

pub fn package_check_cache_artifact_path(
    root: &Path,
    entry_path: &Path,
) -> Result<std::path::PathBuf, Vec<Diagnostic>> {
    cache::package_check_artifact_path_from_entry(root, entry_path)
}

pub fn write_package_interface_artifacts(
    path: &Path,
    artifact_root: &Path,
    package_paths: &[String],
) -> Result<Vec<PathBuf>, Vec<Diagnostic>> {
    let check = check_package_aware_path(path)?;
    let program = check.typed_program;
    let interfaces = program.package_interfaces();
    let requested_packages = if package_paths.is_empty() {
        interfaces
            .packages
            .iter()
            .map(|package| package.path.clone())
            .collect::<Vec<_>>()
    } else {
        package_paths.to_vec()
    };

    let mut written = Vec::new();
    let mut diagnostics = Vec::new();
    for package_path in requested_packages {
        match interfaces.write_persisted_artifact(artifact_root, &package_path, &program.symbols) {
            Ok(path) => written.push(path),
            Err(diagnostic) => diagnostics.push(diagnostic),
        }
    }

    if diagnostics.is_empty() {
        Ok(written)
    } else {
        Err(diagnostics)
    }
}

pub fn write_package_implementation_artifacts(
    path: &Path,
    artifact_root: &Path,
) -> Result<Vec<PathBuf>, Vec<Diagnostic>> {
    let check = check_package_aware_path(path)?;
    let interfaces = check.typed_program.package_interfaces();
    let mut written = Vec::new();
    let mut diagnostics = Vec::new();

    for package in &check.packages.packages {
        let Some(package_id) = check.packages.package_graph.package_id(&package.path) else {
            continue;
        };
        let modules = check
            .module_checks
            .iter()
            .filter(|module| module.package == package_id)
            .map(|module| module.typed_program.clone())
            .collect::<Vec<_>>();
        if modules.is_empty() {
            continue;
        }
        let typed_program =
            typed_hir::merge_modules(&modules, check.packages.package_graph.clone());
        let bytecode_program = bytecode::compile(mir::lower_typed(&typed_program));
        match implementation_artifact::PackageImplementationArtifact::from_bytecode_package(
            &package.path,
            &interfaces,
            &check.typed_program.symbols,
            &check.packages.package_graph,
            bytecode_program,
        )
        .and_then(|artifact| artifact.write_persisted_artifact(artifact_root))
        {
            Ok(path) => written.push(path),
            Err(diagnostic) => diagnostics.push(diagnostic),
        }
    }

    if diagnostics.is_empty() {
        Ok(written)
    } else {
        Err(diagnostics)
    }
}

pub fn write_package_artifacts(
    path: &Path,
    artifact_root: &Path,
) -> Result<Vec<PathBuf>, Vec<Diagnostic>> {
    let mut paths = write_package_interface_artifacts(path, artifact_root, &[])?;
    paths.extend(write_package_implementation_artifacts(path, artifact_root)?);
    paths.push(write_package_check_cache_artifact_for_root(
        path,
        artifact_root,
    )?);
    Ok(paths)
}

pub fn compile_typed_path_against_cached_interface_artifacts(
    path: &Path,
    interface_root: &Path,
    checked_artifact_path: &Path,
) -> Result<TypedHirProgram, Vec<Diagnostic>> {
    let key = cache::compute_package_check_cache_key(path, interface_root)?;
    cache::validate_package_check_artifact(checked_artifact_path, &key)?;
    compile_typed_path_against_interface_artifacts(path, interface_root)
}

pub fn check_package_aware_path_against_cached_interface_artifacts(
    path: &Path,
    interface_root: &Path,
    checked_artifact_path: &Path,
) -> Result<PackageAwareCheck, Vec<Diagnostic>> {
    let key = cache::compute_package_check_cache_key(path, interface_root)?;
    cache::validate_package_check_artifact(checked_artifact_path, &key)?;
    check_package_aware_path_against_interface_artifacts(path, interface_root)
}

pub fn compile_typed_path_against_cached_artifact_root(
    path: &Path,
    artifact_root: &Path,
) -> Result<TypedHirProgram, Vec<Diagnostic>> {
    let checked_artifact_path = cache::package_check_artifact_path_from_entry(artifact_root, path)?;
    compile_typed_path_against_cached_interface_artifacts(
        path,
        artifact_root,
        &checked_artifact_path,
    )
}

pub fn check_package_aware_path_against_cached_artifact_root(
    path: &Path,
    artifact_root: &Path,
) -> Result<PackageAwareCheck, Vec<Diagnostic>> {
    let checked_artifact_path = cache::package_check_artifact_path_from_entry(artifact_root, path)?;
    check_package_aware_path_against_cached_interface_artifacts(
        path,
        artifact_root,
        &checked_artifact_path,
    )
}

fn compile_flattened_typed_program(
    loaded: package::LoadedFlattenedProgram,
) -> Result<TypedHirProgram, Vec<Diagnostic>> {
    let resolve_output = resolver::resolve_program(&loaded.program);
    let type_output = typing::typecheck_program(&loaded.program);
    let mut diagnostics = resolve_output.diagnostics;
    diagnostics.extend(type_output.diagnostics.clone());
    if diagnostics.is_empty() {
        let program = typed_hir::lower(&loaded.program, &type_output, loaded.package_graph);
        let interfaces = program.package_interfaces();
        diagnostics.extend(program.validate_package_references_against_interfaces(&interfaces));
        if diagnostics.is_empty() {
            Ok(program)
        } else {
            Err(diagnostics)
        }
    } else {
        Err(diagnostics)
    }
}

pub fn compile_bytecode_source(source: &str) -> Result<BytecodeProgram, Vec<Diagnostic>> {
    let program = compile_mir_source(source)?;
    Ok(bytecode::compile(program))
}

pub fn compile_bytecode_path(path: &Path) -> Result<BytecodeProgram, Vec<Diagnostic>> {
    let program = compile_mir_path(path)?;
    Ok(bytecode::compile(program))
}

pub fn compile_bytecode_path_for_run_against_artifact_root(
    path: &Path,
    artifact_root: &Path,
) -> Result<BytecodeProgram, Vec<Diagnostic>> {
    let entry = compile_typed_path_against_cached_artifact_root(path, artifact_root)?;
    let entry = bytecode::compile(mir::lower_typed(&entry));
    let package_paths = package::import_paths_from_entry(path)?;
    let mut symbols = symbol::SymbolTable::default();
    let interfaces = PackageInterfaceGraph::read_persisted_artifacts(
        artifact_root,
        &package_paths,
        &mut symbols,
    )?;
    let implementation_artifacts =
        implementation_artifact::read_persisted_artifacts_reserving_program_items(
            artifact_root,
            &interfaces,
            &symbols,
            &[&entry],
        )?;
    let dependencies = implementation_artifact::programs_from_artifacts(implementation_artifacts);
    Ok(bytecode::merge(entry, dependencies))
}

pub fn run_source(source: &str) -> Result<RunOutcome, Vec<Diagnostic>> {
    let program = compile_bytecode_source(source)?;
    runtime::run(&program)
}

pub fn run_path(path: &Path) -> Result<RunOutcome, Vec<Diagnostic>> {
    let program = compile_bytecode_path(path)?;
    runtime::run(&program)
}

pub fn run_path_against_artifact_root(
    path: &Path,
    artifact_root: &Path,
) -> Result<RunOutcome, Vec<Diagnostic>> {
    let program = compile_bytecode_path_for_run_against_artifact_root(path, artifact_root)?;
    runtime::run(&program)
}
