use anyhow::{bail, Result};
use debug_ignore::DebugIgnore;
use itertools::Itertools;
use rust_embed::RustEmbed;
use std::collections::{HashMap, HashSet};
use std::fmt;
use std::fs;
use std::path::Path;

use crate::arg;
use crate::lexer;
use crate::parser;

pub trait ToCode: fmt::Display {}

pub trait ByteResolvable<T> {
    fn resolve_number(&self, data: &mut T) -> Result<u16>;
    fn resolve_bytes(&self, data: &mut T) -> Result<Vec<u8>>;
    fn resolve_string(&self, data: &mut T) -> Result<String>;
}

pub const KB: u16 = 1024;

pub const PROGRAM_SIZE: u16 = 32 * KB;
pub type Program = [u8; PROGRAM_SIZE as usize];

#[derive(RustEmbed)]
#[folder = "lib/"]
#[prefix = "$lib/"]
pub struct Lib;

#[derive(Debug)]
pub struct Macro {
    pub args: Vec<String>,
    pub body: parser::ast::Body,
}

#[derive(Debug)]
pub struct Assembler {
    pub dir: String,
    pub ast: DebugIgnore<parser::ast::AST>,
    pub program: DebugIgnore<Program>,
    pub offset: u16,
    pub body_stack: Vec<parser::ast::Node>,
    pub constants: HashMap<String, arg::Arg>,
    pub macros: DebugIgnore<HashMap<String, Macro>>,
    pub includes: HashSet<String>,
    // pub tmp: Vec<u8>,
    // pub tmp_offset: u16,
    // pub tmp_enabled: bool,
}

impl Assembler {
    pub fn new(dir: String, ast: parser::ast::AST) -> Self {
        Self {
            dir,
            body_stack: ast.body.iter().rev().cloned().collect(),
            ast: ast.into(),
            program: [0; PROGRAM_SIZE as usize].into(),
            offset: 0,
            constants: HashMap::new(),
            macros: HashMap::new().into(),
            includes: HashSet::new(),
            // tmp: vec![],
            // tmp_offset: 0,
            // tmp_enabled: false,
        }
    }

