use std::collections::HashMap;

use super::{
	Result,
	lexer::{Token, TokenKind},
	util::{SpanParserError, TokenError},
};

use TokenKind as K;

#[derive(Clone, Debug)]
struct Macro<'s> {
	contents: Vec<Token<'s>>,
	args: Vec<(usize, usize)>,
	num_args: usize,
}

struct PreProcessor<'s> {
	idx: usize,
	tokens: Vec<Token<'s>>,
	macros: HashMap<&'s str, Macro<'s>>,
}
impl<'s> PreProcessor<'s> {
	fn new(tokens: Vec<Token<'s>>) -> Self {
		let macros = HashMap::new();
		Self {
			idx: 0,
			tokens,
			macros,
		}
	}

	fn next(&mut self) -> Token<'s> {
		if self.idx >= self.tokens.len() {
			Token {
				span: self.tokens[self.idx - 1].span,
				content: "",
				kind: K::Eof,
			}
		} else {
			self.idx += 1;
			self.tokens[self.idx - 1]
		}
	}

	fn expect(&mut self, kind: TokenKind) -> Result<Token<'s>> {
		let token = self.next();
		if token.kind == kind {
			Ok(token)
		} else {
			token.token_err(format!("expected {kind}, got {}", token.kind))
		}
	}

	fn parse_int(&mut self) -> Result<usize> {
		let token = self.expect(K::Integer)?;
		token.content.parse().span_err(token.span)
	}

	fn read_macro(&mut self) -> Result<Macro<'s>> {
		let mut nest_counter = 0;
		let mut contents = vec![];
		let mut args = vec![];
		let mut num_args = 0;
		loop {
			let token = self.next();
			match token.kind {
				K::Argument => {
					let num = self.parse_int()?;
					args.push((num, contents.len()));
					num_args = num_args.max(num);
				}
				K::MacroDefine => {
					nest_counter += 1;
					contents.push(token);
				}
				K::MacroEnd => {
					if nest_counter > 0 {
						nest_counter -= 1;
						contents.push(token);
					} else {
						break;
					}
				}
				K::Eof => return token.token_err("macro definition was not closed"),
				_ => contents.push(token),
			}
		}

		args.sort_by(|(_, l), (_, r)| l.cmp(r).reverse());

		Ok(Macro {
			contents,
			args,
			num_args,
		})
	}

	fn read_macros(&mut self) -> Result<Vec<Token<'s>>> {
		let mut buffer = vec![];
		self.idx = 0;
		loop {
			let token = self.next();
			if token.kind == K::Eof {
				break;
			}
			if token.kind != K::MacroDefine {
				buffer.push(token);
				continue;
			}
			let name = self.expect(K::Identifier)?.content;
			let content = self.read_macro()?;
			self.macros.insert(name, content);
		}
		Ok(buffer)
	}

	fn pass(&mut self) -> Result<Vec<Token<'s>>> {
		let mut buffer = self.read_macros()?;
		let mut idx = 0;
		loop {
			if idx >= buffer.len() {
				break;
			}

			let token = buffer[idx];

			if token.kind != K::Identifier {
				idx += 1;
				continue;
			}

			// TODO: remove clone
			let Some(mac) = self.macros.get(token.content).cloned() else {
				idx += 1;
				continue;
			};

			let mut args = vec![];
			for n in 1..=mac.num_args {
				let arg = buffer[idx + n];
				if matches!(arg.kind, K::Eof | K::LineSeparator) {
					return arg.token_err("missing macro argument");
				}
				args.push(arg);
			}

			let mut content = mac.contents;
			for (n, idx) in mac.args {
				// this works since mac.args is stored by idx descending
				content.insert(idx, args[n - 1]);
			}

			let len = content.len();
			buffer.splice(idx..=(idx + mac.num_args), content);
			idx += len + 1;
		}

		Ok(buffer)
	}

	fn process(&mut self) -> Result<Vec<Token<'s>>> {
		loop {
			let result = self.pass()?;
			if self.tokens == result {
				return Ok(result);
			}
			self.tokens = result;
		}
	}
}

pub fn process(tokens: Vec<Token<'_>>) -> Result<Vec<Token<'_>>> {
	PreProcessor::new(tokens).process()
}