package main

import "fmt"

const (
	Empty  = 0
	Pawn   = 1
	Knight = 2
	Bishop = 3
	Rook   = 4
	Queen  = 5
	King   = 6
)

const (
	White = 1
	Black = 2
)

type Square struct {
	Piece int
	Color int
}

type Position struct {
	Row int
	Col int
}

type Board struct {
	Grid            [8][8]Square
	Castling        map[string]bool
	EnPassantTarget *Position
	MoveHistory     []string
}

func NewBoard() *Board {
	b := &Board{
		Castling: map[string]bool{"K": true, "Q": true, "k": true, "q": true},
	}
	backRank := []int{Rook, Knight, Bishop, Queen, King, Bishop, Knight, Rook}
	for col, p := range backRank {
		b.Grid[0][col] = Square{Piece: p, Color: Black}
		b.Grid[7][col] = Square{Piece: p, Color: White}
	}
	for col := 0; col < 8; col++ {
		b.Grid[1][col] = Square{Piece: Pawn, Color: Black}
		b.Grid[6][col] = Square{Piece: Pawn, Color: White}
	}
	return b
}

func (b *Board) Clone() *Board {
	clone := *b
	clone.Castling = make(map[string]bool)
	for k, v := range b.Castling {
		clone.Castling[k] = v
	}
	if b.EnPassantTarget != nil {
		pt := *b.EnPassantTarget
		clone.EnPassantTarget = &pt
	}
	return &clone
}

func (b *Board) inBounds(r, c int) bool {
	return r >= 0 && r < 8 && c >= 0 && c < 8
}

func (b *Board) isSquareAttacked(pos Position, byColor int) bool {
	r, c := pos.Row, pos.Col
	pawnDir := -1
	if byColor == Black {
		pawnDir = 1
	}
	for dc := -1; dc <= 1; dc += 2 {
		nr, nc := r-pawnDir, c+dc
		if b.inBounds(nr, nc) {
			if b.Grid[nr][nc].Piece == Pawn && b.Grid[nr][nc].Color == byColor {
				return true
			}
		}
	}
	knightMoves := [][2]int{{-2, -1}, {-2, 1}, {-1, -2}, {-1, 2}, {1, -2}, {1, 2}, {2, -1}, {2, 1}}
	for _, dm := range knightMoves {
		nr, nc := r+dm[0], c+dm[1]
		if b.inBounds(nr, nc) {
			if b.Grid[nr][nc].Piece == Knight && b.Grid[nr][nc].Color == byColor {
				return true
			}
		}
	}
	for dr := -1; dr <= 1; dr++ {
		for dc := -1; dc <= 1; dc++ {
			if dr == 0 && dc == 0 {
				continue
			}
			nr, nc := r+dr, c+dc
			if b.inBounds(nr, nc) {
				if b.Grid[nr][nc].Piece == King && b.Grid[nr][nc].Color == byColor {
					return true
				}
			}
		}
	}
	ortho := [][2]int{{-1, 0}, {1, 0}, {0, -1}, {0, 1}}
	for _, dir := range ortho {
		nr, nc := r+dir[0], c+dir[1]
		for b.inBounds(nr, nc) {
			s := b.Grid[nr][nc]
			if s.Piece != Empty {
				if s.Color == byColor && (s.Piece == Rook || s.Piece == Queen) {
					return true
				}
				break
			}
			nr += dir[0]
			nc += dir[1]
		}
	}
	diag := [][2]int{{-1, -1}, {-1, 1}, {1, -1}, {1, 1}}
	for _, dir := range diag {
		nr, nc := r+dir[0], c+dir[1]
		for b.inBounds(nr, nc) {
			s := b.Grid[nr][nc]
			if s.Piece != Empty {
				if s.Color == byColor && (s.Piece == Bishop || s.Piece == Queen) {
					return true
				}
				break
			}
			nr += dir[0]
			nc += dir[1]
		}
	}
	return false
}

func (b *Board) findKing(color int) Position {
	for r := 0; r < 8; r++ {
		for c := 0; c < 8; c++ {
			if b.Grid[r][c].Piece == King && b.Grid[r][c].Color == color {
				return Position{r, c}
			}
		}
	}
	return Position{}
}

