use std::{
	collections::HashMap,
	sync::atomic::{AtomicUsize, Ordering::Relaxed},
};

use bumpalo::Bump;

use crate::{
	parse::{Attribute, Document, Element, Node},
	tokenize::{AttributeTokenType, TokenType},
};

static NEXT_ID: AtomicUsize = AtomicUsize::new(0);

pub struct SymbolTableSet<'a> {
	symbol_tables: HashMap<SymbolTableId, ScopedSymbolTable<'a>>,
}

pub struct ScopedSymbolTable<'a> {
	pub modules: HashMap<ModuleSymbolId, ModuleSymbol<'a>>,
	pub constants: HashMap<ConstantSymbolId, ConstantSymbol<'a>>,
	pub heads: HashMap<HeadSymbolId, HeadSymbol<'a>>,
	pub bodies: HashMap<HeadSymbolId, HeadSymbol<'a>>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct SymbolTableId(pub usize);
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct ModuleSymbolId(pub usize);
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct ConstantSymbolId(pub usize);
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct HeadSymbolId(pub usize);

pub struct ModuleSymbol<'a> {
	pub id: ModuleSymbolId,
	pub name: Option<&'a str>,
	pub symbol_table: Option<&'a ScopedSymbolTable<'a>>,
	pub definition_site: Option<Element<'a>>,
}

pub struct ConstantSymbol<'a> {
	pub id: ConstantSymbolId,
	pub name: Option<&'a str>,
	pub ty: Option<ConstantType<'a>>,
	pub value: Option<ConstantExpression<'a>>,
	pub definition_site: Option<Element<'a>>,
}

pub struct HeadSymbol<'a> {
	pub id: HeadSymbolId,
	pub contents: Option<ConstantSymbolId>,
	pub definition_site: Option<Element<'a>>,
}

pub enum ConstantType<'a> {
	String,
	Block {
		direct_children_tags: Option<&'a [&'a str]>,
		transitive_children_tags: Option<&'a [&'a str]>,
	},
}

pub enum ConstantValue<'a> {
	String(&'a str),
}

pub enum ConstantExpression<'a> {
	Name(&'a str),
	Value(ConstantValue<'a>),
	Fragment(&'a [ConstantSymbolId]),
	Block {
		tag_name_lowercase: &'a str,
		attributes: HashMap<&'a str, Option<ConstantSymbolId>>,
		child: ConstantSymbolId,
	},
}

fn next_id() -> usize {
	NEXT_ID.fetch_add(1, Relaxed)
}

fn element_attributes<'a>(
	symbol_table: &mut ScopedSymbolTable<'a>,
	element: &'a Element<'a>,
) -> HashMap<&'a str, Option<ConstantSymbolId>> {
	let mut attributes = HashMap::new();
	for attribute in element.attributes {
		let Some(name) = attribute.name_lowercase() else {
			continue;
		};
		let value = attribute.value_str();

		let id = value.map(|value| {
			let id = ConstantSymbolId(next_id());
			symbol_table.constants.insert(
				id,
				ConstantSymbol {
					id,
					name: None,
					ty: Some(ConstantType::String),
					value: Some(ConstantExpression::Value(ConstantValue::String(value))),
					definition_site: None,
				},
			);
			id
		});
		attributes.insert(name, id);
	}
	attributes
}

fn fragment<'a>(
	arena: &'a Bump,
	symbol_table: &mut ScopedSymbolTable<'a>,
	children: &'a [Node<'a>],
) -> ConstantExpression<'a> {
	let children = arena.alloc_slice_copy(
		&children
			.iter()
			.filter_map(|node| node_to_constant(arena, symbol_table, node))
			.collect::<Box<_>>(),
	);
	ConstantExpression::Fragment(children)
}

fn node_to_constant<'a>(
	arena: &'a Bump,
	symbol_table: &mut ScopedSymbolTable<'a>,
	node: &'a Node<'a>,
) -> Option<ConstantSymbolId> {
	match node {
		Node::CharacterData(token) => {
			let id = ConstantSymbolId(next_id());
			let TokenType::Data(value) = token.ty else {
				return None;
			};
			symbol_table.constants.insert(
				id,
				ConstantSymbol {
					id,
					name: None,
					ty: Some(ConstantType::String),
					value: Some(ConstantExpression::Value(ConstantValue::String(value))),
					definition_site: None,
				},
			);
			Some(id)
		}
		Node::Element(element) if element.tag_name_lowercase() == Some("kuht-var") => {
			let id = ConstantSymbolId(next_id());
			let name = element
				.attribute("name")
				.and_then(|attribute| attribute.value_str())?;
			let fragment_id = ConstantSymbolId(next_id());
			let fragment = fragment(arena, symbol_table, element.children);
			symbol_table.constants.insert(
				fragment_id,
				ConstantSymbol {
					id,
					name: None,
					// TODO fragment type
					ty: None,
					value: Some(fragment),
					definition_site: None,
				},
			);

			symbol_table.constants.insert(
				id,
				ConstantSymbol {
					id,
					name: None,
					// TODO determine this type
					ty: None,
					value: Some(ConstantExpression::Name(name)),
					// TODO we need some way to get where these come from
					definition_site: None,
				},
			);
			Some(id)
		}
		Node::Element(element) => {
			let id = ConstantSymbolId(next_id());
			let tag_name_lowercase = element.tag_name_lowercase()?;
			let attributes = element_attributes(symbol_table, element);
			let fragment_id = ConstantSymbolId(next_id());
			let fragment = fragment(arena, symbol_table, element.children);
			symbol_table.constants.insert(
				fragment_id,
				ConstantSymbol {
					id,
					name: None,
					// TODO fragment type
					ty: None,
					value: Some(fragment),
					definition_site: None,
				},
			);

			symbol_table.constants.insert(
				id,
				ConstantSymbol {
					id,
					name: None,
					// TODO type here
					ty: None,
					value: Some(ConstantExpression::Block {
						tag_name_lowercase,
						attributes,
						child: fragment_id,
					}),
					definition_site: None,
				},
			);
			Some(id)
		}
		Node::StrayEndTag(_) | Node::Comment(_) => None,
	}
}

