mimium-lang 4.0.1

use std::collections::{HashMap, HashSet};
use std::path::{Path, PathBuf};

use serde::{Deserialize, Serialize};

use super::resolve_include::resolve_include;
use super::statement::Statement;
use crate::ast::Expr;
use crate::ast::statement::into_then_expr;
use crate::interner::{ExprNodeId, Symbol, ToSymbol, TypeNodeId};
use crate::pattern::TypedId;
use crate::types::{PType, RecordTypeField, Type};
use crate::utils::error::{ReportableError, SimpleError};
use crate::utils::metadata::{Location, Span};

use super::StageKind;

/// Visibility modifier for module members
#[derive(Clone, Debug, PartialEq, Default, Serialize, Deserialize)]
pub enum Visibility {
    #[default]
    Private,
    Public,
}

/// Qualified path for module references (e.g., modA::modB::func)
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct QualifiedPath {
    pub segments: Vec<Symbol>,
}

impl QualifiedPath {
    pub fn new(segments: Vec<Symbol>) -> Self {
        Self { segments }
    }

    pub fn single(name: Symbol) -> Self {
        Self {
            segments: vec![name],
        }
    }
}

/// Target of a use statement: single symbol, multiple symbols, or wildcard
#[derive(Clone, Debug, PartialEq)]
pub enum UseTarget {
    /// Single import: `use foo::bar`
    Single,
    /// Multiple imports: `use foo::{bar, baz}`
    Multiple(Vec<Symbol>),
    /// Wildcard import: `use foo::*`
    Wildcard,
}

/// A variant definition for user-defined sum types
#[derive(Clone, Debug, PartialEq)]
pub struct VariantDef {
    /// The name of the variant constructor
    pub name: Symbol,
    /// Optional payload type for the variant
    pub payload: Option<TypeNodeId>,
}

impl VariantDef {
    pub fn new(name: Symbol, payload: Option<TypeNodeId>) -> Self {
        Self { name, payload }
    }
}

#[derive(Clone, Debug, PartialEq)]
pub enum ProgramStatement {
    FnDefinition {
        visibility: Visibility,
        name: Symbol,
        args: (Vec<TypedId>, Location),
        return_type: Option<TypeNodeId>,
        body: ExprNodeId,
    },
    StageDeclaration {
        stage: StageKind,
    },
    GlobalStatement(Statement),
    Import(Symbol),
    /// Module definition: mod name { ... } (inline) or mod name; (external file)
    ModuleDefinition {
        visibility: Visibility,
        name: Symbol,
        /// Body of the module. None means external file module (mod foo;)
        body: Option<Vec<(ProgramStatement, Span)>>,
    },
    /// Use statement: `use path::to::item`, `use path::{a, b}`, or `use path::*`
    /// Can be prefixed with `pub` for re-exporting
    UseStatement {
        /// Visibility of the re-export (pub use for re-exporting)
        visibility: Visibility,
        /// The base path (for `use foo::bar`, this is `[foo, bar]`; for `use foo::{a, b}`, this is `[foo]`)
        path: QualifiedPath,
        /// The import target type
        target: UseTarget,
    },
    /// Type alias: type Alias = BaseType
    TypeAlias {
        visibility: Visibility,
        name: Symbol,
        target_type: crate::interner::TypeNodeId,
    },
    /// Type declaration: type Name = Variant1 | Variant2 | ...
    TypeDeclaration {
        visibility: Visibility,
        name: Symbol,
        variants: Vec<VariantDef>,
        /// Whether this type was declared with `type rec` (allows recursive references)
        is_recursive: bool,
    },
    Comment(Symbol),
    DocComment(Symbol),
    Error,
}

/// Information about a type declaration, including its variants and recursion flag
#[derive(Clone, Debug, PartialEq)]
pub struct TypeDeclInfo {
    pub variants: Vec<VariantDef>,
    /// Whether this type was declared with `type rec` (allows recursive references)
    pub is_recursive: bool,
}

/// Map from type name to type declaration info
pub type TypeDeclarationMap = HashMap<Symbol, TypeDeclInfo>;
/// Map from type alias name to target type
pub type TypeAliasMap = HashMap<Symbol, crate::interner::TypeNodeId>;

#[derive(Clone, Debug, PartialEq, Default)]
pub struct Program {
    pub statements: Vec<(ProgramStatement, Span)>,
}

