//! Gemeinsame Tabellenauflösung getrennter Modulübersetzungen.
use crate::{
    bytecode::{CompiledModule, Instr},
    codegen,
};
use std::collections::{HashMap, HashSet};
use tb_frontend::{
    ast::{Module, Stmt, TypeName},
    forms::FormCatalog,
    hir::{HProcKind, HTy},
    source::{locate_diagnostics, SourceUnit},
    Diagnostic, SourcePos,
};
use tb_runtime::value::TypeInit;
use tb_ui::frm::FormFile;

fn diagnostic(message: impl Into<String>) -> Diagnostic {
    Diagnostic {
        file: None,
        pos: SourcePos::default(),
        message: message.into(),
    }
}

pub fn compile_project(
    name: &str,
    units: &[SourceUnit],
    catalog: &FormCatalog,
    forms: &[FormFile],
) -> Result<CompiledModule, Vec<Diagnostic>> {
    if units.is_empty() || units.len() > u16::MAX as usize {
        return Err(vec![diagnostic("Projekt ohne Module oder zu viele Module")]);
    }
    let mut sources = Vec::new();
    let mut diagnostics = Vec::new();
    let mut names = HashSet::new();
    let parsed: Vec<_> = units
        .iter()
        .enumerate()
        .map(|(id, unit)| {
            let (module, errors) = unit.parse(id as u16, &mut sources);
            if !names.insert(unit.name.to_uppercase()) {
                let mut error = diagnostic(format!("Duplicate definition: module {}", unit.name));
                error.pos = module_pos(&module);
                error.file = unit.segments.first().map(|s| s.file.clone());
                diagnostics.push(error);
            }
            diagnostics.extend(errors);
            module
        })
        .collect();
    locate_diagnostics(&mut diagnostics, &sources);
    if !diagnostics.is_empty() {
        return Err(diagnostics);
    }
    let mut catalog = catalog.clone();
    if catalog.find("SCREEN").is_none() {
        catalog.add(
            "SCREEN",
            tb_frontend::forms::ObjectClass::Screen,
            None,
            false,
        );
    }
    for module in &parsed {
        if module
            .body
            .iter()
            .any(|s| matches!(s, Stmt::MetaForm { .. }))
            && catalog.find(&module.name).is_none()
        {
            catalog.add(
                &module.name,
                tb_frontend::forms::ObjectClass::Form,
                None,
                false,
            );
        }
    }
    let mut exports: Vec<_> = parsed
        .iter()
        .map(|m| tb_frontend::sema::export_declarations(m, &[]))
        .collect();
    // Konstantenabhängigkeiten können beliebig über Module verteilt sein.
    // Jeder erfolgreiche Durchlauf löst mindestens eine weitere Deklaration auf.
    let constant_count: usize = exports
        .iter()
        .flatten()
        .filter(|s| matches!(s, Stmt::ConstDecl { .. }))
        .count();
    for _ in 0..constant_count {
        let mut constants: HashMap<_, Vec<_>> = HashMap::new();
        for stmt in exports.iter().flatten() {
            if let Stmt::ConstDecl { items, .. } = stmt {
                constants.entry(&items[0].0).or_default().push(stmt.clone());
            }
        }
        let constants: Vec<_> = constants
            .into_values()
            .filter(|s| s.len() == 1)
            .flatten()
            .collect();
        let next: Vec<_> = parsed
            .iter()
            .map(|m| tb_frontend::sema::export_declarations(m, &constants))
            .collect();
        if next == exports {
            break;
        }
        exports = next;
    }
    let mut parts = Vec::new();
    let mut commons = Vec::new();
    for module in &parsed {
        let mut module = module.clone();
        import_declarations(&mut module, &parsed, &exports);
        let (hir, errors) = tb_frontend::sema::lower_with_forms(&module, &catalog);
        diagnostics.extend(errors);
        if let Some(hir) = hir {
            parts.push(codegen::compile(&hir));
            commons.push(hir.commons);
        }
    }
    locate_diagnostics(&mut diagnostics, &sources);
    if !diagnostics.is_empty() {
        return Err(diagnostics);
    }
    let mut result = link(name, parts, &parsed, &commons).map_err(|error| {
        let mut errors = vec![error];
        locate_diagnostics(&mut errors, &sources);
        errors
    })?;
    result.sources = sources;
    let objects = FormCatalog {
        objects: result.objects.clone(),
    };
    for form in forms {
        result
            .form_initial
            .extend(form.initial_values(&objects).map_err(|e| {
                vec![diagnostic(format!(
                    "{}: ungültige Forms-Anfangsdaten ({e})",
                    form.root.name
                ))]
            })?);
    }
    result.startup_form = forms
        .first()
        .and_then(|form| objects.find(&form.root.name).map(|(id, _)| id));
    result
        .validate()
        .map_err(|e| vec![diagnostic(e.to_string())])?;
    Ok(result)
}