func (b *Board) isInCheck(color int) bool {
	kingPos := b.findKing(color)
	oppColor := Black
	if color == Black {
		oppColor = White
	}
	return b.isSquareAttacked(kingPos, oppColor)
}

type Move struct {
	From      Position
	To        Position
	Promotion int
}

func (b *Board) generatePseudoLegalMoves(color int) []Move {
	var moves []Move
	for r := 0; r < 8; r++ {
		for c := 0; c < 8; c++ {
			s := b.Grid[r][c]
			if s.Piece == Empty || s.Color != color {
				continue
			}
			from := Position{r, c}
			switch s.Piece {
			case Pawn:
				moves = append(moves, b.pawnMoves(from, color)...)
			case Knight:
				moves = append(moves, b.knightMoves(from, color)...)
			case Bishop:
				moves = append(moves, b.slideMoves(from, color, diagDirs)...)
			case Rook:
				moves = append(moves, b.slideMoves(from, color, orthoDirs)...)
			case Queen:
				all := append(append([][2]int{}, diagDirs...), orthoDirs...)
				moves = append(moves, b.slideMoves(from, color, all)...)
			case King:
				moves = append(moves, b.kingMoves(from, color)...)
			}
		}
	}
	return moves
}

var orthoDirs = [][2]int{{-1, 0}, {1, 0}, {0, -1}, {0, 1}}
var diagDirs = [][2]int{{-1, -1}, {-1, 1}, {1, -1}, {1, 1}}

func (b *Board) pawnMoves(from Position, color int) []Move {
	var moves []Move
	r, c := from.Row, from.Col
	dir := -1
	if color == Black {
		dir = 1
	}
	nr := r + dir
	if b.inBounds(nr, c) && b.Grid[nr][c].Piece == Empty {
		if (color == White && nr == 0) || (color == Black && nr == 7) {
			for _, p := range []int{Queen, Rook, Bishop, Knight} {
				moves = append(moves, Move{From: from, To: Position{nr, c}, Promotion: p})
			}
		} else {
			moves = append(moves, Move{From: from, To: Position{nr, c}})
		}
		startRow := 6
		if color == Black {
			startRow = 1
		}
		if r == startRow {
			nr2 := r + 2*dir
			if b.inBounds(nr2, c) && b.Grid[nr2][c].Piece == Empty {
				moves = append(moves, Move{From: from, To: Position{nr2, c}})
			}
		}
	}
	for dc := -1; dc <= 1; dc += 2 {
		nc := c + dc
		if !b.inBounds(nr, nc) {
			continue
		}
		target := b.Grid[nr][nc]
		if target.Piece != Empty && target.Color != color {
			if (color == White && nr == 0) || (color == Black && nr == 7) {
				for _, p := range []int{Queen, Rook, Bishop, Knight} {
					moves = append(moves, Move{From: from, To: Position{nr, nc}, Promotion: p})
				}
			} else {
				moves = append(moves, Move{From: from, To: Position{nr, nc}})
			}
		}
		if b.EnPassantTarget != nil && b.EnPassantTarget.Row == nr && b.EnPassantTarget.Col == nc {
			moves = append(moves, Move{From: from, To: Position{nr, nc}})
		}
	}
	return moves
}

func (b *Board) knightMoves(from Position, color int) []Move {
	var moves []Move
	kd := [][2]int{{-2, -1}, {-2, 1}, {-1, -2}, {-1, 2}, {1, -2}, {1, 2}, {2, -1}, {2, 1}}
	for _, d := range kd {
		nr, nc := from.Row+d[0], from.Col+d[1]
		if !b.inBounds(nr, nc) {
			continue
		}
		if b.Grid[nr][nc].Color != color {
			moves = append(moves, Move{From: from, To: Position{nr, nc}})
		}
	}
	return moves
}

func (b *Board) slideMoves(from Position, color int, dirs [][2]int) []Move {
	var moves []Move
	for _, d := range dirs {
		nr, nc := from.Row+d[0], from.Col+d[1]
		for b.inBounds(nr, nc) {
			if b.Grid[nr][nc].Piece == Empty {
				moves = append(moves, Move{From: from, To: Position{nr, nc}})
			} else {
				if b.Grid[nr][nc].Color != color {
					moves = append(moves, Move{From: from, To: Position{nr, nc}})
				}
				break
			}
			nr += d[0]
			nc += d[1]
		}
	}
	return moves
}