/// Convert a qualified path to a mangled symbol name.
/// For example, `foo::bar::baz` becomes `foo$bar$baz`.
fn mangle_qualified_name(prefix: &[Symbol], name: Symbol) -> Symbol {
    use crate::interner::ToSymbol;
    if prefix.is_empty() {
        name
    } else {
        let path_str = prefix
            .iter()
            .map(|s| s.as_str())
            .collect::<Vec<_>>()
            .join("$");
        format!("{}${}", path_str, name.as_str()).to_symbol()
    }
}

fn is_reserved_type_param_name(name: Symbol) -> bool {
    let s = name.as_str();
    s.len() == 1 && s.as_bytes()[0].is_ascii_lowercase()
}

/// Convert a full qualified path (all segments) to a mangled symbol name.
/// For example, `[foo, bar, baz]` becomes `foo$bar$baz`.
fn mangle_qualified_path(segments: &[Symbol]) -> Symbol {
    use crate::interner::ToSymbol;
    segments
        .iter()
        .map(|s| s.as_str())
        .collect::<Vec<_>>()
        .join("$")
        .to_symbol()
}

/// Resolve an external file module (`mod foo;` syntax).
/// Looks for `{name}.mmm` in the same directory as the current file.
fn resolve_external_module(
    name: Symbol,
    file_path: &Path,
    span: Span,
    errs: &mut Vec<Box<dyn ReportableError>>,
    module_prefix: &[Symbol],
    module_info: &mut ModuleInfo,
) -> Vec<(Statement, Location)> {
    let module_filename = format!("{}.mmm", name.as_str());
    let (imported, mut new_errs) =
        resolve_include(file_path.to_str().unwrap(), &module_filename, span);
    errs.append(&mut new_errs);

    // Process imported program with the module prefix
    stmts_from_program_with_prefix(
        imported.program.statements,
        imported.resolved_path,
        errs,
        module_prefix,
        module_info,
    )
}

/// Map from mangled symbol name to whether it's public.
/// Only contains entries for module members (not top-level definitions).
pub type VisibilityMap = HashMap<Symbol, bool>;

/// Map from alias name to the mangled name it refers to.
/// Created from `use` statements, e.g., `use foo::bar` creates `bar -> foo$bar`.
pub type UseAliasMap = HashMap<Symbol, Symbol>;

/// Map from mangled function name to its module context (prefix).
/// Used for relative path resolution within modules.
pub type ModuleContextMap = HashMap<Symbol, Vec<Symbol>>;

/// Module-related information collected during parsing.
/// Contains visibility information for module members and use aliases.
#[derive(Clone, Debug, Default)]
pub struct ModuleInfo {
    /// Map from mangled symbol name to whether it's public (only for module members)
    pub visibility_map: VisibilityMap,
    /// Map from alias name to mangled name (from use statements)
    pub use_alias_map: UseAliasMap,
    /// Map from mangled function name to its module context (for relative path resolution)
    pub module_context_map: ModuleContextMap,
    /// List of wildcard import base paths (e.g., `use foo::*` stores "foo")
    pub wildcard_imports: Vec<Symbol>,
    /// Type declarations for user-defined sum types
    pub type_declarations: TypeDeclarationMap,
    /// Type aliases for simple type aliases
    pub type_aliases: TypeAliasMap,
    /// Loaded external modules to avoid duplicate loading when resolving `use` statements
    pub loaded_external_modules: HashSet<Symbol>,
}

impl ModuleInfo {
    pub fn new() -> Self {
        Self::default()
    }
}

/// Resolve a qualified path, trying relative resolution from current module context.
/// Returns `(mangled_name, resolved_path_segments)`.
///
/// For example, if current context is `[outer]` and path is `[inner, secret]`:
/// 1. First tries `inner$secret` (absolute path)
/// 2. If not found, tries `outer$inner$secret` (relative path from current module)
///
/// The `exists` closure should return `true` if the given mangled name exists in the environment.
pub fn resolve_qualified_path<F>(
    path_segments: &[Symbol],
    absolute_mangled: Symbol,
    current_module_context: &[Symbol],
    exists: F,
) -> (Symbol, Vec<Symbol>)
where
    F: Fn(&Symbol) -> bool,
{
    // First, try the absolute path
    if exists(&absolute_mangled) {
        return (absolute_mangled, path_segments.to_vec());
    }

    // If not found and we have a module context, try relative path
    if !current_module_context.is_empty() {
        // Build the relative path: context + path_segments
        let mut relative_path = current_module_context.to_vec();
        relative_path.extend(path_segments.iter().copied());

        let relative_mangled = relative_path
            .iter()
            .map(|s| s.as_str())
            .collect::<Vec<_>>()
            .join("$")
            .to_symbol();

        if exists(&relative_mangled) {
            return (relative_mangled, relative_path);
        }
    }

    // Return absolute path if relative resolution failed
    (absolute_mangled, path_segments.to_vec())
}

