use crate::{
	ast::{Block, Expression, Program, Return, Statement, Type},
	diagnostic::Diagnostic,
	verbose
};
use askama::Template;
use std::{borrow::Cow, fmt::Write as _};

mod assign;
mod conditional;
mod decl;
mod expect_ty;
mod expr;
mod function;
mod resolve_ty;
mod state;
mod ty;

use assign::expand_assign;
use conditional::{expand_if, expand_while};
use decl::expand_decl;
use expect_ty::expect_tuple_exact;
use expr::expand_expr_with_type;
use function::expand_function;
use state::State;
use ty::expand_type;

pub enum CompileResult {
	Ok {
		main_rs: syn::File,
		funs_rs: syn::File,
		types_rs: syn::File
	},
	InternalErr(String, syn::Error)
}

mod tpl {
	use crate::ast::{Generics, Output};
	use askama::Template;
	use proc_macro2::Ident;
	use std::borrow::Cow;

	mod filters {
		use std::fmt::Display;
		pub fn if_some<T: Display>(opt: &Option<T>) -> askama::Result<String> {
			Ok(if let Some(value) = opt.as_ref() {
				value.to_string()
			} else {
				String::new()
			})
		}
	}

	#[derive(Template)]
	#[template(path = "main.rs.j2", escape = "none")]
	pub(super) struct MainRs<'a> {
		pub(super) inputs: Vec<Input<'a>>,
		pub(super) outputs: &'a [Output]
	}

	#[derive(Template)]
	#[template(path = "types.rs.j2", escape = "none")]
	pub(super) struct TypesRs<'a> {
		pub(super) typedefs: Vec<TypeDef<'a>>
	}

	pub(super) struct TypeDef<'a> {
		pub(super) ident: &'a Ident,
		pub(super) generics: Option<&'a Generics<Ident>>,
		pub(super) inner_ty: Cow<'static, str>
	}

	impl<'a> TypeDef<'a> {
		pub(super) fn new(
			ident: &'a Ident,
			generics: Option<&'a Generics<Ident>>,
			inner_ty: Cow<'static, str>
		) -> Self {
			Self {
				ident,
				generics,
				inner_ty
			}
		}
	}

	pub(super) struct Input<'a> {
		pub(super) ident: &'a Ident,
		short_long: &'static str,
		quote: char
	}

	impl<'a> Input<'a> {
		pub(super) fn new(ident: &'a Ident) -> Self {
			let (short_long, quote) = match ident.to_string() {
				str if str.len() == 1 && str != "h" && str != "n" && str != "v" => {
					("short", '\'')
				},
				_ => ("long", '"')
			};
			Self {
				ident,
				short_long,
				quote
			}
		}
	}
}

