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

use common::prelude::Register;

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

pub fn resolve_symbols(nodes: &mut [Node]) -> Result<(), AssembleError> {
    let symbol_table = generate_symbol_table(nodes);

    for node in nodes.iter_mut() {
        match node.opcode() {
            Opcode::Jmp
            | Opcode::Jeq
            | Opcode::Jne
            | Opcode::Jgt
            | Opcode::Jge
            | Opcode::Jlt
            | Opcode::Jle
            | Opcode::Lli
            | Opcode::Lui => {
                if let Token::Symbol(symbol) = node
                    .arg(0)
                    .expect("Expected argument 0 for jump-like opcode")
                {
                    if let Some(address) = symbol_table.get(&symbol) {
                        node.tokens[0] = Token::Immediate(*address);
                    } else {
                        return Err(AssembleError::UndefinedSymbol(symbol));
                    }
                }
            }
            _ => (),
        }
    }

    Ok(())
}

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

    for (i, node) in nodes.iter().enumerate() {
        if let Some(symbol) = node.label() {
            table.insert(symbol, 4 * i as u32);
        }
    }

    table
}

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)
}

pub fn create_sections(nodes: &mut Vec<Node>) -> Result<(), AssembleError> {
    let mut res = Vec::<Node>::with_capacity(nodes.len());

    res.push(node!(None, Opcode::Segment, Token::Immediate(0)));

    for n in nodes.iter() {
        if n.opcode() == Opcode::Data {
            res.push(n.clone());
        }
    }

    let start = res.len() + 1;
    res.insert(
        0,
        node!(
            None,
            Opcode::Jmp,
            Token::Immediate(start as u32 * 4),
            Token::Register(Register::Zero)
        ),
    );
    for n in nodes.iter() {
        if !matches!(n.opcode(), Opcode::Data | Opcode::Include) {
            res.push(n.clone());
        }
    }

    *nodes = res;

    Ok(())
}