/// Lokale Definition gewinnt; außerhalb ihres Moduls muss ein Name eindeutig sein.
fn type_definition<'a>(name: &str, origin: usize, all: &'a [Module]) -> Option<(usize, &'a Stmt)> {
    let mut candidates = all.iter().enumerate().flat_map(|(id, module)| {
        module.body.iter().filter_map(move |stmt| {
            matches!(stmt, Stmt::TypeDecl { name: n, .. } if n == name).then_some((id, stmt))
        })
    });
    if let Some(local) = candidates.clone().find(|(id, _)| *id == origin) {
        return Some(local);
    }
    let first = candidates.next()?;
    candidates.next().is_none().then_some(first)
}

fn same_type(
    name: &str,
    a: usize,
    b: usize,
    all: &[Module],
    visiting: &mut HashSet<(usize, usize, String)>,
) -> bool {
    let (
        Some((a, Stmt::TypeDecl { fields: af, .. })),
        Some((b, Stmt::TypeDecl { fields: bf, .. })),
    ) = (type_definition(name, a, all), type_definition(name, b, all))
    else {
        return false;
    };
    if a == b {
        return true;
    }
    if !visiting.insert((a, b, name.into())) {
        return false;
    }
    let equal = af.len() == bf.len()
        && af.iter().zip(bf).all(|((an, at), (bn, bt))| {
            an == bn
                && match (at, bt) {
                    (TypeName::Udt(an), TypeName::Udt(bn)) => {
                        an == bn && same_type(an, a, b, all, visiting)
                    }
                    _ => at == bt,
                }
        });
    visiting.remove(&(a, b, name.into()));
    equal
}

fn import_type(
    name: &str,
    origin: usize,
    target: usize,
    all: &[Module],
    imported: &mut Vec<Stmt>,
    seen: &mut HashMap<(usize, String), String>,
) -> String {
    let Some((origin, Stmt::TypeDecl { fields, pos, .. })) = type_definition(name, origin, all)
    else {
        return name.into();
    };
    if origin == target
        || (type_definition(name, target, all).is_some_and(|(id, _)| id == target)
            && same_type(name, origin, target, all, &mut HashSet::new()))
    {
        return name.into();
    }
    let key = (origin, name.to_string());
    if let Some(alias) = seen.get(&key) {
        return alias.clone();
    }
    // Öffentliche eindeutige Namen bleiben erhalten. Konfliktbehaftete Abhängigkeiten
    // bekommen einen internen Modulnamen, damit das lokale Layout unangetastet bleibt.
    let alias = if type_definition(name, target, all).is_some_and(|(id, _)| id == origin) {
        name.to_string()
    } else {
        format!("{}!{name}", all[origin].name)
    };
    seen.insert(key, alias.clone());
    let fields = fields
        .iter()
        .map(|(field, ty)| {
            let ty = match ty {
                TypeName::Udt(name) => {
                    TypeName::Udt(import_type(name, origin, target, all, imported, seen))
                }
                _ => ty.clone(),
            };
            (field.clone(), ty)
        })
        .collect();
    imported.push(Stmt::TypeDecl {
        name: alias.clone(),
        fields,
        pos: *pos,
    });
    alias
}

