#![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 common::{
    build::{BuildError, Builder},
    logging::{LogReceiver, Logger},
    prelude::Instruction,
};

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

#[allow(clippy::module_inception)]
// pub mod assembler;
pub mod codegen;
pub mod expand;
pub mod lexer;
pub mod model;
pub mod parser;
pub mod resolver;

// 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,

    logs_rx: LogReceiver,
}

impl From<AssembleError> for BuildError {
    fn from(err: AssembleError) -> Self {
        Self::Generic(err.to_string())
    }
}

impl Builder for Assembler {
    type Output = Vec<u8>;
    type Args = ();

    fn logs(&self) -> Vec<String> {
        self.logs_rx.logs()
    }

    #[must_use]
    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,

            // for logging
            logs_rx: LogReceiver::new(true),
        }
    }

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

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

        thread::spawn(move || {
            if let Ok(res) = assemble(&src, &logger) {
                let buffer: Vec<u8> = res
                    .iter()
                    .flat_map(|instruction| instruction.encode().to_be_bytes())
                    .collect();
                tx.send(Ok(buffer))
                    .expect("Failed to send compilation result from worker thread");
            }
        });

        self.is_running = true;
    }

    fn poll(&mut self) -> Option<Result<Self::Output, common::build::BuildError>> {
        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(std::convert::Into::into))
            }
            Err(mpsc::TryRecvError::Empty) => None,
            Err(mpsc::TryRecvError::Disconnected) => {
                self.is_running = false;
                Some(Err(BuildError::Generic(String::from(
                    "Compilation terminated before a result was returned",
                ))))
            }
        }
    }

    /// Block until compilation is complete and return the result
    fn output(&mut self) -> Result<Self::Output, common::build::BuildError> {
        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(std::convert::Into::into)
        } else {
            self.is_running = false;
            Err(BuildError::Generic(String::from(
                "Compilation terminated before a result was returned",
            )))
        }
    }
}

impl Assembler {}

fn assemble(src: &Path, logger: &Logger) -> 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, &logger)?;

    let mut nodes = program.nodes.clone();

    create_sections(&mut nodes)?;
    resolve_symbols(&mut nodes)?;

    logger.info("Generating assembly output...");

    let instructions = codegen(nodes)?;

    logger.info("Compilation Successful");
    Ok(instructions)
}

fn prepare_dependency(
    path: &Path,
    modules: &mut HashSet<u64>,
    program: &mut Program,
    logger: &Logger,
) -> 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() {
        logger.info(&format!(
            "{: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);

    logger.info(&format!("{:20} {:20}", "Tokenising", filename));
    let tokens = lexer::lexer(src, file_hash)?;

    logger.info(&format!("{:20} {:20}", "Parsing", filename));
    let parsed = Parser::parse_nodes(tokens)?;

    logger.info(&format!("{: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)?;

    logger.info(&format!("{: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 {
        logger.info(&format!(
            "{: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, logger)?;
        }
    }

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