#[allow(clippy::write_with_newline)]
pub fn compile(input: &str) -> Result<CompileResult, Diagnostic> {
	let program: Program = syn::parse_str(input)?;
	let mut state = State::new(program.types, program.functions);

	let main_rs = tpl::MainRs {
		inputs: program
			.inputs
			.iter()
			.map(|arg| tpl::Input::new(&arg.ident))
			.collect(),
		outputs: &program.outputs
	};
	let main_rs = main_rs.render().unwrap();

	let types_rs = tpl::TypesRs {
		typedefs: state
			.types()
			.into_iter()
			.map(|ty| {
				Ok(tpl::TypeDef::new(
					&ty.ident,
					ty.generics.as_ref(),
					expand_type(&state, &ty.ty)?
				))
			})
			.collect::<Result<Vec<_>, Diagnostic>>()?
	};
	let types_rs = types_rs.render().unwrap();

	let mut buf = String::new();
	buf += "#![no_implicit_prelude]\n";
	buf += "#![allow(unused_mut, while_true)]\n\n";
	buf += "use crate::types::*;\n\n";

	buf += "pub(crate) fn run(crate::Values {\n";
	for input in program.inputs {
		write!(buf, "        mut {},\n", input.ident).unwrap();
		state.add_variable(input.ident, Type::Int(Default::default()));
	}
	buf += "}: crate::Values) -> crate::Output {\n";
	for s in program.stmts {
		buf += "    ";
		buf += &expand_stmt(&mut state, s)?;
		buf += "\n";
	}
	buf += "    crate::Output {\n";
	for output in program.outputs {
		for ident in output.idents {
			write!(buf, "        {ident}: ").unwrap();
			buf += &expand_expr_with_type(
				&state,
				&Type::Int(Default::default()),
				Expression::Ident(ident)
			)?;
			buf += ",\n";
		}
	}
	buf += "    }\n";
	buf += "}\n\n";

	while let Some(idx) = state.next_function_instance() {
		state.with_function_instance(idx, |fun, generics, ident| {
			if verbose() {
				eprint!("[compile] Expanding function instance {ident} ...");
			}
			let res = expand_function(&mut buf, &state, fun, generics, ident);
			if verbose() {
				eprintln!("done");
			}
			res
		})?;
	}

	Ok(CompileResult::Ok {
		main_rs: match syn::parse_file(&main_rs) {
			Ok(main_rs) => main_rs,
			Err(err) => return Ok(CompileResult::InternalErr(main_rs, err))
		},
		funs_rs: match syn::parse_file(&buf) {
			Ok(funs_rs) => funs_rs,
			Err(err) => return Ok(CompileResult::InternalErr(buf, err))
		},
		types_rs: match syn::parse_file(&types_rs) {
			Ok(types_rs) => types_rs,
			Err(err) => return Ok(CompileResult::InternalErr(types_rs, err))
		}
	})
}

fn expand_stmt(
	state: &mut State<'_>,
	stmt: Statement
) -> Result<Cow<'static, str>, Diagnostic> {
	Ok(match stmt {
		Statement::Decl(decl) => expand_decl(state, decl)?.into(),
		Statement::Assign(assign) => expand_assign(state, assign)?.into(),
		Statement::If(if_expr) => expand_if(state, if_expr)?.into(),
		Statement::CoinFlip { head, prob, tail } => {
			let mut buf = format!("match ::stdlib::Coin::flip({}) {{", prob);
			buf += "::stdlib::Coin::Head => ";
			buf += &expand_block(state, head)?;
			buf += ", ::stdlib::Coin::Tail => ";
			buf += &expand_block(state, tail)?;
			buf += "}";
			buf.into()
		},
		Statement::While(while_expr) => expand_while(state, while_expr)?.into(),
		Statement::Block(block) => expand_block(state, block)?.into(),
		Statement::Return(Return {
			return_token, expr, ..
		}) => {
			let return_ty = state.return_ty().ok_or_else(|| {
				Diagnostic::new(
					return_token.span,
					"Cannot return",
					"Return statements may only be used in function bodies"
				)
			})?;
			match expr {
				Some(expr) => {
					format!("return {};", expand_expr_with_type(state, return_ty, expr)?)
						.into()
				},
				None => {
					expect_tuple_exact(
						return_ty,
						0,
						return_token.span,
						return_token.span
					)?;
					"return;".into()
				}
			}
		},
		Statement::Continue(cont) => {
			if !state.is_loop() {
				return Err(Diagnostic::new(
					cont.continue_token.span,
					"Not a loop",
					"You may only use `continue` in a loop body"
				));
			}
			"continue;".into()
		},
		Statement::Break(brake) => {
			if !state.is_loop() {
				return Err(Diagnostic::new(
					brake.break_token.span,
					"Not a loop",
					"You may only use `break` in a loop body"
				));
			}
			"break;".into()
		}
	})
}

fn expand_block(state: &State<'_>, block: Block) -> Result<String, Diagnostic> {
	let mut sub = state.substate();
	let mut buf = "{".to_owned();
	for s in block.stmts {
		buf += &expand_stmt(&mut sub, s)?;
	}
	buf += "}";
	Ok(buf)
}