func (b *Board) kingMoves(from Position, color int) []Move {
	var moves []Move
	for dr := -1; dr <= 1; dr++ {
		for dc := -1; dc <= 1; dc++ {
			if dr == 0 && dc == 0 {
				continue
			}
			nr, nc := from.Row+dr, from.Col+dc
			if !b.inBounds(nr, nc) {
				continue
			}
			if b.Grid[nr][nc].Color != color {
				moves = append(moves, Move{From: from, To: Position{nr, nc}})
			}
		}
	}
	oppColor := Black
	if color == White {
		oppColor = Black
	} else {
		oppColor = White
	}
	if color == White && from.Row == 7 && from.Col == 4 {
		if b.Castling["K"] && b.Grid[7][5].Piece == Empty && b.Grid[7][6].Piece == Empty {
			if !b.isSquareAttacked(Position{7, 4}, oppColor) &&
				!b.isSquareAttacked(Position{7, 5}, oppColor) &&
				!b.isSquareAttacked(Position{7, 6}, oppColor) {
				moves = append(moves, Move{From: from, To: Position{7, 6}})
			}
		}
		if b.Castling["Q"] && b.Grid[7][3].Piece == Empty && b.Grid[7][2].Piece == Empty && b.Grid[7][1].Piece == Empty {
			if !b.isSquareAttacked(Position{7, 4}, oppColor) &&
				!b.isSquareAttacked(Position{7, 3}, oppColor) &&
				!b.isSquareAttacked(Position{7, 2}, oppColor) {
				moves = append(moves, Move{From: from, To: Position{7, 2}})
			}
		}
	} else if color == Black && from.Row == 0 && from.Col == 4 {
		if b.Castling["k"] && b.Grid[0][5].Piece == Empty && b.Grid[0][6].Piece == Empty {
			if !b.isSquareAttacked(Position{0, 4}, oppColor) &&
				!b.isSquareAttacked(Position{0, 5}, oppColor) &&
				!b.isSquareAttacked(Position{0, 6}, oppColor) {
				moves = append(moves, Move{From: from, To: Position{0, 6}})
			}
		}
		if b.Castling["q"] && b.Grid[0][3].Piece == Empty && b.Grid[0][2].Piece == Empty && b.Grid[0][1].Piece == Empty {
			if !b.isSquareAttacked(Position{0, 4}, oppColor) &&
				!b.isSquareAttacked(Position{0, 3}, oppColor) &&
				!b.isSquareAttacked(Position{0, 2}, oppColor) {
				moves = append(moves, Move{From: from, To: Position{0, 2}})
			}
		}
	}
	return moves
}

func (b *Board) generateLegalMoves(color int) []Move {
	pseudo := b.generatePseudoLegalMoves(color)
	var legal []Move
	for _, m := range pseudo {
		clone := b.Clone()
		clone.applyMove(m)
		if !clone.isInCheck(color) {
			legal = append(legal, m)
		}
	}
	return legal
}

