use std::iter::Peekable;

use bumpalo::Bump;

use crate::tokenize::{AttributeToken, AttributeTokenType, Token, TokenType};

#[derive(Debug, Clone, Copy)]
pub struct Document<'a> {
	pub children: &'a [Node<'a>],
}

#[derive(Debug, Clone, Copy)]
pub enum Node<'a> {
	Element(Element<'a>),
	CharacterData(Token<'a>),
	Comment(Token<'a>),
	StrayEndTag(Token<'a>),
}

#[derive(Debug, Clone, Copy)]
pub struct Element<'a> {
	pub start_tag: Token<'a>,
	pub attributes: &'a [Attribute<'a>],
	pub children: &'a [Node<'a>],
	pub end_tag: Option<Token<'a>>,
	pub end_tag_attributes: &'a [Attribute<'a>],
}

#[derive(Debug, Clone, Copy)]
pub enum Attribute<'a> {
	NameOnly {
		name: AttributeToken<'a>,
	},
	NamedArgument {
		name: AttributeToken<'a>,
		equal_sign: AttributeToken<'a>,
		value: Option<AttributeToken<'a>>,
	},
	PositionalArgument {
		value: AttributeToken<'a>,
	},
	StrayEqualSign(AttributeToken<'a>),
	NamedArgumentWithTooManyEqualSigns {
		name: AttributeToken<'a>,
		equal_signs: &'a [AttributeToken<'a>],
		value: Option<AttributeToken<'a>>,
	},
}

impl<'a> Element<'a> {
	pub fn is_self_closing(&self) -> bool {
		matches!(self.start_tag.ty, TokenType::Tag { self_closing, .. } if self_closing)
	}

	pub fn tag_name_lowercase(&self) -> Option<&str> {
		let TokenType::Tag {
			lowercase_tag_name, ..
		} = self.start_tag.ty
		else {
			return None;
		};
		Some(lowercase_tag_name)
	}

	pub fn attribute(&self, name: &str) -> Option<&Attribute<'a>> {
		self.attributes
			.iter()
			.find(|attribute| attribute.name_lowercase().is_some_and(|n| n == name))
	}
}

impl<'a> Attribute<'a> {
	pub fn name_token(&self) -> Option<&AttributeToken<'_>> {
		match self {
			Self::NamedArgument { name, .. }
			| Attribute::NameOnly { name }
			| Attribute::NamedArgumentWithTooManyEqualSigns { name, .. } => Some(name),
			Self::PositionalArgument { .. } | Self::StrayEqualSign(_) => None,
		}
	}

	pub fn name_lowercase(&self) -> Option<&str> {
		let name = self.name_token()?;
		let AttributeTokenType::Identifier { lowercase_name, .. } = name.ty else {
			return None;
		};
		Some(lowercase_name)
	}

	pub fn value_token(&self) -> Option<&AttributeToken<'a>> {
		match self {
			Self::NamedArgument { value, .. }
			| Self::NamedArgumentWithTooManyEqualSigns { value, .. } => value.as_ref(),
			Self::PositionalArgument { value } => Some(value),
			Self::NameOnly { .. } | Self::StrayEqualSign(..) => None,
		}
	}

	pub fn value_str(&self) -> Option<&str> {
		let value = self.value_token()?;
		let AttributeTokenType::String { value, .. } = value.ty else {
			return None;
		};
		Some(value)
	}
}

fn parse_named_attribute<'a>(
	attribute_tokens: &mut Peekable<impl Iterator<Item = &'a AttributeToken<'a>>>,
	name: AttributeToken<'a>,
	equal_sign: AttributeToken<'a>,
	arena: &'a Bump,
) -> Attribute<'a> {
	if let Some(equal_sign_2) = attribute_tokens
		.next_if(|t| matches!(t.ty, AttributeTokenType::EqualSign))
		.copied()
	{
		let mut equal_signs = vec![equal_sign, equal_sign_2];
		while let Some(equal_sign) = attribute_tokens
			.next_if(|t| matches!(t.ty, AttributeTokenType::EqualSign))
			.copied()
		{
			equal_signs.push(equal_sign);
		}
		let equal_signs = arena.alloc_slice_copy(&equal_signs);

		let value = attribute_tokens.next().copied();
		Attribute::NamedArgumentWithTooManyEqualSigns {
			name,
			equal_signs,
			value,
		}
	} else {
		let value = attribute_tokens.next().copied();
		Attribute::NamedArgument {
			name,
			equal_sign,
			value,
		}
	}
}

