#![allow(dead_code, unused)]

use std::{
    collections::HashSet,
    fmt, fs,
    hash::{DefaultHasher, Hash, Hasher},
    path::{Path, PathBuf},
    sync::{
        Arc, Mutex,
        mpsc::{self, Receiver, Sender},
    },
    thread,
};
use std::collections::HashMap;
use std::io::Cursor;
use common::formats::binary_dse::{DseExecutable, SymbolEntry};
use common::prelude::Instruction;
use crate::assembler::resolver2::{resolve_address, resolve_dependencies, resolve_symbols, split_sections, SectionId};

// Module declarations
#[macro_use]
pub mod macros;

#[allow(clippy::module_inception)]
pub mod codegen;
pub mod expand;

mod lexer;
mod model;
pub mod parser;
// pub mod resolver;
pub mod resolver2;



// Re-exports
pub use self::{
    codegen::codegen,
    expand::expand_pseudo_ops,
    lexer::lexer,
    model::{Module, Node, Opcode, Symbol, Token, TokenType},
    parser::{Parser, Program},
    // resolver::{create_sections, resolve_dependencies, resolve_symbols},
};

pub struct Assembler {
    src_path: PathBuf,
    result_tx: mpsc::Sender<Result<Vec<u8>, AssembleError>>,
    result_rx: Option<mpsc::Receiver<Result<Vec<u8>, AssembleError>>>,
    is_running: bool,
}

impl Assembler {
    #[must_use]
    pub fn new(src_path: impl Into<PathBuf>) -> Self {
        let (tx, rx) = mpsc::channel();
        Self {
            src_path: src_path.into(),
            result_tx: tx,
            result_rx: Some(rx),
            is_running: false,
        }
    }

    /// Start the compilation process in a separate thread
    pub fn start(&mut self, args: ()) {
        if self.is_running {
            return;
        }

        let src = self.src_path.clone();
        let tx = self.result_tx.clone();

        thread::spawn(move || match assemble(&src) {
            Ok(exe) => {
                let mut buffer = Cursor::new(Vec::new());
                match exe.write(&mut buffer) {
                    Ok(()) => {
                        tx.send(Ok(buffer.into_inner()))
                            .expect("Failed to send compilation result from worker thread");
                    }
                    Err(e) => {
                        tx.send(Err(AssembleError::Generic))
                            .expect("...");
                    }
                }
            }
            Err(err) => { /* unchanged */ }
        });

        self.is_running = true;
    }

    /// Polls the result channel to check if a compilation result is ready.
    ///
    /// This method attempts a non-blocking read (`try_recv`) from the result receiver.
    /// It returns `Some(Ok(Vec<u8>))` if the compilation successfully finished and
    /// the binary output is available, or `None` if no result was received yet.
    ///
    /// If it receives an error via `mpsc::TryRecvError::Disconnected`, it indicates that
    /// the worker thread terminated unexpectedly, returning a specific error string.
    pub fn poll(&mut self) -> Option<Result<Vec<u8>, String>> {
        if !self.is_running {
            return None;
        }

        match self
            .result_rx
            .as_ref()
            .expect("result_rx should be Some while compilation is running")
            .try_recv()
        {
            Ok(result) => {
                self.is_running = false;
                Some(result.map_err(|e| e.to_string()))
            }
            Err(mpsc::TryRecvError::Empty) => None,
            Err(mpsc::TryRecvError::Disconnected) => {
                self.is_running = false;
                Some(Err(String::from(
                    "Compilation terminated before a result was returned",
                )))
            }
        }
    }

    /// Block until compilation is complete and return the result
    pub fn output(&mut self) -> Result<Vec<u8>, String> {
        if let Ok(result) = self
            .result_rx
            .take()
            .expect("result_rx should be Some while waiting for compilation result")
            .recv()
        {
            self.is_running = false;
            result.map_err(|e| e.to_string())
        } else {
            self.is_running = false;
            Err(String::from(
                "Compilation terminated before a result was returned",
            ))
        }
    }
}

impl Assembler {}

// fn assemble(src: &Path) -> Result<Vec<Instruction>, AssembleError> {
//     let mut modules = HashSet::new();
//     let mut program = Program::new();
//
//     let hash = quick_hash(src);
//
//     if modules.contains(&hash) {
//         return Ok(vec![]);
//     }
//
//     prepare_dependency(src, &mut modules, &mut program)?;
//
//     let mut nodes = program.nodes.clone();
//
//     println!("PRE LINK");
//     for node in &nodes {
//         println!("{:?}", node);
//     }
//
//     create_sections(&mut nodes)?;
//     resolve_symbols(&mut nodes)?;
//
//     println!("Generating assembly output...");
//
//     let instructions = codegen(nodes)?;
//
//     println!("Compilation Successful");
//     Ok(instructions)
// }

