eschac 0.6.1

//! Building chess positions.
//!
//! [`Setup`] is a builder for the [`Position`] type.

use crate::bitboard::*;
use crate::board::*;
use crate::lookup;
use crate::position::*;

/// **A builder type for chess positions.**
///
/// This type is useful to edit a position without having to ensure it stays legal at every step.
/// It must be validated and converted to a [`Position`] using [`Setup::into_position`] before
/// generating moves.
///
/// ## Text description
///
/// Forsyth-Edwards Notation (FEN) is typically used to describe chess positions as text. eschac
/// uses a slightly different notation, which simply removes the last two fields of the FEN record
/// (i.e. the halfmove clock and the fullmove number) as the [`Position`] type does not keep track
/// of those. For example, the starting position is recorded as
/// `rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq -`.
///
/// [`Display`](core::fmt::Display) and [`FromStr`](alloc::str::FromStr) are purposely not
/// implemented on [`Setup`] as there does not exist a truely canonical format for writing chess
/// positions. The format described above is implemented in
/// [`from_text_record`](Setup::from_text_record) and [`to_text_record`](Setup::to_text_record).
#[derive(Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct Setup {
	bitboards: [Bitboard; 5],
}

impl Setup {
	/// Creates an empty board, i.e. `8/8/8/8/8/8/8/8 w - -`.
	#[inline]
	pub fn new() -> Self {
		Self::from_bitboards_inner([Bitboard(0); 5])
	}

	/// Reads a position from a text record
	/// (see [`from_ascii_record`](Setup::from_ascii_record)).
	#[inline]
	pub fn from_text_record(record: &str) -> Result<Self, ParseRecordError> {
		Self::from_ascii_record(record.as_bytes())
	}

	/// Reads a position from a text record.
	pub fn from_ascii_record(record: &[u8]) -> Result<Self, ParseRecordError> {
		let mut s = record.iter().copied().peekable();
		(|| {
			let mut setup = Setup::new();
			let mut accept_empty_square = true;
			let mut rank: u8 = 7;
			let mut file: u8 = 0;
			for c in s.by_ref() {
				if c == b'/' {
					(file == 8).then_some(())?;
					rank = rank.checked_sub(1)?;
					file = 0;
					accept_empty_square = true;
				} else if (b'1'..=b'8').contains(&c) && accept_empty_square {
					file = file + c - b'0';
					(file <= 8).then_some(())?;
					accept_empty_square = false;
				} else if c == b' ' {
					break;
				} else {
					let role = Role::from_ascii(c)?;
					let color = match c.is_ascii_uppercase() {
						true => Color::White,
						false => Color::Black,
					};
					setup = setup.set(
						Square::from_coords(File::new(file)?, Rank::new(rank)?),
						Some(Piece { role, color }),
					);
					file += 1;
					accept_empty_square = true;
				}
			}
			(rank == 0).then_some(())?;
			(file == 8).then_some(())?;

			setup = match s.next()? {
				b'w' => setup.set_turn(Color::White),
				b'b' => setup.set_turn(Color::Black),
				_ => return None,
			};
			(s.next()? == b' ').then_some(())?;

			if s.next_if_eq(&b'-').is_none() {
				if s.next_if_eq(&b'K').is_some() {
					setup = setup.set_castling_rights(Color::White, Castle::Short, true);
				}
				if s.next_if_eq(&b'Q').is_some() {
					setup = setup.set_castling_rights(Color::White, Castle::Long, true);
				}
				if s.next_if_eq(&b'k').is_some() {
					setup = setup.set_castling_rights(Color::Black, Castle::Short, true);
				}
				if s.next_if_eq(&b'q').is_some() {
					setup = setup.set_castling_rights(Color::Black, Castle::Long, true);
				}
			}
			(s.next()? == b' ').then_some(())?;

			match s.next()? {
				b'-' => (),
				file => {
					setup = setup.set_en_passant_target_square(Some(Square::from_coords(
						File::from_ascii(file)?,
						Rank::from_ascii(s.next()?)?,
					)));
				}
			}
			s.next().is_none().then_some(())?;

			Some(setup)
		})()
		.ok_or_else(|| ParseRecordError {
			byte: record.len() - s.len(),
		})
	}