func (b *Board) applyMove(m Move) {
	from, to := m.From, m.To
	movingPiece := b.Grid[from.Row][from.Col]
	if movingPiece.Piece == Pawn && to.Col != from.Col && b.Grid[to.Row][to.Col].Piece == Empty {
		b.Grid[from.Row][to.Col] = Square{}
	}
	b.Grid[to.Row][to.Col] = movingPiece
	b.Grid[from.Row][from.Col] = Square{}
	if movingPiece.Piece == Pawn && m.Promotion != 0 {
		b.Grid[to.Row][to.Col].Piece = m.Promotion
	}
	if movingPiece.Piece == King && MathAbs(to.Col-from.Col) == 2 {
		row := from.Row
		if to.Col == 6 {
			b.Grid[row][5] = b.Grid[row][7]
			b.Grid[row][7] = Square{}
		} else if to.Col == 2 {
			b.Grid[row][3] = b.Grid[row][0]
			b.Grid[row][0] = Square{}
		}
	}
	if movingPiece.Piece == King {
		if movingPiece.Color == White {
			b.Castling["K"] = false
			b.Castling["Q"] = false
		} else {
			b.Castling["k"] = false
			b.Castling["q"] = false
		}
	}
	if movingPiece.Piece == Rook {
		if from.Row == 7 && from.Col == 0 {
			b.Castling["Q"] = false
		}
		if from.Row == 7 && from.Col == 7 {
			b.Castling["K"] = false
		}
		if from.Row == 0 && from.Col == 0 {
			b.Castling["q"] = false
		}
		if from.Row == 0 && from.Col == 7 {
			b.Castling["k"] = false
		}
	}
	if to.Row == 7 && to.Col == 0 {
		b.Castling["Q"] = false
	}
	if to.Row == 7 && to.Col == 7 {
		b.Castling["K"] = false
	}
	if to.Row == 0 && to.Col == 0 {
		b.Castling["q"] = false
	}
	if to.Row == 0 && to.Col == 7 {
		b.Castling["k"] = false
	}
	b.EnPassantTarget = nil
	if movingPiece.Piece == Pawn {
		if MathAbs(to.Row-from.Row) == 2 {
			b.EnPassantTarget = &Position{(from.Row + to.Row) / 2, from.Col}
		}
	}
}

func MathAbs(x int) int {
	if x < 0 {
		return -x
	}
	return x
}

func (b *Board) Display(checkColor int) string {
	pieceChars := map[int]string{
		Pawn: "P", Knight: "N", Bishop: "B", Rook: "R", Queen: "Q", King: "K",
	}
	var out string
	out += "  a  b  c  d  e  f  g  h\n"
	for r := 0; r < 8; r++ {
		out += fmt.Sprintf("%d", 8-r)
		for c := 0; c < 8; c++ {
			s := b.Grid[r][c]
			if s.Piece == Empty {
				out += "  ."
			} else {
				ch := pieceChars[s.Piece]
				if s.Color == White {
					out += "  " + ch
				} else {
					out += "  " + string(rune(97+int(ch[0])-65))
				}
			}
		}
		out += fmt.Sprintf("  %d\n", 8-r)
	}
	out += "  a  b  c  d  e  f  g  h\n"
	if b.isInCheck(checkColor) {
		out += fmt.Sprintf("\n  !! %s IS IN CHECK !!\n", colorName(checkColor))
	}
	return out
}

func colorName(c int) string {
	if c == White {
		return "WHITE"
	}
	return "BLACK"
}

func (b *Board) ParseMove(input string) (*Move, error) {
	if len(input) < 4 || len(input) > 5 {
		return nil, fmt.Errorf("invalid move format: %s (use e.g. e2e4)", input)
	}
	fromCol := int(input[0] - 'a')
	fromRow := 8 - int(input[1] - '0')
	toCol := int(input[2] - 'a')
	toRow := 8 - int(input[3] - '0')
	if fromCol < 0 || fromCol > 7 || toCol < 0 || toCol > 7 {
		return nil, fmt.Errorf("invalid columns")
	}
	if fromRow < 0 || fromRow > 7 || toRow < 0 || toRow > 7 {
		return nil, fmt.Errorf("invalid rows")
	}
	promotion := 0
	if len(input) == 5 {
		switch input[4] {
		case 'q', 'Q':
			promotion = Queen
		case 'r', 'R':
			promotion = Rook
		case 'b', 'B':
			promotion = Bishop
		case 'n', 'N':
			promotion = Knight
		default:
			return nil, fmt.Errorf("invalid promotion piece: %c", input[4])
		}
	}
	m := Move{
		From:      Position{fromRow, fromCol},
		To:        Position{toRow, toCol},
		Promotion: promotion,
	}
	return &m, nil
}

func MoveToAlgebraic(m Move) string {
	from := string(rune('a'+m.From.Col)) + string(rune('0'+(8-m.From.Row)))
	to := string(rune('a'+m.To.Col)) + string(rune('0'+(8-m.To.Row)))
	s := from + to
	if m.Promotion != 0 {
		pieceChars := map[int]byte{Queen: 'q', Rook: 'r', Bishop: 'b', Knight: 'n'}
		s += string(pieceChars[m.Promotion])
	}
	return s
}