fn assemble(src: &Path) -> Result<DseExecutable, AssembleError> {
    let mut modules = HashSet::new();
    let mut program = Program::new();

    prepare_dependency(src, &mut modules, &mut program)?;

    let nodes = program.nodes.clone();

    println!("PRE LINK");
    for node in &nodes {
        println!("{:?}", node);
    }

    let (data_nodes, mut text_nodes) = split_sections(nodes);
    let symbol_table = resolve_symbols(&data_nodes, &mut text_nodes)?;
    let text_len = text_nodes.len() as u32;

    println!("Generating assembly output...");
    let instructions = codegen(text_nodes)?;
    let text: Vec<u32> = instructions
        .iter()
        .map(|i| i.0).collect();

    let data = build_data_bytes(&data_nodes)?;

    // Build string table + SymbolEntry list from the resolved symbol table
    let mut string_table = Vec::new();
    let mut symbols = Vec::new();
    for (symbol, entry) in &symbol_table {
        let name_offset = string_table.len() as u32;
        string_table.extend_from_slice(symbol.name.as_bytes());
        string_table.push(0); // null terminator

        symbols.push(SymbolEntry {
            name_offset,
            value: resolve_address(*entry, text_len),
            section: match entry.1 {
                SectionId::Text => 2,
                SectionId::Data => 1,
            },
            binding: 1, // global — refine later if adding local/weak
            _pad: 0,
        });
    }

    // Entry point: first symbol to jump to is '_init'
    let entry_point = symbol_table
        .iter()
        .find(|(sym, _)| sym.name == "_init")
        .map(|(_, entry)| resolve_address(*entry, text_len))
        .ok_or(AssembleError::UndefinedSymbol(Symbol {
            name: "_init".to_string(),
            module: Module::Resolved(0),
        }))?;

    println!("Compilation Successful");

    Ok(DseExecutable::new(
        symbols.as_slice(),
        &string_table,
        &data,
        &text,
        entry_point,
    ).unwrap())
}

fn build_data_bytes(data_nodes: &[Node]) -> Result<Vec<u8>, AssembleError> {
    let mut bytes = Vec::with_capacity(data_nodes.len() * 4);
    for node in data_nodes {
        let Token::Immediate(value) = node.arg(0)? else {
            return Err(AssembleError::InvalidArg);
        };
        bytes.extend_from_slice(&value.to_le_bytes());
    }
    Ok(bytes)
}


fn prepare_dependency(
    path: &Path,
    modules: &mut HashSet<u64>,
    program: &mut Program,
) -> Result<(), AssembleError> {
    let filename = path
        .file_name()
        .and_then(|n| n.to_str())
        .expect("Failed to get file name from path");

    if let Ok(path) = path.canonicalize() {
        println!("{:20} {:20} [{}]", "Building", filename, path.display());
    }

    let src =
        fs::read_to_string(path).map_err(|_| AssembleError::InvalidFile(path.to_path_buf()))?;
    let file_hash = quick_hash(path);

    println!("{:20} {:20}", "Tokenising", filename);
    let tokens = lexer::lexer(src, file_hash)?;
    // let tokens = Lexer::new(src, file_hash).run()?;

    println!("{:20} {:20}", "Parsing", filename);
    let parsed = Parser::parse_nodes(tokens)?;

    println!("{:20} {:20}", "Resolving Deps", filename);
    // Get the parent directory of the source file to use as the base directory
    let base_dir = path
        .parent()
        .ok_or_else(|| AssembleError::InvalidFile(path.to_path_buf()))?;
    let mut nodes = expand_pseudo_ops(parsed, file_hash)?;
    nodes = resolve_dependencies(nodes, base_dir)?;

    let deps = Parser::get_dependencies(&nodes, path)?;

    println!("{:20} {:20}", "Expanding Pseudo-ops", filename);

    // add a section instruction
    nodes.insert(
        0,
        node!(None, Opcode::Segment, Token::Immediate(file_hash as u32)),
    );

    // for n in &nodes {
    //     println!("{n}");
    // }

    program.add_module(nodes);

    for dep in deps {
        println!(
            "{:20} {:20}",
            "Including",
            dep.file_name()
                .and_then(|f| f.to_str())
                .expect("Dependency path has no file name or is not valid UTF-8")
        );

        let dep_hash = quick_hash(&dep);
        if modules.insert(dep_hash) {
            prepare_dependency(dep.as_path(), modules, program)?;
        }
    }

    Ok(())
}

#[derive(Debug, Clone)]
pub enum AssembleError {
    Generic,
    UnexpectedEof,
    InvalidFile(PathBuf),
    UnexpectedToken(Token, TokenType),
    InvalidArg,
    UndefinedSymbol(Symbol),
    /// Contains the nth element missing from the instruction.
    MissingArgument(u8),
}

impl fmt::Display for AssembleError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Generic => write!(f, "Generic error"),
            Self::UnexpectedToken(tok, expected) => {
                write!(f, "Unexpected token {tok:?}, expected {expected:?}")
            }
            Self::UnexpectedEof => write!(f, "Unexpected end of file"),
            Self::InvalidFile(path) => write!(f, "Invalid file `{}`", path.display()),
            Self::InvalidArg => write!(f, "Invalid argument"),
            Self::UndefinedSymbol(symbol) => {
                write!(f, "Undefined symbol {symbol}")
            }
            Self::MissingArgument(n) => {
                write!(f, "Missing argument #{n} from instruction arguments.")
            }
        }
    }
}

fn quick_hash(value: &Path) -> u64 {
    let mut hasher = DefaultHasher::new();
    value
        .canonicalize()
        .expect("Failed to canonicalize path for quick_hash")
        .to_str()
        .hash(&mut hasher);

    hasher.finish()
}