	/// Returns the inner bitboard representation of the position.
	///
	/// In order, those bitboards contain:
	/// - squares occupied by pawns, bishops & queens
	/// - squares occupied by knights, bishops & kings
	/// - squares occupied by rooks, queens & kings
	/// - squares occupied by black pieces
	/// - color to play, castling rights and en passant target square[^1]
	///
	/// [^1]: The color to play is represented with the `B1` square: if it's set then it's
	/// black's turn, else it's white's turn. The castling rights are represented by the
	/// squares of the corresponding rook. Finally, there can be one en passant target square
	/// on either the third or sixth rank.
	#[inline]
	pub fn bitboards(&self) -> [Bitboard; 5] {
		self.bitboards
	}

	/// Returns the piece on the square, if any.
	#[inline]
	pub fn get(&self, square: Square) -> Option<Piece> {
		let setup = self;
		Some(Piece {
			role: setup.role(square)?,
			color: setup.color(square),
		})
	}

	/// Returns the color to play.
	#[inline]
	pub fn turn(&self) -> Color {
		let [_, _, _, _, x] = self.bitboards();
		Color::new(((x.0 >> 1) & 1) as u8).unwrap()
	}

	/// Returns `true` if castling is available for the given color and side.
	#[inline]
	pub fn castling_rights(&self, color: Color, castle: Castle) -> bool {
		let [_, _, _, _, x] = self.bitboards();
		let square = match (color, castle) {
			(Color::White, Castle::Short) => sq("h1"),
			(Color::White, Castle::Long) => sq("a1"),
			(Color::Black, Castle::Short) => sq("h8"),
			(Color::Black, Castle::Long) => sq("a8"),
		};
		x.contains(square)
	}

	/// Returns the optional en passant target square.
	#[inline]
	pub fn en_passant_target_square(&self) -> Option<Square> {
		let [_, _, _, _, x] = self.bitboards();
		(x & (Rank::Third.bitboard() | Rank::Sixth.bitboard())).first()
	}

	/// Returns the text record of the position
	/// (see [`write_text_record`](Self::write_text_record)).
	pub fn to_text_record(&self) -> TextRecord {
		TextRecord {
			setup: self.clone(),
		}
	}

	/// Writes the text record of the position.
	#[inline]
	pub fn write_text_record<W: core::fmt::Write>(&self, w: &mut W) -> core::fmt::Result {
		for rank in Rank::all().into_iter().rev() {
			let mut count = 0;
			for file in File::all() {
				match self.get(Square::from_coords(file, rank)) {
					Some(piece) => {
						if count > 0 {
							w.write_char(char::from_u32('0' as u32 + count).unwrap())?;
						}
						count = 0;
						w.write_char(match piece.color {
							Color::White => piece.role.to_char_uppercase(),
							Color::Black => piece.role.to_char_lowercase(),
						})?;
					}
					None => {
						count += 1;
					}
				}
			}
			if count > 0 {
				w.write_char(char::from_u32('0' as u32 + count).unwrap())?;
			}
			if rank != Rank::First {
				w.write_char('/')?;
			}
		}

		w.write_char(' ')?;

		w.write_char(match self.turn() {
			Color::White => 'w',
			Color::Black => 'b',
		})?;

		w.write_char(' ')?;

		let mut no_castle_available = true;
		if self.castling_rights(Color::White, Castle::Short) {
			w.write_char('K')?;
			no_castle_available = false;
		}
		if self.castling_rights(Color::White, Castle::Long) {
			w.write_char('Q')?;
			no_castle_available = false;
		}
		if self.castling_rights(Color::Black, Castle::Short) {
			w.write_char('k')?;
			no_castle_available = false;
		}
		if self.castling_rights(Color::Black, Castle::Long) {
			w.write_char('q')?;
			no_castle_available = false;
		}
		if no_castle_available {
			w.write_char('-')?;
		}

		w.write_char(' ')?;

		match self.en_passant_target_square() {
			Some(sq) => {
				w.write_char(sq.file().to_char())?;
				w.write_char(sq.rank().to_char())?;
			}
			None => {
				w.write_char('-')?;
			}
		}

		Ok(())
	}

