use std::collections::HashMap;
use std::path::Path;

use crate::assembler::model::{Module, Node, Opcode, Symbol, Token};
use crate::assembler::{quick_hash, AssembleError};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SectionId {
    Data,
    Text,
}

/// Splits a flat (possibly multi-module) node list into separate data and
/// text streams. Segment/Include markers are dropped here — module
/// disambiguation lives on `Symbol.module`, not node position, so we don't
/// need them past this point.
pub fn split_sections(nodes: Vec<Node>) -> (Vec<Node>, Vec<Node>) {
    let mut data = Vec::new();
    let mut text = Vec::new();

    for node in nodes {
        match node.opcode() {
            Opcode::Data => data.push(node),
            Opcode::Segment | Opcode::Include => {}
            _ => text.push(node),
        }
    }

    (data, text)
}

pub fn generate_symbol_table(
    data_nodes: &[Node],
    text_nodes: &[Node],
) -> HashMap<Symbol, (u32, SectionId)> {
    let mut table = HashMap::new();

    for (i, node) in data_nodes.iter().enumerate() {
        if let Some(symbol) = node.label() {
            table.insert(symbol, (4 * i as u32, SectionId::Data));
        }
    }
    for (i, node) in text_nodes.iter().enumerate() {
        if let Some(symbol) = node.label() {
            table.insert(symbol, (4 * i as u32, SectionId::Text));
        }
    }

    table
}

/// v1 loader convention: TEXT is placed at address 0, DATA immediately
/// follows it. This is a stopgap until real relocations exist — it's the
/// same assumption your current jmp-hack already makes, just explicit now.
pub fn resolve_address((offset, section): (u32, SectionId), text_len_words: u32) -> u32 {
    match section {
        SectionId::Text => offset,
        SectionId::Data => (text_len_words * 4) + offset,
    }
}

/// Patches every symbol reference in `text_nodes` with its resolved absolute
/// address, and returns the symbol table so the caller can build the
/// executable's symbol/string tables from it.
pub fn resolve_symbols(
    data_nodes: &[Node],
    text_nodes: &mut [Node],
) -> Result<HashMap<Symbol, (u32, SectionId)>, AssembleError> {
    let symbol_table = generate_symbol_table(data_nodes, text_nodes);
    let text_len = text_nodes.len() as u32;

    for node in text_nodes.iter_mut() {
        match node.opcode() {
            Opcode::Jmp | Opcode::Call | Opcode::Jic | Opcode::Jnc | Opcode::Lli | Opcode::Lui => {
                if let Token::Symbol(symbol) = node
                    .arg(0)
                    .expect("Expected argument 0 for jump-like opcode")
                {
                    let Some(entry) = symbol_table.get(&symbol) else {
                        return Err(AssembleError::UndefinedSymbol(symbol));
                    };
                    node.tokens[0] = Token::Immediate(resolve_address(*entry, text_len));
                }
            }
            Opcode::Jez | Opcode::Jnz => {
                if let Token::Symbol(symbol) = node
                    .arg(1)
                    .expect("Expected argument 1 for jump-like opcode")
                {
                    let Some(entry) = symbol_table.get(&symbol) else {
                        return Err(AssembleError::UndefinedSymbol(symbol));
                    };
                    node.tokens[1] = Token::Immediate(resolve_address(*entry, text_len));
                }
            }
            _ => (),
        }
    }

    Ok(symbol_table)
}

// resolve_dependencies is unchanged — keep exactly as you have it.
pub fn resolve_dependencies(
    mut nodes: Vec<Node>,
    base_dir: &Path,
) -> Result<Vec<Node>, AssembleError> {
    // First we get a list of imports.
    let mut dependencies = Vec::new();
    for node in &nodes {
        if node.opcode() == Opcode::Include {
            // we want the path, and the name
            let name = if let Token::Symbol(name) = node
                .arg(0)
                .expect("Expected argument #0 for Include directive.")
            {
                name.name.clone()
            } else {
                unreachable!()
            }; //node.2.get(0).unwrap()

            let Ok(Token::StringLit(path)) = node.arg(1) else {
                unreachable!()
            };

            let full_path = base_dir.join(path);
            let canonical_path = full_path
                .canonicalize()
                .map_err(|_| AssembleError::InvalidFile(full_path.clone()))?;

            let hash = quick_hash(&canonical_path);

            dependencies.push((name, hash));
        }
    }

    let mut changes = Vec::<(u32, u32, Symbol)>::new();
    // now we resolve the symbols on all the nodes
    // we need to check all operands for unresolved signals
    for (i, node) in nodes.clone().iter().enumerate() {
        let Node {
            tokens: operands, ..
        } = node;
        for (j, token) in operands.iter().enumerate() {
            if let Token::Symbol(symbol) = token {
                for d in &dependencies {
                    if let Module::Unresolved(name) = symbol.module.clone() {
                        if name != d.0 {
                            continue;
                        }

                        let symbol = Symbol {
                            name: symbol.name.clone(),
                            module: Module::Resolved(d.1),
                        };
                        changes.push((i as u32, j as u32, symbol));
                    }
                }
            }
        }
    }

    for (i, j, symbol) in changes {
        nodes[i as usize].tokens[j as usize] = Token::Symbol(symbol);
    }

    Ok(nodes)
}