fn stmts_from_program(
    program: Program,
    file_path: PathBuf,
    errs: &mut Vec<Box<dyn ReportableError>>,
    module_info: &mut ModuleInfo,
) -> Vec<(Statement, Location)> {
    stmts_from_program_with_prefix(program.statements, file_path, errs, &[], module_info)
}

fn stmts_from_program_with_prefix(
    statements: Vec<(ProgramStatement, Span)>,
    file_path: PathBuf,
    errs: &mut Vec<Box<dyn ReportableError>>,
    module_prefix: &[Symbol],
    module_info: &mut ModuleInfo,
) -> Vec<(Statement, Location)> {
    // Track the current stage so that module/use wrappers can restore it correctly.
    let mut current_stage = StageKind::Main;
    let mut result = Vec::new();

    for (stmt, span) in statements {
        let stmts: Option<Vec<(Statement, Location)>> = match stmt {
            ProgramStatement::FnDefinition {
                visibility,
                name,
                args,
                return_type,
                body,
            } => {
                let loc = Location::new(span, file_path.clone());
                let argloc = args.1.clone();
                let argsty = args
                    .clone()
                    .0
                    .into_iter()
                    .map(RecordTypeField::from)
                    .collect::<Vec<_>>();
                let argty = match argsty.as_slice() {
                    [] => Type::Primitive(PType::Unit).into_id_with_location(argloc.clone()),
                    [arg] => arg.ty,
                    _ => Type::Record(argsty).into_id_with_location(argloc),
                };
                let fnty = Type::Function {
                    arg: argty,
                    ret: return_type.unwrap_or(Type::Unknown.into_id_with_location(loc.clone())),
                }
                .into_id_with_location(loc.clone());
                // Use mangled name if inside a module
                let mangled_name = mangle_qualified_name(module_prefix, name);
                // Track visibility for all functions (both module members and top-level)
                module_info
                    .visibility_map
                    .insert(mangled_name, visibility == Visibility::Public);
                // Track module context for relative path resolution
                if !module_prefix.is_empty() {
                    module_info
                        .module_context_map
                        .insert(mangled_name, module_prefix.to_vec());
                }
                Some(vec![(
                    Statement::LetRec(
                        TypedId::new(mangled_name, fnty),
                        Expr::Lambda(args.0, return_type, body).into_id(loc.clone()),
                    ),
                    loc,
                )])
            }
            ProgramStatement::GlobalStatement(statement) => {
                if !module_prefix.is_empty() {
                    collect_statement_bindings(&statement)
                        .into_iter()
                        .for_each(|name| {
                            module_info
                                .module_context_map
                                .insert(name, module_prefix.to_vec());
                        });
                }
                Some(vec![(statement, Location::new(span, file_path.clone()))])
            }
            ProgramStatement::Comment(_) | ProgramStatement::DocComment(_) => None,
            ProgramStatement::Import(filename) => {
                let (imported, mut new_errs) =
                    resolve_include(file_path.to_str().unwrap(), filename.as_str(), span.clone());
                errs.append(&mut new_errs);
                let res =
                    stmts_from_program(imported.program, imported.resolved_path, errs, module_info);
                Some(res)
            }
            ProgramStatement::StageDeclaration { stage } => {
                current_stage = stage.clone();
                Some(vec![(
                    Statement::DeclareStage(stage),
                    Location::new(span, file_path.clone()),
                )])
            }
            ProgramStatement::ModuleDefinition {
                visibility: _,
                name,
                body,
            } => {
                let module_symbol = mangle_qualified_name(module_prefix, name);
                module_info.loaded_external_modules.insert(module_symbol);

                // Flatten module contents with qualified names
                let mut new_prefix = module_prefix.to_vec();
                new_prefix.push(name);

                let inner_stmts = match body {
                    Some(inline_body) => {
                        // Inline module: mod foo { ... }
                        stmts_from_program_with_prefix(
                            inline_body,
                            file_path.clone(),
                            errs,
                            &new_prefix,
                            module_info,
                        )
                    }
                    None => {
                        // External file module: mod foo;
                        resolve_external_module(
                            name,
                            &file_path,
                            span.clone(),
                            errs,
                            &new_prefix,
                            module_info,
                        )
                    }
                };

                // Wrap module contents with stage boundary:
                // - Start with #stage(main) to isolate module from consumer's stage context
                // - End with the consumer's current stage to restore it after module processing
                let module_loc = Location::new(span, file_path.clone());
                let start_decl = (Statement::DeclareStage(StageKind::Main), module_loc.clone());
                let restore_decl = (Statement::DeclareStage(current_stage.clone()), module_loc);
                let result = [vec![start_decl], inner_stmts, vec![restore_decl]].concat();

                Some(result)
            }
            ProgramStatement::UseStatement {
                visibility,
                path,
                target,
            } => {
                let imported_stmts = if let Some(base_module) = path.segments.first().copied() {
                    let local_module_symbol = mangle_qualified_name(module_prefix, base_module);
                    let absolute_module_symbol = mangle_qualified_name(&[], base_module);
                    if module_info
                        .loaded_external_modules
                        .contains(&local_module_symbol)
                        || module_info
                            .loaded_external_modules
                            .contains(&absolute_module_symbol)
                    {
                        vec![]
                    } else {
                        module_info
                            .loaded_external_modules
                            .insert(absolute_module_symbol);
                        let new_prefix = vec![base_module];
                        let inner_stmts = resolve_external_module(
                            base_module,
                            &file_path,
                            span.clone(),
                            errs,
                            &new_prefix,
                            module_info,
                        );
                        let module_loc = Location::new(span.clone(), file_path.clone());
                        let start_decl =
                            (Statement::DeclareStage(StageKind::Main), module_loc.clone());
                        let restore_decl =
                            (Statement::DeclareStage(current_stage.clone()), module_loc);
                        [vec![start_decl], inner_stmts, vec![restore_decl]].concat()
                    }
                } else {
                    vec![]
                };

                process_use_statement(&visibility, &path, &target, module_prefix, module_info);
                (!imported_stmts.is_empty()).then_some(imported_stmts)
            }
            ProgramStatement::TypeAlias {
                visibility,
                name,
                target_type,
            } => {
                if is_reserved_type_param_name(name) {
                    errs.push(Box::new(SimpleError {
                        message: format!(
                            "type name '{}' is reserved for explicit type parameters (single lowercase letter)",
                            name.as_str()
                        ),
                        span: Location::new(span.clone(), file_path.clone()),
                    }));
                    continue;
                }
                // Store type alias for later use in type environment
                let mangled_name = mangle_qualified_name(module_prefix, name);
                module_info.type_aliases.insert(mangled_name, target_type);
                // Track visibility for all type aliases (both module members and top-level)
                module_info
                    .visibility_map
                    .insert(mangled_name, visibility == Visibility::Public);
                // Track module context for relative path resolution
                if !module_prefix.is_empty() {
                    module_info
                        .module_context_map
                        .insert(mangled_name, module_prefix.to_vec());
                }
                None
            }
            ProgramStatement::TypeDeclaration {
                visibility,
                name,
                variants,
                is_recursive,
            } => {
                if is_reserved_type_param_name(name) {
                    errs.push(Box::new(SimpleError {
                        message: format!(
                            "type name '{}' is reserved for explicit type parameters (single lowercase letter)",
                            name.as_str()
                        ),
                        span: Location::new(span.clone(), file_path.clone()),
                    }));
                    continue;
                }
                // Store type declaration for later use in type environment
                let mangled_name = mangle_qualified_name(module_prefix, name);
                module_info.type_declarations.insert(
                    mangled_name,
                    TypeDeclInfo {
                        variants,
                        is_recursive,
                    },
                );
                // Track visibility for type declarations
                module_info
                    .visibility_map
                    .insert(mangled_name, visibility == Visibility::Public);
                if !module_prefix.is_empty() {
                    // Track module context for relative path resolution
                    module_info
                        .module_context_map
                        .insert(mangled_name, module_prefix.to_vec());
                }
                None
            }
            ProgramStatement::Error => Some(vec![(
                Statement::Error,
                Location::new(span, file_path.clone()),
            )]),
        };
        if let Some(stmts) = stmts {
            result.extend(stmts);
        }
    }
    result
}