	/// Sets the occupancy of a square.
	#[inline]
	pub fn set(&self, square: Square, piece: Option<Piece>) -> Self {
		let mask = square.bitboard();
		let [mut a, mut b, mut c, mut d, e] = self.bitboards();
		if let Some(piece) = piece {
			match piece.role {
				Role::Pawn => {
					a |= mask;
					b &= !mask;
					c &= !mask;
				}
				Role::Knight => {
					a &= !mask;
					b |= mask;
					c &= !mask;
				}
				Role::Bishop => {
					a |= mask;
					b |= mask;
					c &= !mask;
				}
				Role::Rook => {
					a &= !mask;
					b &= !mask;
					c |= mask;
				}
				Role::Queen => {
					a |= mask;
					b &= !mask;
					c |= mask;
				}
				Role::King => {
					a &= !mask;
					b |= mask;
					c |= mask;
				}
			}
			match piece.color {
				Color::White => d &= !mask,
				Color::Black => d |= mask,
			}
		}
		Self::from_bitboards_inner([a, b, c, d, e])
	}

	/// Sets the color to play.
	#[inline]
	pub fn set_turn(&self, color: Color) -> Self {
		let [a, b, c, d, mut e] = self.bitboards();
		e.0 &= !2;
		e.0 |= (color as u8 as u64) << 1;
		Self::from_bitboards_inner([a, b, c, d, e])
	}

	/// Switches the color to play.
	#[inline]
	pub fn switch_turn(&self) -> Self {
		let [a, b, c, d, mut e] = self.bitboards();
		e.0 ^= 2;
		Self::from_bitboards_inner([a, b, c, d, e])
	}

	/// Sets the castling rights for the given color and side.
	#[inline]
	pub fn set_castling_rights(&self, color: Color, castle: Castle, value: bool) -> Self {
		let [a, b, c, d, mut e] = self.bitboards();
		let square = match (color, castle) {
			(Color::White, Castle::Short) => sq("h1"),
			(Color::White, Castle::Long) => sq("a1"),
			(Color::Black, Castle::Short) => sq("h8"),
			(Color::Black, Castle::Long) => sq("a8"),
		};
		match value {
			false => e.remove(square),
			true => e.insert(square),
		}
		Self::from_bitboards_inner([a, b, c, d, e])
	}

	/// Sets the en passant target square.
	#[inline]
	pub fn set_en_passant_target_square(&self, square: Option<Square>) -> Self {
		let [a, b, c, d, mut e] = self.bitboards();
		let mask = Rank::Third.bitboard() | Rank::Sixth.bitboard();
		e &= !mask;
		if let Some(square) = square {
			e |= mask & square.bitboard();
		}
		Self::from_bitboards_inner([a, b, c, d, e])
	}

	/// Returns the mirror image of the position.
	///
	/// The mirror of a position is the position obtained after reflecting the placement of pieces
	/// horizontally, inverting the color of all the pieces, inverting the turn, and reflecting the
	/// castling rights as well as the en passant target square.
	///
	/// For example, the mirror image of `rnbqkbnr/pppppppp/8/8/4P3/8/PPPP1PPP/RNBQKBNR b Kq e3`
	/// is `rnbqkbnr/pppp1ppp/8/4p3/8/8/PPPPPPPP/RNBQKBNR w Qk e6`.
	#[inline]
	pub fn mirror(&self) -> Self {
		let mask = sq("b1").bitboard();
		let [a, b, c, d, e] = self.bitboards();
		let [a, b, c, d] = [a, b, c, d].map(Bitboard::mirror);
		let turn = (e & mask) ^ mask;
		let e = (e & !mask).mirror() | turn;
		Self::from_bitboards_inner([a, b, c, (a | b | c) ^ d, e])
	}