fn parse_attribute<'a>(
	attribute_tokens: &mut Peekable<impl Iterator<Item = &'a AttributeToken<'a>>>,
	arena: &'a Bump,
) -> Option<Attribute<'a>> {
	let attribute_token = *attribute_tokens.next()?;
	match attribute_token.ty {
		AttributeTokenType::Identifier { .. } => {
			let name = attribute_token;
			if let Some(equal_sign) = attribute_tokens
				.next_if(|t| matches!(t.ty, AttributeTokenType::EqualSign))
				.copied()
			{
				parse_named_attribute(attribute_tokens, name, equal_sign, arena)
			} else {
				Attribute::NameOnly { name }
			}
		}
		AttributeTokenType::String { .. } => Attribute::PositionalArgument {
			value: attribute_token,
		},
		AttributeTokenType::EqualSign => Attribute::StrayEqualSign(attribute_token),
	}
	.into()
}

fn parse_attributes<'a>(
	attribute_tokens: &'a [AttributeToken<'a>],
	arena: &'a Bump,
) -> &'a [Attribute<'a>] {
	let mut attributes = Vec::new();
	let mut attribute_tokens = attribute_tokens.iter().peekable();
	while let Some(attribute) = parse_attribute(&mut attribute_tokens, arena) {
		attributes.push(attribute);
	}
	arena.alloc_slice_copy(&attributes)
}

fn parse_element<'a>(
	tokens: &mut Peekable<impl Iterator<Item = &'a Token<'a>>>,
	stack: &[&'a str],
	arena: &'a Bump,
) -> Option<Node<'a>> {
	let start_tag = *tokens.next()?;
	let TokenType::Tag {
		is_end,
		lowercase_tag_name,
		original_tag_name: _,
		self_closing,
		attribute_tokens,
		greater_than_exists: _,
	} = start_tag.ty
	else {
		return None;
	};

	if is_end {
		return Some(Node::StrayEndTag(start_tag));
	}

	let attributes = parse_attributes(attribute_tokens, arena);

	if self_closing {
		return Some(Node::Element(Element {
			start_tag,
			attributes,
			children: &[],
			end_tag: None,
			end_tag_attributes: &[],
		}));
	}

	let mut children = Vec::new();
	let mut stack = Vec::from(stack);
	stack.push(lowercase_tag_name);

	while let Some(token) = tokens.peek() {
		match token.ty {
			TokenType::Comment { .. } => {
				let token = tokens.next().expect("a token");
				children.push(Node::Comment(*token))
			}
			TokenType::Data(_) => {
				let token = tokens.next().expect("a token");
				children.push(Node::CharacterData(*token))
			}
			TokenType::Tag {
				is_end,
				lowercase_tag_name: end_tag_name,
				attribute_tokens,
				..
			} if is_end && end_tag_name == lowercase_tag_name => {
				let token = tokens.next().expect("a token");
				let children = arena.alloc_slice_copy(&children);
				let end_tag_attributes = parse_attributes(attribute_tokens, arena);
				return Some(Node::Element(Element {
					start_tag,
					attributes,
					children,
					end_tag: Some(*token),
					end_tag_attributes,
				}));
			}
			TokenType::Tag {
				is_end,
				lowercase_tag_name,
				..
			} if is_end && stack.contains(&lowercase_tag_name) => {
				tokens.next().expect("a token");
				break;
			}
			TokenType::Tag { is_end, .. } if is_end => {
				let token = tokens.next().expect("a token");
				children.push(Node::StrayEndTag(*token))
			}
			TokenType::Tag { .. } => {
				children.push(parse_element(tokens, &[], arena).expect("an element"));
			}
		}
	}

	let children = arena.alloc_slice_copy(&children);

	Some(Node::Element(Element {
		start_tag,
		attributes,
		children,
		end_tag: None,
		end_tag_attributes: &[],
	}))
}

pub fn parse_document<'a>(tokens: &'a [Token<'a>], arena: &'a Bump) -> Document<'a> {
	let mut children = Vec::new();
	let mut tokens = tokens.iter().peekable();
	while let Some(token) = tokens.peek() {
		match token.ty {
			TokenType::Comment { .. } => {
				let token = tokens.next().expect("we just saw a token from peeking");
				children.push(Node::Comment(*token))
			}
			TokenType::Data(_) => {
				let token = tokens.next().expect("we just saw a token from peeking");
				children.push(Node::CharacterData(*token))
			}
			TokenType::Tag { .. } => {
				children.push(parse_element(&mut tokens, &[], arena).expect("an element"));
			}
		}
	}

	Document {
		children: arena.alloc_slice_copy(&children),
	}
}