// TODO replace <import> with <module "base" file="../base.kuht" />
// TODO support builtin modules
// TODO also support <import "foo::bar" as="baz" />
fn add_import<'a>(symbol_table: &mut ScopedSymbolTable<'a>, element: &Element<'a>) {
	let id = ModuleSymbolId(next_id());
	let name = if let Some(Attribute::PositionalArgument { value }) = element.attributes.first()
		&& let AttributeTokenType::String { value, .. } = value.ty
	{
		Some(value)
	} else {
		None
	};
	let definition_site = Some(*element);

	symbol_table.modules.insert(
		id,
		ModuleSymbol {
			id,
			name,
			symbol_table: None,
			definition_site,
		},
	);
}

fn add_constant<'a>(symbol_table: &mut ScopedSymbolTable<'a>, element: &'a Element<'a>) {
	let id = ConstantSymbolId(next_id());
	let name = element.attribute("name").and_then(|attr| attr.value_str());
	let value = element.attribute("value").and_then(|attr| attr.value_str());
	symbol_table.constants.insert(
		id,
		ConstantSymbol {
			id,
			name,
			ty: Some(ConstantType::String),
			value: value.map(|value| ConstantExpression::Value(ConstantValue::String(value))),
			definition_site: Some(*element),
		},
	);
}

fn add_block<'a>(
	arena: &'a Bump,
	symbol_table: &mut ScopedSymbolTable<'a>,
	element: &'a Element<'a>,
) {
	let id = ConstantSymbolId(next_id());
	let name = element.attribute("name").and_then(|attr| attr.value_str());
	let fragment = fragment(arena, symbol_table, element.children);
	symbol_table.constants.insert(
		id,
		ConstantSymbol {
			id,
			name,
			// fill this in
			ty: None,
			value: Some(fragment),
			definition_site: Some(*element),
		},
	);
}

fn add_head<'a>(
	arena: &'a Bump,
	symbol_table: &mut ScopedSymbolTable<'a>,
	element: &'a Element<'a>,
) {
	let id = HeadSymbolId(next_id());
	let fragment_id = ConstantSymbolId(next_id());
	let fragment = fragment(arena, symbol_table, element.children);
	symbol_table.constants.insert(
		fragment_id,
		ConstantSymbol {
			id: fragment_id,
			name: None,
			ty: None,
			value: Some(fragment),
			definition_site: None,
		},
	);
	symbol_table.heads.insert(
		id,
		HeadSymbol {
			id,
			contents: Some(fragment_id),
			definition_site: Some(*element),
		},
	);
}

fn add_body<'a>(
	arena: &'a Bump,
	symbol_table: &mut ScopedSymbolTable<'a>,
	element: &'a Element<'a>,
) {
	let id = HeadSymbolId(next_id());
	let fragment_id = ConstantSymbolId(next_id());
	let fragment = fragment(arena, symbol_table, element.children);
	symbol_table.constants.insert(
		fragment_id,
		ConstantSymbol {
			id: fragment_id,
			name: None,
			ty: None,
			value: Some(fragment),
			definition_site: None,
		},
	);
	symbol_table.bodies.insert(
		id,
		HeadSymbol {
			id,
			contents: Some(fragment_id),
			definition_site: Some(*element),
		},
	);
}

pub fn symbol_table<'a>(arena: &'a Bump, document: Document<'a>) -> ScopedSymbolTable<'a> {
	let mut symbol_table = ScopedSymbolTable {
		modules: HashMap::new(),
		constants: HashMap::new(),
		heads: HashMap::new(),
		bodies: HashMap::new(),
	};

	for node in document.children {
		match node {
			Node::Element(element) if element.tag_name_lowercase() == Some("import") => {
				add_import(&mut symbol_table, element)
			}
			Node::Element(element) if element.tag_name_lowercase() == Some("const") => {
				add_constant(&mut symbol_table, element)
			}
			Node::Element(element) if element.tag_name_lowercase() == Some("block") => {
				add_block(arena, &mut symbol_table, element)
			}
			Node::Element(element) if element.tag_name_lowercase() == Some("head") => {
				add_head(arena, &mut symbol_table, element)
			}
			Node::Element(element) if element.tag_name_lowercase() == Some("body") => {
				add_body(arena, &mut symbol_table, element)
			}
			_ => {}
		}
	}

	symbol_table
}

impl<'a> SymbolTableSet<'a> {
	pub fn new() -> Self {
		Self {
			symbol_tables: HashMap::new(),
		}
	}

	pub fn add_document(&mut self, arena: &'a Bump, document: Document<'a>) {
		self.symbol_tables
			.insert(SymbolTableId(next_id()), symbol_table(arena, document));
	}
}