	/// Checks for the validity of the position, returning a set of potential reasons.
	pub fn illegal_reasons(&self) -> IllegalPositionReasons {
		let mut reasons = IllegalPositionReasons::new();

		let [a, b, c, _, _] = self.bitboards();
		let blockers = a | b | c;
		let pieces = self.pieces();

		if Color::all()
			.into_iter()
			.any(|color| pieces.get(color).king().is_empty())
		{
			reasons.insert(IllegalPositionReason::MissingKing);
		}

		if pieces.get(!self.turn()).king().any(|enemy_king| {
			let pieces = pieces.get(self.turn());
			!(lookup::king(enemy_king) & pieces.king()
				| lookup::bishop(enemy_king, blockers) & (pieces.queen() | pieces.bishop())
				| lookup::rook(enemy_king, blockers) & (pieces.queen() | pieces.rook())
				| lookup::knight(enemy_king) & pieces.knight()
				| lookup::pawn_attack(!self.turn(), enemy_king) & pieces.pawn())
			.is_empty()
		}) {
			reasons.insert(IllegalPositionReason::HangingKing);
		}

		if Color::all().into_iter().any(|color| {
			!(pieces.get(color).pawn() & (Rank::First.bitboard() | Rank::Eighth.bitboard()))
				.is_empty()
		}) {
			reasons.insert(IllegalPositionReason::PawnOnBackRank);
		}

		if Color::all().into_iter().any(|color| {
			let pieces = pieces.get(color);
			(pieces.pawn().len()
				+ pieces.queen().len().saturating_sub(1)
				+ pieces.bishop().len().saturating_sub(2)
				+ pieces.knight().len().saturating_sub(2)
				+ pieces.rook().len().saturating_sub(2))
			.saturating_sub(8)
				+ pieces.king().len().saturating_sub(1)
				> 0
		}) {
			reasons.insert(IllegalPositionReason::TooMuchMaterial);
		}

		if Color::all().into_iter().any(|color| {
			Castle::all().into_iter().any(|castle| {
				self.castling_rights(color, castle)
					&& !(pieces
						.get(color)
						.king()
						.contains(Square::from_coords(File::E, color.home_rank()))
						&& pieces.get(color).rook().contains(Square::from_coords(
							castle.rook_origin_file(),
							color.home_rank(),
						)))
			})
		}) {
			reasons.insert(IllegalPositionReason::InvalidCastlingRights);
		}

		if self.en_passant_target_square().is_some_and(|en_passant| {
			let (target_rank, pawn_rank) = match self.turn() {
				Color::White => (Rank::Sixth, Rank::Fifth),
				Color::Black => (Rank::Third, Rank::Fourth),
			};
			let pawn_square = Square::from_coords(en_passant.file(), pawn_rank);
			en_passant.rank() != target_rank
				|| blockers.contains(en_passant)
				|| !pieces.get(!self.turn()).pawn().contains(pawn_square)
		}) {
			reasons.insert(IllegalPositionReason::InvalidEnPassant);
		}

		if self.en_passant_target_square().is_some_and(|en_passant| {
			let x = en_passant.trans_unchecked(!self.turn().forward());
			let y = pieces.get(!self.turn()).bishop() | pieces.get(!self.turn()).queen();
			let king_square = pieces.get(self.turn()).king().first().unwrap();
			let mask = lookup::line(king_square, x);
			lookup::bishop(king_square, blockers).contains(x)
				&& !(y & mask & lookup::bishop(king_square, blockers ^ x.bitboard())).is_empty()
		}) {
			reasons.insert(IllegalPositionReason::ImpossibleEnPassantPin);
		}

		reasons
	}