    pub fn resolve_node(&mut self) -> Result<Option<Vec<u8>>> {
        let node = match self.body_stack.pop() {
            Some(x) => x,
            None => bail!("Body stack empty. No more nodes to resolve"),
        };

        match node {
            parser::ast::Node::Comment(_) => Ok(None),
            parser::ast::Node::Label(x) => {
                self.constants.insert(x, arg::Arg::Number(self.offset));

                Ok(None)
            }
            parser::ast::Node::Call { name, arg } => {
                let opcode: u8 = match name.as_str() {
                    "nop" => 0x00,
                    "push" => 0x01,
                    "pop" => 0x02,
                    "ts" => 0x03,
                    "tsr" => 0x04,
                    "tss" => 0x05,
                    "tlr" => 0x06,
                    "tlrc" => 0x07,
                    "tls" => 0x08,
                    "dbg" => 0x09,
                    "alu" => 0x0a,
                    "get" => 0x0b,
                    "set" => 0x0c,
                    _ => bail!("Unknown opcode: {}", name),
                };
                let a = match arg {
                    Some(x) => x.resolve_number(self)?,
                    None => 0,
                };
                if a == 0 {
                    Ok(Some(vec![opcode | 0b10000000]))
                } else {
                    Ok(Some(vec![opcode, (a >> 8) as u8, a as u8]))
                }
            }
            parser::ast::Node::MacroCall { name, args } => match name.as_str() {
                "debug" => {
                    for arg in args {
                        let assembly = arg.to_string().replace('\n', "\\n");
                        let num = arg.resolve_number(self)?;
                        let bytes = arg.resolve_bytes(self)?;

                        println!("{}", assembly);
                        println!("  => 0b{:0>16b}", num);
                        println!("  => {}", num);
                        println!("  => 0x{:0>4x}", num);
                        println!(
                            "  => [{}]",
                            bytes.iter().map(|n| format!("{:0>2x}", n)).join(" ")
                        );
                    }

                    Ok(None)
                }
                "debug_assembler" => {
                    println!("{:#?}", self);

                    Ok(None)
                }
                "define" => {
                    let name = match &args[0] {
                        arg::Arg::Variable(x) | arg::Arg::String(x) => x,
                        _ => {
                            bail!("First argument of define macro needs to be a variable")
                        }
                    };
                    self.constants.insert(name.clone(), args[1].clone());

                    Ok(None)
                }
                "define_eval" => {
                    let name = match &args[0] {
                        arg::Arg::Variable(x) => x,
                        _ => {
                            bail!("First argument of define macro needs to be a variable")
                        }
                    };
                    let a = arg::Arg::Number(args[1].resolve_number(self)?);
                    self.constants.insert(name.clone(), a);

                    Ok(None)
                }
                "delete" => {
                    for arg in args {
                        match arg {
                            arg::Arg::Variable(x) => {
                                self.constants.remove(&x);
                            }
                            _ => bail!("Arguments need to be variables"),
                        }
                    }

                    Ok(None)
                }
                "macro" => {
                    let name = match &args[0] {
                        arg::Arg::Variable(x) | arg::Arg::String(x) => x,
                        _ => {
                            bail!("First argument of define macro needs to be a literal-like value")
                        }
                    };
                    let args = args[1..]
                        .iter()
                        .map(|a| match a {
                            arg::Arg::Variable(x) => Ok(x.to_owned()),
                            _ => bail!("Macro arguments need to be variables"),
                        })
                        .collect::<Result<Vec<_>>>()?;

                    let mut depth: usize = 1;
                    let body: Vec<_> = self
                        .body_stack
                        .iter()
                        .rev()
                        .take_while(|n| match n {
                            parser::ast::Node::MacroCall { name, args: _ } => match name.as_str() {
                                "macro" => {
                                    depth += 1;
                                    true
                                }
                                "endmacro" => {
                                    depth -= 1;
                                    depth != 0
                                }
                                _ => true,
                            },
                            _ => true,
                        })
                        .cloned()
                        .collect();
                    self.body_stack
                        .drain((self.body_stack.len() - body.len() - 1)..)
                        .for_each(drop);
                    self.macros.insert(name.clone(), Macro { args, body });

                    Ok(None)
                }
                "if" => {
                    let cond = (args[0].resolve_number(self)? & 1) == 1;

                    let mut if_depth: usize = 1;
                    let body: Vec<_> = self
                        .body_stack
                        .iter()
                        .rev()
                        .take_while(|n| match n {
                            parser::ast::Node::MacroCall { name, args: _ } => match name.as_str() {
                                "if" => {
                                    if_depth += 1;
                                    true
                                }
                                "endif" => {
                                    if_depth -= 1;
                                    if_depth != 0
                                }
                                _ => true,
                            },
                            _ => true,
                        })
                        .cloned()
                        .collect();
                    self.body_stack
                        .drain((self.body_stack.len() - body.len() - 1)..)
                        .for_each(drop);

                    let mut iter = body.iter().enumerate();
                    if let Some(pos) = loop {
                        if let Some(pos) = iter.next() {
                            if let parser::ast::Node::MacroCall { name, args: _ } = pos.1 {
                                match name.as_str() {
                                    "if" => {
                                        break None;
                                    }
                                    "else" => {
                                        break Some(pos);
                                    }
                                    _ => {}
                                }
                            }
                        } else {
                            break None;
                        }
                    } {
                        self.body_stack.extend(if cond {
                            body[..pos.0].to_vec()
                        } else {
                            body[(pos.0 + 1)..].to_vec()
                        });
                    } else if cond {
                        self.body_stack.extend(body.into_iter().rev());
                    }

                    Ok(None)
                }
                "org" => {
                    self.offset = args[0].resolve_number(self)?;

                    Ok(None)
                }
                "org_add" => {
                    self.offset = self.offset.wrapping_add(args[0].resolve_number(self)?);

                    Ok(None)
                }
                "org_sub" => {
                    self.offset = self.offset.wrapping_sub(args[0].resolve_number(self)?);

                    Ok(None)
                }
                "bytes" => {
                    let mut data: Vec<u8> = vec![];
                    for arg in args {
                        data.append(&mut arg.resolve_bytes(self)?);
                    }

                    if data.is_empty() {
                        Ok(None)
                    } else {
                        Ok(Some(data))
                    }
                }
                "embed" => {
                    let arg = args[0].resolve_string(self)?;
                    let data = if Path::new(&arg).starts_with("$lib") {
                        match Lib::get(arg.as_str()) {
                            Some(x) => x.data.to_vec(),
                            None => bail!("Virtual lib file couldn't be found with path: {}", arg),
                        }
                    } else {
                        let path = Path::new(&self.dir).join(arg);
                        fs::read(path)?
                    };

                    Ok(Some(data))
                }
                "requires" => {
                    let arg = args[0].resolve_string(self)?;
                    if !self.includes.contains(&arg) {
                        bail!("File requires include of not included file: {}", arg);
                    }
                    Ok(None)
                }
                "include" => {
                    let arg = args[0].resolve_string(self)?;
                    let source = if Path::new(&arg).starts_with("$lib") {
                        match Lib::get(arg.as_str()) {
                            Some(x) => std::str::from_utf8(&x.data)?.to_string(),
                            None => bail!("Virtual lib file couldn't be found with path: {}", arg),
                        }
                    } else {
                        let path = Path::new(&self.dir).join(&arg);
                        fs::read_to_string(path)?
                    };

                    self.includes.insert(arg);
                    let body = parser::parse(lexer::lex(&source)?)?.body;
                    self.body_stack
                        .extend(body.into_iter().rev().collect::<Vec<_>>());

                    Ok(None)
                }
                _ => match self.macros.get(&name) {
                    Some(m) => {
                        if m.args.len() != args.len() {
                            bail!(
                                "Macro call argument signature does not match: (.macro {name}{macro_args}) != (.{name}{call_args})",
                                macro_args = if m.args.is_empty() {
                                    "".to_string()
                                } else {
                                    format!(" {}", m.args.join(", "))
                                },
                                call_args = if args.is_empty() {
                                    "".to_string()
                                } else {
                                    format!(" {}", args.iter().map(|a| a.to_string()).join(", "))
                                }
                            );
                        }

                        for (i, arg) in args.iter().enumerate() {
                            self.constants.insert(m.args[i].to_owned(), arg.to_owned());
                        }

                        if !args.is_empty() {
                            self.body_stack.push(parser::ast::Node::MacroCall {
                                name: "delete".to_string(),
                                args: args
                                    .iter()
                                    .enumerate()
                                    .map(|(i, _)| arg::Arg::Variable(m.args[i].to_owned()))
                                    .collect::<Vec<_>>(),
                            });
                        }
                        self.body_stack
                            .append(&mut m.body.iter().rev().cloned().collect());

                        Ok(None)
                    }
                    None => bail!("Unknown macro: {}", name),
                },
            },
        }
    }

    pub fn assemble(&mut self) -> Result<Program> {
        let offset_name = "OFFSET".to_string();
        while !self.body_stack.is_empty() {
            self.constants
                .insert(offset_name.clone(), arg::Arg::Number(self.offset));
            if let Some(x) = self.resolve_node()? {
                if self.offset + (x.len() as u16) > 32 * 1024 {
                    bail!("Offset out of bounds: 0x{:0>4x} > 0x8000", self.offset);
                }
                for byte in x {
                    self.program[self.offset as usize] = byte;
                    self.offset += 1;
                }
            }
        }

        Ok(*self.program)
    }
}