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::Span;
use crate::deps::resolver::ProjectId;
use crate::frontends::pbs::ast::FileNode;
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 prometeu_bytecode::v0::{ConstantPoolEntry, DebugInfo, FunctionMeta};
use serde::{Deserialize, Serialize};
use std::collections::{BTreeMap, HashMap};
use std::path::Path;

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

#[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 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>,
}

#[derive(Debug)]
pub enum CompileError {
    Frontend(crate::common::diagnostics::DiagnosticBundle),
    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::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>
) -> Result<CompiledModule, CompileError> {
    let mut file_manager = FileManager::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, FileNode, String)> = Vec::new(); // (module_path, ast, file_stem)

    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, file_id);
        let ast = parser.parse_file()?;
        
        let mut collector = SymbolCollector::new();
        let (ts, vs) = collector.collect(&ast)?;
        
        let module_path = source_rel.parent()
            .unwrap_or(Path::new(""))
            .to_string_lossy()
            .replace('\\', "/");
        
        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 Err(existing) = ms.type_symbols.insert(sym) {
                return Err(DiagnosticBundle::error(
                    format!("Duplicate type symbol '{}' in module '{}'", existing.name, module_path),
                    Some(existing.span)
                ).into());
            }
        }
        for sym in vs.symbols.into_values() {
            if let Err(existing) = ms.value_symbols.insert(sym) {
                return Err(DiagnosticBundle::error(
                    format!("Duplicate value symbol '{}' in module '{}'", existing.name, module_path),
                    Some(existing.span)
                ).into());
            }
        }
        
        let file_stem = source_abs.file_stem().unwrap().to_string_lossy().to_string();
        parsed_files.push((module_path, ast, file_stem));
    }

    // 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.replace("src/main/modules/", "");
                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: key.symbol_name.clone(),
                        kind: key.kind.clone(),
                        namespace: match key.kind {
                            SymbolKind::Function | SymbolKind::Service => Namespace::Value,
                            SymbolKind::Struct | SymbolKind::Contract | SymbolKind::ErrorType => Namespace::Type,
                            _ => Namespace::Value,
                        },
                        visibility: Visibility::Pub,
                        ty: meta.ty.clone(),
                        is_host: meta.is_host,
                        span: Span::new(0, 0, 0),
                        origin: Some(synthetic_module_path.clone()),
                    };
                    
                    if sym.namespace == Namespace::Type {
                        ms.type_symbols.insert(sym).ok();
                    } else {
                        ms.value_symbols.insert(sym).ok();
                    }
                }
            }
        }
    }

    // 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, ast, _) in &parsed_files {
        let ms = module_symbols_map.get(module_path).unwrap();
        let mut resolver = Resolver::new(ms, &module_provider);
        resolver.resolve(ast)?;
        
        // 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);
        typechecker.check(ast)?;
    }

    // 4. Lower to IR
    let mut combined_program = crate::ir_core::Program {
        const_pool: crate::ir_core::ConstPool::new(),
        modules: Vec::new(),
        field_offsets: HashMap::new(),
        field_types: HashMap::new(),
    };

    for (module_path, ast, file_stem) 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(ms, imported);
        let program = lowerer.lower_file(ast, &file_stem)?;
        
        // Combine program into combined_program
        if combined_program.modules.is_empty() {
            combined_program = program;
        } else {
            // TODO: Real merge
        }
    }

    // 4. Emit fragments
    let vm_module = core_to_vm::lower_program(&combined_program)
        .map_err(|e| CompileError::Internal(format!("Lowering error: {}", e)))?;
    
    let fragments = emit_fragments(&vm_module)
        .map_err(|e| CompileError::Internal(format!("Emission error: {}", e)))?;

    // 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 {
                exports.insert(ExportKey {
                    module_path: module_path.clone(),
                    symbol_name: sym.name.clone(),
                    kind: sym.kind.clone(),
                }, ExportMetadata { 
                    func_idx: None,
                    is_host: sym.is_host,
                    ty: sym.ty.clone(),
                });
            }
        }
        for sym in ms.value_symbols.symbols.values() {
            if sym.visibility == Visibility::Pub {
                // Find func_idx if it's a function or service
                let func_idx = vm_module.functions.iter().position(|f| f.name == sym.name).map(|i| i as u32);
                
                exports.insert(ExportKey {
                    module_path: module_path.clone(),
                    symbol_name: sym.name.clone(),
                    kind: sym.kind.clone(),
                }, ExportMetadata { 
                    func_idx,
                    is_host: sym.is_host,
                    ty: sym.ty.clone(),
                });
            }
        }
    }

    // 6. 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,
        target: step.target,
        exports,
        imports,
        const_pool: fragments.const_pool,
        code: fragments.code,
        function_metas: fragments.functions,
        debug_info: fragments.debug_info,
    })
}

#[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_id = ProjectId { name: "root".to_string(), version: "0.1.0".to_string() };
        let step = BuildStep {
            project_id: project_id.clone(),
            project_dir: project_dir.clone(),
            target: BuildTarget::Main,
            sources: vec![PathBuf::from("src/main/modules/main.pbs")],
            deps: BTreeMap::new(),
        };
        
        let compiled = compile_project(step, &HashMap::new()).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: "src/main/modules".to_string(),
            symbol_name: "Vec2".to_string(),
            kind: SymbolKind::Struct,
        };
        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.
    }
}