	/// Tries to convert the position into the [`Position`] type.
	///
	/// Some unreachable positions are rejected, see [`IllegalPositionReason`] for details.
	#[inline]
	pub fn into_position(self) -> Result<Position, IllegalPosition> {
		let reasons = self.illegal_reasons();
		if reasons.is_empty() {
			Ok(Position::from_setup_inner(self))
		} else {
			Err(IllegalPosition {
				setup: self,
				reasons,
			})
		}
	}

	/// Converts the position into the [`Position`] type without checking if it is valid.
	///
	/// # Safety
	///
	/// If the position is not valid, using it might cause undefined behavior.
	#[inline]
	pub unsafe fn into_position_unchecked(self) -> Position {
		Position::from_setup_inner(self)
	}

	/// Returns the position encoded by the given bitboards, or `None` if the bitboards are not
	/// compatible with the format described in [`bitboards`](Self::bitboards).
	#[inline]
	pub fn from_bitboards(bitboards: [Bitboard; 5]) -> Option<Self> {
		if sanity_check(bitboards) {
			Some(Self { bitboards })
		} else {
			None
		}
	}

	/// Returns the position encoded by the given bitboards, without checking if the format is
	/// respected.
	///
	/// # Safety
	///
	/// If the bitboards are ill-formed, then using the position obtained from this method might
	/// cause undefined behavior.
	#[inline]
	pub unsafe fn from_bitboards_unchecked(bitboards: [Bitboard; 5]) -> Self {
		Self::from_bitboards_inner(bitboards)
	}

	/// Returns the squares occupied by each type & color of piece.
	#[inline]
	pub(crate) fn pieces(&self) -> ByColor<ByRole<Bitboard>> {
		let [a, b, c, d, _] = self.bitboards();
		let t = [
			a & !b & !c,
			!a & b & !c,
			a & b & !c,
			!a & !b & c,
			a & !b & c,
			!a & b & c,
		];
		ByColor([ByRole(t.map(|x| x & !d)), ByRole(t.map(|x| x & d))])
	}

	#[inline]
	pub(crate) fn from_bitboards_inner(bitboards: [Bitboard; 5]) -> Self {
		debug_assert!(sanity_check(bitboards));
		Self { bitboards }
	}

	#[inline]
	pub(crate) fn role(&self, square: Square) -> Option<Role> {
		let mask = square.bitboard();
		let [a, b, c, _, _] = self.bitboards();
		let bit0 = (a & mask).0 >> square.index();
		let bit1 = (b & mask).0 >> square.index();
		let bit2 = (c & mask).0 >> square.index();
		Role::new((bit0 | bit1 << 1 | bit2 << 2) as u8)
	}

	#[inline]
	pub(crate) fn color(&self, square: Square) -> Color {
		let [_, _, _, x, _] = self.bitboards();
		match x.contains(square) {
			false => Color::White,
			true => Color::Black,
		}
	}
}

#[inline]
pub(crate) fn sanity_check(bitboards: [Bitboard; 5]) -> bool {
	let [a, b, c, d, e] = bitboards;
	(a & b & c).is_empty()
		&& (d & !(a | b | c)).is_empty()
		&& (e & !Bitboard(0x8100FF0000FF0083)).is_empty()
		&& (e & Bitboard(0x0000FF0000FF0000)).len() <= 1
}

impl core::fmt::Debug for Setup {
	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> Result<(), core::fmt::Error> {
		f.debug_tuple("Setup")
			.field(&self.to_text_record())
			.finish()
	}
}

impl Default for Setup {
	#[inline]
	fn default() -> Self {
		Self::new()
	}
}

pub struct TextRecord {
	setup: Setup,
}
impl core::fmt::Display for TextRecord {
	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
		self.setup.write_text_record(f)
	}
}
impl core::fmt::Debug for TextRecord {
	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
		write!(f, "\"{self}\"")
	}
}