fn collect_pattern_bindings(pat: &crate::pattern::Pattern, out: &mut Vec<Symbol>) {
    match pat {
        crate::pattern::Pattern::Single(name) => out.push(*name),
        crate::pattern::Pattern::Tuple(items) => {
            items.iter().for_each(|p| collect_pattern_bindings(p, out));
        }
        crate::pattern::Pattern::Record(fields) => {
            fields
                .iter()
                .for_each(|(_, p)| collect_pattern_bindings(p, out));
        }
        crate::pattern::Pattern::Placeholder | crate::pattern::Pattern::Error => {}
    }
}

fn collect_statement_bindings(stmt: &Statement) -> Vec<Symbol> {
    match stmt {
        Statement::Let(typed_pat, _) => {
            let mut symbols = vec![];
            collect_pattern_bindings(&typed_pat.pat, &mut symbols);
            symbols
        }
        Statement::LetRec(id, _) => vec![id.id],
        Statement::Single(expr) => collect_expr_bindings(*expr),
        _ => vec![],
    }
}

fn collect_expr_bindings(expr: ExprNodeId) -> Vec<Symbol> {
    match expr.to_expr() {
        Expr::Let(typed_pat, _body, then_opt) => {
            let mut symbols = vec![];
            collect_pattern_bindings(&typed_pat.pat, &mut symbols);
            if let Some(then) = then_opt {
                symbols.extend(collect_expr_bindings(then));
            }
            symbols
        }
        Expr::LetRec(id, _body, then_opt) => {
            let mut symbols = vec![id.id];
            if let Some(then) = then_opt {
                symbols.extend(collect_expr_bindings(then));
            }
            symbols
        }
        _ => vec![],
    }
}