fn import_declarations(module: &mut Module, all: &[Module], exports: &[Vec<Stmt>]) {
    let target = all.iter().position(|m| m.name == module.name).unwrap();
    let mut procedures: HashSet<_> = module.procs.iter().map(|p| p.sig.name.clone()).collect();
    let mut local_constants = HashSet::new();
    let mut local_types = HashSet::new();
    for stmt in &module.body {
        match stmt {
            Stmt::Declare { sig, .. } => {
                procedures.insert(sig.name.clone());
            }
            Stmt::ConstDecl { items, .. } => {
                local_constants.extend(items.iter().map(|i| i.0.clone()))
            }
            Stmt::TypeDecl { name, .. } => {
                local_types.insert(name.clone());
            }
            _ => {}
        }
    }
    let mut foreign: HashMap<_, Vec<_>> = HashMap::new();
    for (origin, declarations) in exports.iter().enumerate().filter(|(id, _)| *id != target) {
        for stmt in declarations {
            let key = match stmt {
                Stmt::Declare { sig, .. } => (0, sig.name.clone()),
                Stmt::ConstDecl { items, .. } => (1, items[0].0.clone()),
                Stmt::TypeDecl { name, .. } => (2, name.clone()),
                _ => continue,
            };
            foreign.entry(key).or_default().push((origin, stmt));
        }
    }
    let mut names: Vec<_> = foreign.keys().cloned().collect();
    names.sort();
    let mut imported = Vec::new();
    let mut seen_types = HashMap::new();
    for key in names {
        let candidates = &foreign[&key];
        if candidates.len() != 1 {
            continue;
        }
        let (origin, stmt) = candidates[0];
        match stmt {
            Stmt::Declare { sig, pos } if !procedures.contains(&sig.name) => {
                let mut sig = sig.clone();
                for param in &mut sig.params {
                    if let Some(TypeName::Udt(name)) = &param.as_type {
                        param.as_type = Some(TypeName::Udt(import_type(
                            name,
                            origin,
                            target,
                            all,
                            &mut imported,
                            &mut seen_types,
                        )));
                    }
                }
                imported.push(Stmt::Declare { sig, pos: *pos });
            }
            Stmt::ConstDecl { items, .. } if !local_constants.contains(&items[0].0) => {
                imported.push(stmt.clone())
            }
            Stmt::TypeDecl { name, .. } if !local_types.contains(name) => {
                import_type(name, origin, target, all, &mut imported, &mut seen_types);
            }
            _ => {}
        }
    }
    imported.append(&mut module.body);
    let mut types = Vec::new();
    imported.retain(|stmt| {
        if matches!(stmt, Stmt::TypeDecl { .. }) {
            types.push(stmt.clone());
            false
        } else {
            true
        }
    });
    let mut known = HashSet::new();
    let mut ordered = Vec::new();
    while !types.is_empty() {
        let next = types.iter().position(|stmt| match stmt {
            Stmt::TypeDecl { fields, .. } => fields.iter().all(|(_, ty)| match ty {
                TypeName::Udt(name) => known.contains(name),
                _ => true,
            }),
            _ => unreachable!(),
        });
        let Some(index) = next else { break };
        let stmt = types.remove(index);
        if let Stmt::TypeDecl { name, .. } = &stmt {
            known.insert(name.clone());
        }
        ordered.push(stmt);
    }
    ordered.extend(types); // Sema diagnostiziert fehlende oder zyklische Typen.
    ordered.extend(imported);
    module.body = ordered;
}

