use crate::backend::emit_fragments;
use crate::building::plan::{BuildStep, BuildTarget};
use crate::common::diagnostics::DiagnosticBundle;
use crate::common::files::FileManager;
use crate::common::spans::{FileId, Span};
use crate::deps::resolver::ProjectKey;
use prometeu_analysis::ids::ProjectId;
use crate::frontends::pbs::ast::ParsedAst;
use crate::frontends::pbs::collector::SymbolCollector;
use crate::frontends::pbs::lowering::Lowerer;
use crate::frontends::pbs::parser::Parser;
use crate::frontends::pbs::resolver::{ModuleProvider, Resolver};
use crate::frontends::pbs::symbols::{ModuleSymbols, Namespace, Symbol, SymbolKind, Visibility};
use crate::frontends::pbs::typecheck::TypeChecker;
use crate::frontends::pbs::types::PbsType;
use crate::lowering::core_to_vm;
use crate::semantics::export_surface::ExportSurfaceKind;
use prometeu_bytecode::{ConstantPoolEntry, DebugInfo, FunctionMeta};
use prometeu_analysis::NameInterner;
use serde::{Deserialize, Serialize};
use std::collections::{BTreeMap, HashMap};

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
pub struct ExportKey {
    pub module_path: String,
    pub symbol_name: String,
    pub kind: ExportSurfaceKind,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExportMetadata {
    pub func_idx: Option<u32>,
    pub is_host: bool,
    pub ty: Option<PbsType>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
pub struct ImportKey {
    pub dep_alias: String,
    pub module_path: String,
    pub symbol_name: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ImportMetadata {
    pub key: ImportKey,
    pub relocation_pcs: Vec<u32>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CompiledModule {
    pub project_id: ProjectId,
    pub project_key: ProjectKey,
    pub target: BuildTarget,
    pub exports: BTreeMap<ExportKey, ExportMetadata>,
    pub imports: Vec<ImportMetadata>,
    pub const_pool: Vec<ConstantPoolEntry>,
    pub code: Vec<u8>,
    pub function_metas: Vec<FunctionMeta>,
    pub debug_info: Option<DebugInfo>,
    pub symbols: Vec<crate::common::symbols::Symbol>,
}

#[derive(Debug)]
pub enum CompileError {
    Frontend(crate::common::diagnostics::DiagnosticBundle),
    DuplicateExport {
        symbol: String,
        first_dep: String,
        second_dep: String,
    },
    Io(std::io::Error),
    Internal(String),
}

impl std::fmt::Display for CompileError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            CompileError::Frontend(d) => write!(f, "Frontend error: {:?}", d),
            CompileError::DuplicateExport { symbol, first_dep, second_dep } => {
                write!(f, "duplicate export: symbol `{}`\n  first defined in dependency `{}`\n  again defined in dependency `{}`", symbol, first_dep, second_dep)
            }
            CompileError::Io(e) => write!(f, "IO error: {}", e),
            CompileError::Internal(s) => write!(f, "Internal error: {}", s),
        }
    }
}

impl std::error::Error for CompileError {}

impl From<std::io::Error> for CompileError {
    fn from(e: std::io::Error) -> Self {
        CompileError::Io(e)
    }
}

impl From<crate::common::diagnostics::DiagnosticBundle> for CompileError {
    fn from(d: crate::common::diagnostics::DiagnosticBundle) -> Self {
        CompileError::Frontend(d)
    }
}

struct ProjectModuleProvider {
    modules: HashMap<String, ModuleSymbols>,
}

impl ModuleProvider for ProjectModuleProvider {
    fn get_module_symbols(&self, from_path: &str) -> Option<&ModuleSymbols> {
        self.modules.get(from_path)
    }
}

pub fn compile_project(
    step: BuildStep,
    dep_modules: &HashMap<ProjectId, CompiledModule>,
    file_manager: &mut FileManager,
) -> Result<CompiledModule, CompileError> {
    let mut interner = NameInterner::new();
    // 1. Parse all files and group symbols by module
    let mut module_symbols_map: HashMap<String, ModuleSymbols> = HashMap::new();
    let mut parsed_files: Vec<(String, ParsedAst)> = Vec::new(); // (module_path, ast)

    for source_rel in &step.sources {
        let source_abs = step.project_dir.join(source_rel);
        let source_code = std::fs::read_to_string(&source_abs)?;
        let file_id = file_manager.add(source_abs.clone(), source_code.clone());
        
        let mut parser = Parser::new(&source_code, FileId(file_id as u32), &mut interner);
        let parsed = parser.parse_file()?;

        let mut collector = SymbolCollector::new(&interner);
        let (ts, vs) = collector.collect(&parsed.arena, parsed.root)?;
        
        let full_path = source_rel.to_string_lossy().replace('\\', "/");
        let logical_module_path = if let Some(stripped) = full_path.strip_prefix("src/main/modules/") {
            stripped
        } else if let Some(stripped) = full_path.strip_prefix("src/test/modules/") {
            stripped
        } else {
            &full_path
        };

        let module_path = std::path::Path::new(logical_module_path)
            .parent()
            .map(|p| p.to_string_lossy().replace('\\', "/"))
            .unwrap_or_else(|| "".to_string());
        
        let ms = module_symbols_map.entry(module_path.clone()).or_insert_with(ModuleSymbols::new);
        
        // Merge symbols
        for sym in ts.symbols.into_values() {
            if let Some(existing) = ms.type_symbols.get(sym.name) {
                return Err(DiagnosticBundle::error(
                    "E_RESOLVE_DUPLICATE_SYMBOL",
                    format!(
                        "Duplicate type symbol '{}' in module '{}'",
                        interner.resolve(existing.name),
                        module_path
                    ),
                    existing.span.clone(),
                ).into());
            }
            let _ = ms.type_symbols.insert(sym);
        }
        for sym in vs.symbols.into_values() {
            if let Some(existing) = ms.value_symbols.get(sym.name) {
                return Err(DiagnosticBundle::error(
                    "E_RESOLVE_DUPLICATE_SYMBOL",
                    format!(
                        "Duplicate value symbol '{}' in module '{}'",
                        interner.resolve(existing.name),
                        module_path
                    ),
                    existing.span.clone(),
                ).into());
            }
            let _ = ms.value_symbols.insert(sym);
        }
        
        parsed_files.push((module_path, parsed));
    }

    // 2. Synthesize ModuleSymbols for dependencies
    let mut all_visible_modules = module_symbols_map.clone();
    for (alias, project_id) in &step.deps {
        if let Some(compiled) = dep_modules.get(project_id) {
            for (key, meta) in &compiled.exports {
                // Support syntax: "alias/module" and "@alias:module"
                let key_module_path = &key.module_path;
                let synthetic_paths = [
                    format!("{}/{}", alias, key_module_path),
                    format!("@{}:{}", alias, key_module_path),
                ];
                
                for synthetic_module_path in synthetic_paths {
                    let ms = all_visible_modules.entry(synthetic_module_path.clone()).or_insert_with(ModuleSymbols::new);
                    
                    let sym = Symbol {
                        name: interner.intern(&key.symbol_name),
                        kind: match key.kind {
                            ExportSurfaceKind::Service => SymbolKind::Service,
                            ExportSurfaceKind::DeclareType => {
                                match &meta.ty {
                                    Some(PbsType::Contract(_)) => SymbolKind::Contract,
                                    Some(PbsType::ErrorType(_)) => SymbolKind::ErrorType,
                                    _ => SymbolKind::Struct,
                                }
                            }
                            ExportSurfaceKind::Function => SymbolKind::Function,
                        },
                        namespace: key.kind.namespace(),
                        visibility: Visibility::Pub,
                        ty: meta.ty.clone(),
                        is_host: meta.is_host,
                        span: Span::new(FileId::INVALID, 0, 0),
                        origin: Some(synthetic_module_path.clone()),
                    };
                    
                    if sym.namespace == Namespace::Type {
                        if let Some(existing) = ms.type_symbols.get(sym.name) {
                            return Err(CompileError::DuplicateExport {
                                symbol: interner.resolve(sym.name).to_string(),
                                first_dep: existing.origin.clone().unwrap_or_else(|| "unknown".to_string()),
                                second_dep: sym.origin.unwrap_or_else(|| "unknown".to_string()),
                            });
                        }
                        let _ = ms.type_symbols.insert(sym.clone());
                    } else {
                        if let Some(existing) = ms.value_symbols.get(sym.name) {
                            return Err(CompileError::DuplicateExport {
                                symbol: interner.resolve(sym.name).to_string(),
                                first_dep: existing.origin.clone().unwrap_or_else(|| "unknown".to_string()),
                                second_dep: sym.origin.unwrap_or_else(|| "unknown".to_string()),
                            });
                        }
                        let _ = ms.value_symbols.insert(sym.clone());
                    }
                }
            }
        }
    }

    // 3. Resolve and TypeCheck each file
    let module_provider = ProjectModuleProvider { 
        modules: all_visible_modules,
    };
    
    // We need to collect imported symbols for Lowerer
    let mut file_imported_symbols: HashMap<String, ModuleSymbols> = HashMap::new(); // keyed by module_path

    for (module_path, parsed) in &parsed_files {
        let ms = module_symbols_map.get(module_path).unwrap();
        // Ensure primitive names are interned before creating resolver/bootstrap
        {
            let primitives = ["int", "bool", "float", "string", "bounded", "void"];
            for p in primitives { interner.intern(p); }
        }
        let mut resolver = Resolver::new(ms, &module_provider, &interner);
        resolver.bootstrap_types(&interner);
        resolver.resolve(&parsed.arena, parsed.root)?;
        
        // Capture imported symbols
        file_imported_symbols.insert(module_path.clone(), resolver.imported_symbols.clone());

        // TypeChecker also needs &mut ModuleSymbols
        let mut ms_mut = module_symbols_map.get_mut(module_path).unwrap();
        let imported = file_imported_symbols.get(module_path).unwrap();
        let mut typechecker = TypeChecker::new(&mut ms_mut, imported, &module_provider, &interner);
        typechecker.check(&parsed.arena, parsed.root)?;
    }

    // 4. Lower ALL modules to VM and emit a single combined bytecode image for this project
    // Rationale: services and functions can live in multiple modules; exports must refer to
    // correct function indices within this CompiledModule. We aggregate all VM functions
    // into a single ir_vm::Module and assemble once using the public API `emit_fragments`.

    // Combined VM module (we will merge const pools and remap ConstIds accordingly)
    let mut combined_vm = crate::ir_vm::Module::new(step.project_key.name.clone());
    combined_vm.const_pool = crate::ir_core::ConstPool::new();
    // Track origin module_path for each function we append to combined_vm
    let mut combined_func_origins: Vec<String> = Vec::new();

    // Helper to insert a constant value into combined pool and return its new id
    let mut insert_const = |pool: &mut crate::ir_core::ConstPool, val: &crate::ir_core::ConstantValue| -> crate::ir_vm::types::ConstId {
        let new_id = pool.insert(val.clone());
        crate::ir_vm::types::ConstId(new_id.0)
    };

    // Accumulate VM functions from each source file, remapping ConstIds as we go
    for (module_path, parsed) in &parsed_files {
        let ms = module_symbols_map.get(module_path).unwrap();
        let imported = file_imported_symbols.get(module_path).unwrap();
        let lowerer = Lowerer::new(&parsed.arena, ms, imported, &interner);
        let program = lowerer.lower_file(parsed.root, module_path)?;

        let vm_module = core_to_vm::lower_program(&program)
            .map_err(|e| CompileError::Internal(format!("Lowering error ({}): {}", module_path, e)))?;

        // Build remap for this module's const ids
        let mut const_map: Vec<crate::ir_vm::types::ConstId> = Vec::with_capacity(vm_module.const_pool.constants.len());
        for c in &vm_module.const_pool.constants {
            let new_id = insert_const(&mut combined_vm.const_pool, c);
            const_map.push(new_id);
        }

        // Clone functions and remap any PushConst const ids
        for mut f in vm_module.functions.into_iter() {
            for instr in &mut f.body {
                if let crate::ir_vm::instr::InstrKind::PushConst(old_id) = instr.kind {
                    let mapped = const_map.get(old_id.0 as usize).cloned().unwrap_or(old_id);
                    instr.kind = crate::ir_vm::instr::InstrKind::PushConst(mapped);
                }
            }
            combined_func_origins.push(module_path.clone());
            combined_vm.functions.push(f);
        }
    }

    // Assemble once for the whole project using the public API
    let fragments = emit_fragments(&combined_vm)
        .map_err(|e| CompileError::Internal(format!("Emission error: {}", e)))?;

    // Post-fix FunctionMeta slots from VM IR (some emitters may default to 0)
    let mut fixed_function_metas = fragments.functions.clone();
    for (i, fm) in fixed_function_metas.iter_mut().enumerate() {
        if let Some(vm_func) = combined_vm.functions.get(i) {
            fm.param_slots = vm_func.param_slots;
            fm.local_slots = vm_func.local_slots;
            fm.return_slots = vm_func.return_slots;
        }
    }

    // Note: Entry point validation for the root project is now performed at the orchestrator level,
    // after compilation and before linking, using enriched debug info. We skip it here to avoid
    // double validation and mismatches with name annotations.

    // 5. Collect exports
    let mut exports = BTreeMap::new();
    for (module_path, ms) in &module_symbols_map {
        for sym in ms.type_symbols.symbols.values() {
            if sym.visibility == Visibility::Pub {
                if let Some(surface_kind) = ExportSurfaceKind::from_symbol_kind(sym.kind) {
                    exports.insert(ExportKey {
                        module_path: module_path.clone(),
                        symbol_name: interner.resolve(sym.name).to_string(),
                        kind: surface_kind,
                    }, ExportMetadata { 
                        func_idx: None,
                        is_host: sym.is_host,
                        ty: sym.ty.clone(),
                    });
                }
            }
        }
        // Build a set of public function names declared in this module (value namespace)
        let mut pub_fn_names: std::collections::HashSet<String> = std::collections::HashSet::new();
        for sym in ms.value_symbols.symbols.values() {
            if sym.visibility == Visibility::Pub {
                if let Some(surface_kind) = ExportSurfaceKind::from_symbol_kind(sym.kind) {
                    if matches!(surface_kind, ExportSurfaceKind::Function) {
                        pub_fn_names.insert(interner.resolve(sym.name).to_string());
                    }
                }
            }
        }

        for sym in ms.value_symbols.symbols.values() {
            if sym.visibility == Visibility::Pub {
                if let Some(surface_kind) = ExportSurfaceKind::from_symbol_kind(sym.kind) {
                    // Encontrar TODAS as funções no módulo VM combinado originadas deste module_path
                    // cujo nome seja igual ao do símbolo público; para cada uma, exportar
                    // tanto o nome simples quanto o alias com aridade.
                    let name_simple = interner.resolve(sym.name).to_string();
                    let mut any_found = false;
                    for (i, f) in combined_vm.functions.iter().enumerate() {
                        if combined_func_origins.get(i).map(|s| s.as_str()) != Some(module_path.as_str()) { continue; }
                        if f.name != name_simple { continue; }
                        any_found = true;
                        let meta = ExportMetadata { 
                            func_idx: Some(i as u32),
                            is_host: sym.is_host,
                            ty: sym.ty.clone(),
                        };
                        // Simple name
                        exports.insert(ExportKey {
                            module_path: module_path.clone(),
                            symbol_name: name_simple.clone(),
                            kind: surface_kind,
                        }, meta.clone());
                        // name/arity
                        let arity = f.params.len();
                        let export_name_arity = format!("{}/{}", name_simple, arity);
                        exports.insert(ExportKey {
                            module_path: module_path.clone(),
                            symbol_name: export_name_arity,
                            kind: surface_kind,
                        }, meta);
                    }
                    // Caso nada tenha sido encontrado no VM (ex.: métodos ainda não materializados),
                    // publique ao menos o nome simples sem func_idx (mantém compatibilidade de surface)
                    if !any_found {
                        exports.insert(ExportKey {
                            module_path: module_path.clone(),
                            symbol_name: name_simple.clone(),
                            kind: surface_kind,
                        }, ExportMetadata { func_idx: None, is_host: sym.is_host, ty: sym.ty.clone() });
                    }
                }
            }
        }
    }

    // 6. Collect symbols
    let project_symbols = crate::common::symbols::collect_symbols(
        &step.project_key.name,
        &module_symbols_map,
        file_manager,
        &interner,
    );

    // 6.b) Enriquecer debug_info com metadados de função (offset/len) para análise externa
    let mut dbg = fragments.debug_info.clone().unwrap_or_default();
    // Adiciona pares (func_idx, (code_offset, code_len)) ao campo function_names como anotações extras
    // Sem quebrar o formato, usamos o name como "name@offset+len" para tooling/analysis.json
    let mut enriched_function_names = Vec::new();
    for (i, (fid, name)) in fragments
        .debug_info
        .as_ref()
        .map(|d| d.function_names.clone())
        .unwrap_or_default()
        .into_iter()
        .enumerate()
    {
        let meta = &fixed_function_metas[i];
        let annotated = format!("{}@{}+{}", name, meta.code_offset, meta.code_len);
        enriched_function_names.push((fid, annotated));
    }
    if !enriched_function_names.is_empty() {
        dbg.function_names = enriched_function_names;
    }

    // 7. Collect imports from unresolved labels
    let mut imports = Vec::new();
    for (label, pcs) in fragments.unresolved_labels {
        if label.starts_with('@') {
            // Format: @dep_alias::module_path:symbol_name
            let parts: Vec<&str> = label[1..].splitn(2, "::").collect();
            if parts.len() == 2 {
                let dep_alias = parts[0].to_string();
                let rest = parts[1];
                let sub_parts: Vec<&str> = rest.rsplitn(2, ':').collect();
                if sub_parts.len() == 2 {
                    let symbol_name = sub_parts[0].to_string();
                    let module_path = sub_parts[1].to_string();
                    
                    imports.push(ImportMetadata {
                        key: ImportKey {
                            dep_alias,
                            module_path,
                            symbol_name,
                        },
                        relocation_pcs: pcs,
                    });
                }
            }
        }
    }

    Ok(CompiledModule {
        project_id: step.project_id,
        project_key: step.project_key,
        target: step.target,
        exports,
        imports,
        const_pool: fragments.const_pool,
        code: fragments.code,
        function_metas: fixed_function_metas,
        debug_info: Some(dbg),
        symbols: project_symbols,
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::fs;
    use std::path::PathBuf;
    use tempfile::tempdir;

    #[test]
    fn test_compile_root_only_project() {
        let dir = tempdir().unwrap();
        let project_dir = dir.path().to_path_buf();
        
        fs::create_dir_all(project_dir.join("src/main/modules")).unwrap();
        
        let main_code = r#"
            pub declare struct Vec2(x: int, y: int)
            
            fn add(a: int, b: int): int {
                return a + b;
            }
            
            mod fn frame(): void {
                let x = add(1, 2);
            }
        "#;
        
        fs::write(project_dir.join("src/main/modules/main.pbs"), main_code).unwrap();
        
        let project_key = ProjectKey { name: "root".to_string(), version: "0.1.0".to_string() };
        let project_id = ProjectId(0);
        let step = BuildStep {
            project_id,
            project_key: project_key.clone(),
            project_dir: project_dir.clone(),
            target: BuildTarget::Main,
            sources: vec![PathBuf::from("src/main/modules/main.pbs")],
            deps: BTreeMap::new(),
        };
        
        let mut file_manager = FileManager::new();
        let compiled = compile_project(step, &HashMap::new(), &mut file_manager).expect("Failed to compile project");
        
        assert_eq!(compiled.project_id, project_id);
        assert_eq!(compiled.target, BuildTarget::Main);
        
        // Vec2 should be exported
        let vec2_key = ExportKey {
            module_path: "".to_string(),
            symbol_name: "Vec2".to_string(),
            kind: ExportSurfaceKind::DeclareType,
        };
        assert!(compiled.exports.contains_key(&vec2_key));
        
        // frame is NOT exported (top-level fn cannot be pub in v0)
        // Wait, I put "pub fn frame" in the test code. SymbolCollector should have ignored pub.
        // Actually, SymbolCollector might error on pub fn.
    }
}