/// Process a use statement, registering aliases in module_info.
/// - Single: `use foo::bar` → alias `bar` → `foo$bar`
/// - Multiple: `use foo::{a, b}` → alias `a` → `foo$a`, `b` → `foo$b`
/// - Wildcard: `use foo::*` → import all public members from foo
/// If visibility is Public, the imported names are re-exported as public.
fn process_use_statement(
    visibility: &Visibility,
    path: &QualifiedPath,
    target: &UseTarget,
    module_prefix: &[Symbol],
    module_info: &mut ModuleInfo,
) {
    let resolve_use_mangled = |segments: &[Symbol], info: &ModuleInfo| {
        let absolute_mangled = mangle_qualified_path(segments);
        let (resolved, _) =
            resolve_qualified_path(segments, absolute_mangled, module_prefix, |name| {
                info.visibility_map.contains_key(name)
                    || info.use_alias_map.contains_key(name)
                    || info.module_context_map.contains_key(name)
                    || info.type_aliases.contains_key(name)
                    || info.type_declarations.contains_key(name)
            });
        resolved
    };

    fn register_alias(
        module_info: &mut ModuleInfo,
        visibility: &Visibility,
        module_prefix: &[Symbol],
        alias_name: Symbol,
        mangled: Symbol,
    ) {
        module_info.use_alias_map.insert(alias_name, mangled);

        if *visibility == Visibility::Public {
            let exported_name = mangle_qualified_name(module_prefix, alias_name);
            module_info.visibility_map.insert(exported_name, true);
            module_info.use_alias_map.insert(exported_name, mangled);
        }
    }

    match target {
        UseTarget::Single => {
            // use foo::bar creates an alias: bar -> foo$bar
            if let Some(alias_name) = path.segments.last().copied() {
                let mangled = resolve_use_mangled(&path.segments, module_info);
                register_alias(module_info, visibility, module_prefix, alias_name, mangled);
            }
        }
        UseTarget::Multiple(names) => {
            // use foo::{bar, baz} creates:
            //   bar -> foo$bar
            //   baz -> foo$baz
            for name in names {
                let mut full_path = path.segments.clone();
                full_path.push(*name);
                let mangled = resolve_use_mangled(&full_path, module_info);
                register_alias(module_info, visibility, module_prefix, *name, mangled);
            }
        }
        UseTarget::Wildcard => {
            // use foo::* imports all public members from foo
            // We need to defer this until we know all the public members
            // For now, store the base path for later resolution
            let base_mangled = if path.segments.is_empty() {
                // use ::* at module level means import from parent (current prefix)
                module_prefix
                    .iter()
                    .map(|s| s.as_str())
                    .collect::<Vec<_>>()
                    .join("$")
            } else {
                mangle_qualified_path(&path.segments).as_str().to_string()
            };

            // Store wildcard import for later resolution
            // The key is the base path (e.g., "foo"), value is the prefix for resolving members
            module_info.wildcard_imports.push(base_mangled.to_symbol());

            // Note: For pub use foo::*, we can't know all exported names at this point.
            // Wildcard re-exports would require a second pass or runtime resolution.
            // For now, wildcard imports with pub are stored but individual re-exports
            // need to be resolved later when the symbols are actually accessed.
        }
    }
}

pub(crate) fn expr_from_program(
    program: Program,
    file_path: PathBuf,
) -> (ExprNodeId, ModuleInfo, Vec<Box<dyn ReportableError>>) {
    let mut errs = vec![];
    let mut module_info = ModuleInfo::new();
    let stmts = stmts_from_program(program, file_path.clone(), &mut errs, &mut module_info);

    let res = into_then_expr(stmts.as_slice()).unwrap_or(Expr::Error.into_id_without_span());

    (res, module_info, errs)
}