/// An invalid position.
///
/// This is an illegal position that can't be represented with the [`Position`] type.
#[derive(Debug)]
pub struct IllegalPosition {
	pub setup: Setup,
	pub reasons: IllegalPositionReasons,
}
impl core::fmt::Display for IllegalPosition {
	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
		use core::fmt::Write;
		let setup = &self.setup;
		write!(f, "`{}` is illegal:", setup.to_text_record())?;
		let mut first = true;
		for reason in self.reasons {
			if !first {
				f.write_char(',')?;
			}
			first = false;
			f.write_char(' ')?;
			f.write_str(reason.to_str())?;
		}
		Ok(())
	}
}
impl core::error::Error for IllegalPosition {}

/// A set of [`IllegalPositionReason`]s.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct IllegalPositionReasons {
	set: u8,
}

impl IllegalPositionReasons {
	/// Returns `true` if the given reason appears in the set.
	pub fn contains(&self, reason: IllegalPositionReason) -> bool {
		self.set & (1 << reason as u8) != 0
	}

	/// Returns `true` if the set is empty.
	pub fn is_empty(&self) -> bool {
		self.set == 0
	}

	fn new() -> Self {
		Self { set: 0 }
	}

	fn insert(&mut self, reason: IllegalPositionReason) {
		self.set |= 1 << reason as u8;
	}
}

impl core::fmt::Debug for IllegalPositionReasons {
	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
		f.debug_list().entries(*self).finish()
	}
}

impl Iterator for IllegalPositionReasons {
	type Item = IllegalPositionReason;
	fn next(&mut self) -> Option<Self::Item> {
		if self.set == 0 {
			None
		} else {
			let i = self.set.trailing_zeros();
			self.set &= !(1 << i);
			IllegalPositionReason::new(i as u8)
		}
	}
}

/// Reasons for an illegal position to be rejected by eschac.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum IllegalPositionReason {
	/// One of the colors misses its king.
	MissingKing,
	/// The opponent's king is in check.
	HangingKing,
	/// There is a pawn on the first or eighth rank.
	PawnOnBackRank,
	/// Some castling rights are invalid regarding the positions of the rooks and kings.
	InvalidCastlingRights,
	/// The en passant target square is invalid, either because:
	/// - it is not on the correct rank
	/// - it is occupied
	/// - it is not behind an opponent's pawn
	InvalidEnPassant,
	/// There is an impossible number of pieces.
	///
	/// Enforcing this enables to put an upper limit on the number of legal moves on any position,
	/// allowing to reduce the size of [`Moves`].
	TooMuchMaterial,
	/// The pawn that can be taken en passant is pinned diagonally to the playing king.
	///
	/// This can't happen on a legal position, as it would imply that the king could have be taken
	/// on that move. Enforcing this makes it unnecessary to test for a discovery check on the
	/// diagonal when taking en passant.
	ImpossibleEnPassantPin,
}

impl IllegalPositionReason {
	fn new(index: u8) -> Option<Self> {
		let reason = match index {
			0 => Self::MissingKing,
			1 => Self::HangingKing,
			2 => Self::PawnOnBackRank,
			3 => Self::InvalidCastlingRights,
			4 => Self::InvalidEnPassant,
			5 => Self::TooMuchMaterial,
			6 => Self::ImpossibleEnPassantPin,
			_ => return None,
		};
		debug_assert!(reason as u8 == index);
		Some(reason)
	}

	fn to_str(self) -> &'static str {
		match self {
			Self::MissingKing => "missing king",
			Self::HangingKing => "hanging king",
			Self::PawnOnBackRank => "pawn on back rank",
			Self::InvalidCastlingRights => "invalid castling rights",
			Self::InvalidEnPassant => "invalid en passant",
			Self::TooMuchMaterial => "too much material",
			Self::ImpossibleEnPassantPin => "illegal en passant",
		}
	}
}

impl core::fmt::Display for IllegalPositionReason {
	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
		f.write_str(self.to_str())
	}
}

/// An error when trying to parse a position record.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct ParseRecordError {
	/// The index where the error was detected.
	pub byte: usize,
}
impl core::fmt::Display for ParseRecordError {
	fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
		write!(f, "invalid text record (at byte {})", self.byte)
	}
}
impl core::error::Error for ParseRecordError {}