use rust_decimal::Decimal;

use crate::{
	ast::{Command, Expression, ExpressionSuffix, List, ListItem, Program, Statement},
	tokens::{Token, TokenType},
};

#[derive(Debug, Clone, PartialEq)]
pub enum Ty {
	Any,
	Union(Vec<Ty>),
	Pair(Box<Ty>, Box<Ty>),
	Function {
		args: Vec<Ty>,
		return_ty: Box<Ty>,
	},
	String {
		value: Option<String>,
	},
	Number {
		minimum: Option<Decimal>,
		maximum: Option<Decimal>,
	},
	Atom {
		value: Option<String>,
	},
}

#[derive(Debug, Clone, PartialEq)]
pub struct Register {
	id: usize,
	ty: Ty,
}

#[derive(Debug, Clone, PartialEq)]
pub enum IrStatement {
	LoadImmediateNumber {
		destination: Register,
		immediate: Decimal,
	},
	LoadImmediateString {
		destination: Register,
		immediate: String,
	},
	LoadImmediateAtom {
		destination: Register,
		immediate: String,
	},
	CreateConsPair {
		destination: Register,
		left: Register,
		right: Register,
	},
	Copy {
		destination: Register,
		source: Register,
	},
	EvaluateAtom {
		destination: Register,
		atom: Register,
	},
	CallFunction {
		destination: Register,
		function: Register,
		args: Vec<Register>,
	},
}

pub struct Scope {
	registers: Vec<Register>,
	statements: Vec<IrStatement>,
}

fn new_register(scope: &mut Scope, ty: Ty) -> Register {
	let id = scope.registers.len();
	let register = Register { id, ty };
	scope.registers.push(register.clone());
	register
}

pub fn compile_program(program: &Program) -> Scope {
	let mut scope = Scope {
		registers: Vec::new(),
		statements: Vec::new(),
	};

	for command in &*program.commands {
		compile_command(&mut scope, command);
	}

	scope
}

fn compile_command(scope: &mut Scope, command: &Command) {
	match command {
		Command::Statement(statement) => compile_statement(scope, statement),
		Command::Expression(_) => unimplemented!("top-level lists are not currently supported"),
	};
}

fn compile_statement(scope: &mut Scope, statement: &Statement) {
	let args = statement
		.args
		.iter()
		.map(|expr| compile_expression(scope, expr))
		.collect();

	let function_atom = new_register(scope, Ty::Atom { value: None });
	let function = new_register(scope, Ty::Any);
	let destination = new_register(scope, Ty::Any);

	let TokenType::Identifier(function_name) = &statement.function_token.ty else {
		panic!("expected a function name")
	};

	scope.statements.push(IrStatement::LoadImmediateAtom {
		destination: function_atom.clone(),
		immediate: function_name.to_string(),
	});

	scope.statements.push(IrStatement::EvaluateAtom {
		destination: function.clone(),
		atom: function_atom,
	});

	scope.statements.push(IrStatement::CallFunction {
		destination,
		function,
		args,
	});
}

fn compile_expression(scope: &mut Scope, expression: &Expression) -> Register {
	let expression_suffix = compile_expression_suffix(scope, &expression.suffix);

	if expression.is_quoted() {
		let head = new_register(
			scope,
			Ty::Pair(
				Box::new(Ty::Function {
					args: vec![Ty::Any],
					return_ty: Box::new(Ty::Any),
				}),
				Box::new(Ty::Any),
			),
		);
		let tail = new_register(scope, Ty::Any);
		let quote = new_register(
			scope,
			Ty::Atom {
				value: Some("QUOTE".to_string()),
			},
		);
		let null = new_register(
			scope,
			Ty::Atom {
				value: Some("NIL".to_string()),
			},
		);

		scope.statements.push(IrStatement::LoadImmediateAtom {
			destination: quote.clone(),
			immediate: "QUOTE".to_string(),
		});
		scope.statements.push(IrStatement::LoadImmediateAtom {
			destination: null.clone(),
			immediate: "NIL".to_string(),
		});
		scope.statements.push(IrStatement::CreateConsPair {
			destination: tail.clone(),
			left: expression_suffix,
			right: null,
		});
		scope.statements.push(IrStatement::CreateConsPair {
			destination: head.clone(),
			left: quote,
			right: tail,
		});

		head
	} else {
		expression_suffix
	}
}

fn compile_expression_suffix(scope: &mut Scope, suffix: &ExpressionSuffix) -> Register {
	match suffix {
		ExpressionSuffix::Nothing => panic!("invalid value"),
		ExpressionSuffix::Atom(atom) => compile_atom(scope, atom),
		ExpressionSuffix::List(list) => compile_list(scope, list),
	}
}

fn compile_atom(scope: &mut Scope, suffix: &Token) -> Register {
	match &suffix.ty {
		TokenType::Identifier(atom) => {
			let register = new_register(
				scope,
				Ty::Atom {
					value: Some(atom.to_string()),
				},
			);

			scope.statements.push(IrStatement::LoadImmediateAtom {
				destination: register.clone(),
				immediate: atom.to_string(),
			});

			register
		}
		TokenType::String {
			content,
			terminated: _,
		} => {
			let register = new_register(
				scope,
				Ty::String {
					value: Some(content.to_string()),
				},
			);

			scope.statements.push(IrStatement::LoadImmediateString {
				destination: register.clone(),
				immediate: content.to_string(),
			});

			register
		}
		TokenType::Number(number) => {
			let register = new_register(
				scope,
				Ty::Number {
					minimum: Some(*number),
					maximum: Some(*number),
				},
			);

			scope.statements.push(IrStatement::LoadImmediateNumber {
				destination: register.clone(),
				immediate: *number,
			});

			register
		}
		_ => {
			panic!("expected a string, number, or identifier");
		}
	}
}

fn compile_list(scope: &mut Scope, list: &List) -> Register {
	let contains_dot = list
		.items
		.iter()
		.any(|item| matches!(item, ListItem::Dot(_)));
	let mut head = (!contains_dot).then(|| {
		let register = new_register(
			scope,
			Ty::Atom {
				value: Some("NIL".to_string()),
			},
		);
		scope.statements.push(IrStatement::LoadImmediateAtom {
			destination: register.clone(),
			immediate: "NIL".to_string(),
		});
		register
	});

	for item in list.items.iter().rev() {
		let ListItem::Expression(item) = item else {
			continue;
		};
		let Some(head) = head.as_mut() else {
			head = Some(compile_expression(scope, item));
			continue;
		};

		let expression = compile_expression(scope, item);
		let register = new_register(
			scope,
			Ty::Pair(Box::new(expression.ty.clone()), Box::new(head.ty.clone())),
		);

		scope.statements.push(IrStatement::CreateConsPair {
			destination: register.clone(),
			left: expression,
			right: head.clone(),
		});
		*head = register;
	}

	head.expect("list should be non-empty")
}