fn proc_pos(module: &Module, name: &str) -> SourcePos {
    module
        .body
        .iter()
        .find_map(|stmt| match stmt {
            Stmt::Declare { sig, pos } if sig.name == name => Some(*pos),
            _ => None,
        })
        .or_else(|| {
            module
                .procs
                .iter()
                .find(|p| p.sig.name == name)
                .map(|p| p.pos)
        })
        .unwrap_or_else(|| module_pos(module))
}
fn module_pos(module: &Module) -> SourcePos {
    module
        .body
        .iter()
        .map(tb_frontend::sema::stmt_pos)
        .find(|p| p.line > 0)
        .or_else(|| module.procs.first().map(|p| p.pos))
        .unwrap_or_default()
}

fn remap_type(ty: &mut TypeInit, ids: &[u16]) {
    if let TypeInit::Udt(id) = ty {
        *id = ids[*id as usize];
    }
}
fn remap_signature(ty: &mut HTy, ids: &[u16]) {
    if let HTy::Udt(id) = ty {
        *id = ids[*id as usize];
    }
}

fn link(
    name: &str,
    mut parts: Vec<CompiledModule>,
    ast: &[Module],
    module_commons: &[Vec<tb_frontend::hir::HCommon>],
) -> Result<CompiledModule, Diagnostic> {
    let at = |pos, message| Diagnostic {
        file: None,
        pos,
        message,
    };
    let mut result = codegen::compile(&tb_frontend::analyze_source(name, "").hir.unwrap());
    result.modules = parts
        .iter()
        .map(|p| (p.name.clone(), p.option_base))
        .collect();
    result.objects = parts[0].objects.clone();
    result.strings.clear();
    result.procs.clear();
    result.option_base = parts[0].option_base;
    let mut proc_maps = Vec::new();
    let mut count = 1usize;
    let mut definitions: HashMap<String, Vec<u16>> = HashMap::new();
    for (part, module) in parts.iter().zip(ast) {
        let defined: HashSet<_> = module.procs.iter().map(|p| p.sig.name.as_str()).collect();
        let mut map = vec![0];
        for p in part.procs.iter().skip(1) {
            let id = u16::try_from(count)
                .map_err(|_| at(proc_pos(module, &p.name), "Zu viele Prozeduren".into()))?;
            count += 1;
            map.push(id);
            if defined.contains(p.name.as_str()) || p.kind == HProcKind::DefFn {
                definitions.entry(p.name.clone()).or_default().push(id);
            }
        }
        proc_maps.push(map);
    }
    // DECLARE-Platzhalter auf die tatsächliche, eindeutig bestimmte Definition binden.
    for (module_id, part) in parts.iter().enumerate() {
        let defined: HashSet<_> = ast[module_id]
            .procs
            .iter()
            .map(|p| p.sig.name.as_str())
            .collect();
        for (id, proc) in part.procs.iter().enumerate().skip(1) {
            if !defined.contains(proc.name.as_str()) && proc.kind != HProcKind::DefFn {
                if let Some(candidates) = definitions.get(&proc.name) {
                    if candidates.len() != 1 {
                        return Err(at(
                            proc_pos(&ast[module_id], &proc.name),
                            format!("Ambiguous subprogram: {}", proc.name),
                        ));
                    }
                    proc_maps[module_id][id] = candidates[0];
                }
            }
        }
    }
    let splits: Vec<_> = parts
        .iter()
        .map(|p| {
            p.procs[0]
                .code
                .iter()
                .position(|i| matches!(i, Instr::Stmt(0)))
                .map_or(0, |i| i.saturating_sub(1))
        })
        .collect();
    let mut init_starts = Vec::new();
    let mut body_starts = Vec::new();
    let mut pc = 0;
    for split in &splits {
        init_starts.push(pc);
        pc += split;
    }
    for (part, split) in parts.iter().zip(&splits) {
        body_starts.push(pc);
        pc += part.procs[0].code.len() - split - 1;
    }
    let mut main = parts[0].procs[0].clone();
    main.code.clear();
    main.name = name.into();
    let mut initializers = Vec::new();
    let mut bodies = Vec::new();
    let mut common: HashMap<_, (u16, tb_frontend::hir::HCommon)> = HashMap::new();
    for (module_id, part) in parts.iter_mut().enumerate() {
        let mut types = Vec::new();
        for udt in &part.udts {
            let mut udt = udt.clone();
            udt.name = udt.name.rsplit('!').next().unwrap().to_string();
            for ty in &mut udt.fields {
                remap_type(ty, &types);
            }
            let id = result
                .udts
                .iter()
                .position(|u| u.name == udt.name && u.fields == udt.fields)
                .unwrap_or_else(|| {
                    result.udts.push(udt);
                    result.udts.len() - 1
                });
            types.push(
                u16::try_from(id)
                    .map_err(|_| at(module_pos(&ast[module_id]), "Zu viele TYPEs".into()))?,
            );
        }
        let commons: HashMap<_, _> = module_commons[module_id]
            .iter()
            .map(|c| (c.slot, c))
            .collect();
        let mut globals = Vec::new();
        for (slot, (ty, name)) in part.globals_init.iter().zip(&part.global_names).enumerate() {
            let mut ty = ty.clone();
            remap_type(&mut ty, &types);
            let declaration = commons.get(&(slot as u16));
            let key = declaration.map(|c| (c.block.clone(), c.key.clone()));
            let mut common_ty = declaration.map(|c| c.ty.clone());
            if let Some(ty) = &mut common_ty {
                remap_signature(ty, &types);
            }
            let id = if let Some((id, previous)) = key.as_ref().and_then(|key| common.get_mut(key))
            {
                let declaration = declaration.unwrap();
                let compatible_dims = match (&previous.dims, &declaration.dims) {
                    (None, None) => true,
                    (Some(a), Some(b)) if a.is_empty() || b.is_empty() => true,
                    (Some(a), Some(b)) => {
                        a.len() == b.len()
                            && a.iter().zip(b).all(|((al, ah), (bl, bh))| {
                                al.zip(*bl).is_none_or(|(a, b)| a == b)
                                    && ah.zip(*bh).is_none_or(|(a, b)| a == b)
                            })
                    }
                    _ => false,
                };
                if Some(&previous.ty) != common_ty.as_ref() || !compatible_dims {
                    return Err(at(
                        declaration.pos,
                        format!("COMMON type or bounds mismatch: {name}"),
                    ));
                }
                if let (Some(previous), Some(current)) = (&mut previous.dims, &declaration.dims) {
                    if previous.is_empty() {
                        *previous = current.clone();
                    } else {
                        for ((lo, hi), (new_lo, new_hi)) in previous.iter_mut().zip(current) {
                            *lo = lo.or(*new_lo);
                            *hi = hi.or(*new_hi);
                        }
                    }
                }
                *id
            } else {
                let id = u16::try_from(result.globals_init.len()).map_err(|_| {
                    at(
                        module_pos(&ast[module_id]),
                        "Zu viele globale Variablen".into(),
                    )
                })?;
                result.globals_init.push(ty);
                result.global_names.push(if ast.len() == 1 {
                    name.clone()
                } else {
                    format!("{}!{name}", part.name)
                });
                if let Some(key) = key {
                    let mut declaration = (*declaration.unwrap()).clone();
                    declaration.ty = common_ty.unwrap();
                    common.insert(key, (id, declaration));
                }
                id
            };
            globals.push(id);
        }
        let string_offset = result.strings.len();
        if string_offset + part.strings.len() >= u16::MAX as usize {
            return Err(at(
                module_pos(&ast[module_id]),
                "Zu viele Stringkonstanten".into(),
            ));
        }
        result.strings.append(&mut part.strings);
        let data_offset = result.data.len() as u32;
        result.data.append(&mut part.data);
        let jump_offset = result.jump_tables.len();
        if jump_offset + part.jump_tables.len() > u16::MAX as usize {
            return Err(at(
                module_pos(&ast[module_id]),
                "Zu viele Sprungtabellen".into(),
            ));
        }
        let main_pc = |pc: u32| if (pc as usize) < splits[module_id] { init_starts[module_id] + pc as usize } else { body_starts[module_id] + pc as usize - splits[module_id] } as u32;
        for (proc_id, proc) in part.procs.iter_mut().enumerate() {
            proc.module = module_id as u16;
            for ty in &mut proc.locals_init {
                remap_type(ty, &types);
            }
            for param in &mut proc.params {
                remap_signature(&mut param.ty, &types);
            }
            if let Some(ty) = &mut proc.ret_ty {
                remap_signature(ty, &types);
            }
            if ast.len() > 1 {
                proc.name = format!("{}!{}", part.name, proc.name);
            }
            for instruction in &mut proc.code {
                use Instr::*;
                match instruction {
                    PushStr(id)
                    | Unsupported(id)
                    | LoadDynamicObjectProperty(id)
                    | StoreDynamicObjectProperty(id) => *id += string_offset as u16,
                    Input(_, _, id, _) if *id != u16::MAX => *id += string_offset as u16,
                    LoadGlobal(id) | StoreGlobal(id) | MakeRefGlobal(id) => {
                        *id = globals[*id as usize]
                    }
                    LoadArr(global, id, _, ty)
                    | DimArr(global, id, _, ty)
                    | CommonArr(global, id, _, ty)
                    | RedimArr(global, id, _, ty) => {
                        if *global {
                            *id = globals[*id as usize];
                        }
                        remap_type(ty, &types);
                    }
                    EraseSlot(true, id) => *id = globals[*id as usize],
                    GetPut(_, _, 7, id) | PushUdtId(id) => *id = types[*id as usize],
                    Call(id, _) => *id = proc_maps[module_id][*id as usize],
                    Restore(id) => *id += data_offset,
                    OnErrorGoto(pc) => *pc = main_pc(*pc),
                    Jump(pc)
                    | JumpIfFalse(pc)
                    | JumpIfTrue(pc)
                    | Gosub(pc)
                    | RetGosubTo(pc)
                    | OnErrorLocal(pc)
                    | ResumeLabel(pc)
                    | TrapDefine(_, pc)
                        if proc_id == 0 =>
                    {
                        *pc = main_pc(*pc)
                    }
                    OnJump(id, _) => {
                        if proc_id == 0 {
                            for target in &mut part.jump_tables[*id as usize] {
                                *target = main_pc(*target);
                            }
                        }
                        *id += jump_offset as u16;
                    }
                    _ => {}
                }
            }
        }
        initializers.extend_from_slice(&part.procs[0].code[..splits[module_id]]);
        bodies.extend_from_slice(
            &part.procs[0].code[splits[module_id]..part.procs[0].code.len() - 1],
        );
        result.procs.extend(part.procs.iter().skip(1).cloned());
        for mut e in part.event_procs.drain(..) {
            e.proc = proc_maps[module_id][e.proc as usize];
            result.event_procs.push(e);
        }
        result.jump_tables.append(&mut part.jump_tables);
    }
    initializers.extend(bodies);
    initializers.push(Instr::End);
    main.code = initializers;
    result.procs.insert(0, main);
    // Vollständige Signaturen zwischen DECLARE-Platzhalter und Ziel vergleichen.
    for (module, part) in parts.iter().enumerate() {
        for (id, proc) in part.procs.iter().enumerate().skip(1) {
            let target = &result.procs[proc_maps[module][id] as usize];
            if proc.ret_ty != target.ret_ty
                || proc.kind != target.kind
                || proc.params.len() != target.params.len()
                || proc
                    .params
                    .iter()
                    .zip(&target.params)
                    .any(|(a, b)| a.ty != b.ty || a.array != b.array || a.by_ref != b.by_ref)
            {
                return Err(at(
                    proc_pos(&ast[module], proc.name.rsplit('!').next().unwrap()),
                    format!("Parameter type mismatch: {}", proc.name),
                ));
            }
        }
    }
    Ok(result)
}