package main

import "errors"

// ПРИМЕЧАНИЕ О ПРАВИЛАХ: реализованы полные правила шахмат, включая
// рокировку (в обе стороны, со всеми условиями: ни король, ни ладья
// не должны были ходить, поля между ними свободны, король не под
// шахом и не проходит через атакованное поле) и взятие на проходе.
// Превращение пешки УПРОЩЕНО: пешка всегда превращается в ферзя
// (автоматически, без выбора фигуры) — это стандартное упрощение
// для большинства цифровых реализаций, так как подавляющее
// большинство реальных превращений и так выбирают ферзя. Из
// официальных правил ничьей реализованы мат/пат, недостаточность
// материала и упрощённый аналог правила 75 полуходов без взятия и
// хода пешкой (автоматическое завершение партии, без необходимости
// заявлять ничью); троекратное повторение позиции НЕ реализовано
// (потребовало бы хранить историю всех позиций).

// ChessColor — цвет стороны.
type ChessColor int

const (
	ChessWhite ChessColor = iota
	ChessBlack
)

func (c ChessColor) Opponent() ChessColor {
	if c == ChessWhite {
		return ChessBlack
	}
	return ChessWhite
}

// ChessPieceType — тип фигуры.
type ChessPieceType int

const (
	ChessPawn ChessPieceType = iota
	ChessKnight
	ChessBishop
	ChessRook
	ChessQueen
	ChessKing
)

// ChessPiece — фигура на доске.
type ChessPiece struct {
	Color ChessColor
	Type  ChessPieceType
}

// ChessPos — координаты клетки (0-7, 0-7). Row 0 = 8-я горизонталь
// (чёрные сзади), Row 7 = 1-я горизонталь (белые сзади). Col 0-7 =
// вертикали a-h.
type ChessPos struct {
	Row, Col int
}

func (p ChessPos) inBounds() bool {
	return p.Row >= 0 && p.Row < 8 && p.Col >= 0 && p.Col < 8
}

var (
	ErrChessNoPiece     = errors.New("на этой клетке нет вашей фигуры")
	ErrChessInvalidMove = errors.New("такой ход недопустим")
	ErrChessGameOver    = errors.New("партия уже завершена")
)

// chessMoveOption — один вариант хода для конкретной фигуры.
type chessMoveOption struct {
	To                ChessPos
	IsEnPassant       bool
	IsCastleKingside  bool
	IsCastleQueenside bool
	IsPromotion       bool
}

// ChessResult — итог завершившейся партии.
type ChessResult struct {
	Winner     ChessColor
	Checkmate  bool
	Draw       bool
	DrawReason string // "stalemate", "insufficient_material", "no_progress"
}

// maxChessHalfmovesWithoutProgress — защита от бесконечных партий
// без взятий и ходов пешками (упрощённый аналог официального
// правила 75 полуходов).
const maxChessHalfmovesWithoutProgress = 150

// ChessGameState — полное состояние партии в шахматы.
type ChessGameState struct {
	Board [8][8]*ChessPiece
	Turn  ChessColor

	WhiteKingMoved          bool
	WhiteRookQueensideMoved bool // ладья a1
	WhiteRookKingsideMoved  bool // ладья h1
	BlackKingMoved          bool
	BlackRookQueensideMoved bool // ладья a8
	BlackRookKingsideMoved  bool // ладья h8

	EnPassantTarget *ChessPos // клетка "за" пешкой, только что шагнувшей на 2 — цель взятия на проходе
	HalfmoveClock   int       // полуходов без взятия/хода пешкой

	Result *ChessResult
}

// NewChessGame создаёт новую партию со стандартной расстановкой;
// первый ход — за белыми.
func NewChessGame() *ChessGameState {
	g := &ChessGameState{Turn: ChessWhite}
	backRank := []ChessPieceType{ChessRook, ChessKnight, ChessBishop, ChessQueen, ChessKing, ChessBishop, ChessKnight, ChessRook}
	for col, t := range backRank {
		g.Board[0][col] = &ChessPiece{Color: ChessBlack, Type: t}
		g.Board[7][col] = &ChessPiece{Color: ChessWhite, Type: t}
	}
	for col := 0; col < 8; col++ {
		g.Board[1][col] = &ChessPiece{Color: ChessBlack, Type: ChessPawn}
		g.Board[6][col] = &ChessPiece{Color: ChessWhite, Type: ChessPawn}
	}
	return g
}

func (p ChessPos) add(dr, dc int) ChessPos { return ChessPos{Row: p.Row + dr, Col: p.Col + dc} }

var chessBishopDirs = []ChessPos{{Row: -1, Col: -1}, {Row: -1, Col: 1}, {Row: 1, Col: -1}, {Row: 1, Col: 1}}
var chessRookDirs = []ChessPos{{Row: -1, Col: 0}, {Row: 1, Col: 0}, {Row: 0, Col: -1}, {Row: 0, Col: 1}}
var chessKnightOffsets = []ChessPos{{-2, -1}, {-2, 1}, {-1, -2}, {-1, 2}, {1, -2}, {1, 2}, {2, -1}, {2, 1}}

// slidingMoves перечисляет ходы скользящей фигуры (ладья/слон/ферзь)
// вдоль набора направлений, останавливаясь на первой фигуре (со
// взятием, если это вражеская).
func (g *ChessGameState) slidingMoves(pos ChessPos, dirs []ChessPos, color ChessColor) []chessMoveOption {
	var out []chessMoveOption
	for _, d := range dirs {
		for step := 1; ; step++ {
			np := pos.add(d.Row*step, d.Col*step)
			if !np.inBounds() {
				break
			}
			occupant := g.Board[np.Row][np.Col]
			if occupant == nil {
				out = append(out, chessMoveOption{To: np})
				continue
			}
			if occupant.Color != color {
				out = append(out, chessMoveOption{To: np})
			}
			break
		}
	}
	return out
}

// pseudoLegalMoves перечисляет ходы фигуры без учёта того, остаётся
// ли после них свой король под шахом (эта проверка — уровнем выше,
// в LegalMovesFrom).
func (g *ChessGameState) pseudoLegalMoves(pos ChessPos) []chessMoveOption {
	piece := g.Board[pos.Row][pos.Col]
	if piece == nil {
		return nil
	}
	switch piece.Type {
	case ChessKnight:
		var out []chessMoveOption
		for _, o := range chessKnightOffsets {
			np := pos.add(o.Row, o.Col)
			if !np.inBounds() {
				continue
			}
			occupant := g.Board[np.Row][np.Col]
			if occupant == nil || occupant.Color != piece.Color {
				out = append(out, chessMoveOption{To: np})
			}
		}
		return out
	case ChessBishop:
		return g.slidingMoves(pos, chessBishopDirs, piece.Color)
	case ChessRook:
		return g.slidingMoves(pos, chessRookDirs, piece.Color)
	case ChessQueen:
		out := g.slidingMoves(pos, chessBishopDirs, piece.Color)
		return append(out, g.slidingMoves(pos, chessRookDirs, piece.Color)...)
	case ChessKing:
		return g.kingMoves(pos, piece.Color)
	case ChessPawn:
		return g.pawnMoves(pos, piece.Color)
	}
	return nil
}

func (g *ChessGameState) kingMoves(pos ChessPos, color ChessColor) []chessMoveOption {
	var out []chessMoveOption
	for _, d := range append(append([]ChessPos{}, chessBishopDirs...), chessRookDirs...) {
		np := pos.add(d.Row, d.Col)
		if !np.inBounds() {
			continue
		}
		occupant := g.Board[np.Row][np.Col]
		if occupant == nil || occupant.Color != color {
			out = append(out, chessMoveOption{To: np})
		}
	}
	out = append(out, g.castlingMoves(pos, color)...)
	return out
}

func (g *ChessGameState) castlingMoves(kingPos ChessPos, color ChessColor) []chessMoveOption {
	if g.inCheck(color) {
		return nil // нельзя рокироваться, находясь под шахом
	}
	row := kingPos.Row
	var kingMoved, rookQMoved, rookKMoved bool
	if color == ChessWhite {
		kingMoved, rookQMoved, rookKMoved = g.WhiteKingMoved, g.WhiteRookQueensideMoved, g.WhiteRookKingsideMoved
	} else {
		kingMoved, rookQMoved, rookKMoved = g.BlackKingMoved, g.BlackRookQueensideMoved, g.BlackRookKingsideMoved
	}
	if kingMoved {
		return nil
	}
	var out []chessMoveOption
	enemy := color.Opponent()

	// короткая рокировка (kingside): ладья h, король e->g
	if !rookKMoved && g.Board[row][5] == nil && g.Board[row][6] == nil {
		if rook := g.Board[row][7]; rook != nil && rook.Type == ChessRook && rook.Color == color {
			if !g.isSquareAttacked(ChessPos{row, 5}, enemy) && !g.isSquareAttacked(ChessPos{row, 6}, enemy) {
				out = append(out, chessMoveOption{To: ChessPos{row, 6}, IsCastleKingside: true})
			}
		}
	}
	// длинная рокировка (queenside): ладья a, король e->c
	if !rookQMoved && g.Board[row][1] == nil && g.Board[row][2] == nil && g.Board[row][3] == nil {
		if rook := g.Board[row][0]; rook != nil && rook.Type == ChessRook && rook.Color == color {
			if !g.isSquareAttacked(ChessPos{row, 3}, enemy) && !g.isSquareAttacked(ChessPos{row, 2}, enemy) {
				out = append(out, chessMoveOption{To: ChessPos{row, 2}, IsCastleQueenside: true})
			}
		}
	}
	return out
}

func (g *ChessGameState) pawnMoves(pos ChessPos, color ChessColor) []chessMoveOption {
	var out []chessMoveOption
	forward := 1
	startRow := 1
	lastRow := 7
	if color == ChessWhite {
		forward = -1
		startRow = 6
		lastRow = 0
	}

	one := pos.add(forward, 0)
	if one.inBounds() && g.Board[one.Row][one.Col] == nil {
		out = append(out, chessMoveOption{To: one, IsPromotion: one.Row == lastRow})
		if pos.Row == startRow {
			two := pos.add(forward*2, 0)
			if g.Board[two.Row][two.Col] == nil {
				out = append(out, chessMoveOption{To: two})
			}
		}
	}
	for _, dc := range []int{-1, 1} {
		capture := pos.add(forward, dc)
		if !capture.inBounds() {
			continue
		}
		if occupant := g.Board[capture.Row][capture.Col]; occupant != nil && occupant.Color != color {
			out = append(out, chessMoveOption{To: capture, IsPromotion: capture.Row == lastRow})
			continue
		}
		if g.EnPassantTarget != nil && *g.EnPassantTarget == capture {
			out = append(out, chessMoveOption{To: capture, IsEnPassant: true})
		}
	}
	return out
}

// isSquareAttacked проверяет, атакована ли клетка pos какой-либо
// фигурой цвета byColor.
func (g *ChessGameState) isSquareAttacked(pos ChessPos, byColor ChessColor) bool {
	dir := 1
	if byColor == ChessWhite {
		dir = -1
	}
	for _, dc := range []int{-1, 1} {
		p := pos.add(-dir, dc)
		if p.inBounds() {
			if occ := g.Board[p.Row][p.Col]; occ != nil && occ.Color == byColor && occ.Type == ChessPawn {
				return true
			}
		}
	}
	for _, o := range chessKnightOffsets {
		p := pos.add(o.Row, o.Col)
		if p.inBounds() {
			if occ := g.Board[p.Row][p.Col]; occ != nil && occ.Color == byColor && occ.Type == ChessKnight {
				return true
			}
		}
	}
	for _, d := range append(append([]ChessPos{}, chessBishopDirs...), chessRookDirs...) {
		p := pos.add(d.Row, d.Col)
		if p.inBounds() {
			if occ := g.Board[p.Row][p.Col]; occ != nil && occ.Color == byColor && occ.Type == ChessKing {
				return true
			}
		}
	}
	for _, d := range chessBishopDirs {
		for step := 1; ; step++ {
			p := pos.add(d.Row*step, d.Col*step)
			if !p.inBounds() {
				break
			}
			occ := g.Board[p.Row][p.Col]
			if occ == nil {
				continue
			}
			if occ.Color == byColor && (occ.Type == ChessBishop || occ.Type == ChessQueen) {
				return true
			}
			break
		}
	}
	for _, d := range chessRookDirs {
		for step := 1; ; step++ {
			p := pos.add(d.Row*step, d.Col*step)
			if !p.inBounds() {
				break
			}
			occ := g.Board[p.Row][p.Col]
			if occ == nil {
				continue
			}
			if occ.Color == byColor && (occ.Type == ChessRook || occ.Type == ChessQueen) {
				return true
			}
			break
		}
	}
	return false
}

func (g *ChessGameState) findKing(color ChessColor) ChessPos {
	for r := 0; r < 8; r++ {
		for c := 0; c < 8; c++ {
			if p := g.Board[r][c]; p != nil && p.Color == color && p.Type == ChessKing {
				return ChessPos{r, c}
			}
		}
	}
	return ChessPos{-1, -1} // не должно происходить при корректном состоянии
}

func (g *ChessGameState) inCheck(color ChessColor) bool {
	return g.isSquareAttacked(g.findKing(color), color.Opponent())
}

// cloneChess создаёт независимую глубокую копию — нужна для
// проверки "не остаётся ли король под шахом после хода" и для
// минимакса бота.
func cloneChess(g *ChessGameState) *ChessGameState {
	clone := *g
	for r := 0; r < 8; r++ {
		for c := 0; c < 8; c++ {
			if p := g.Board[r][c]; p != nil {
				cp := *p
				clone.Board[r][c] = &cp
			} else {
				clone.Board[r][c] = nil
			}
		}
	}
	if g.EnPassantTarget != nil {
		ep := *g.EnPassantTarget
		clone.EnPassantTarget = &ep
	}
	if g.Result != nil {
		res := *g.Result
		clone.Result = &res
	}
	return &clone
}

// LegalMovesFrom возвращает допустимые ходы фигуры в pos —
// псевдолегальные ходы, отфильтрованные так, чтобы после них
// собственный король не оставался под шахом.
func (g *ChessGameState) LegalMovesFrom(pos ChessPos) []chessMoveOption {
	if g.Result != nil {
		return nil
	}
	piece := g.Board[pos.Row][pos.Col]
	if piece == nil || piece.Color != g.Turn {
		return nil
	}
	var out []chessMoveOption
	for _, opt := range g.pseudoLegalMoves(pos) {
		clone := cloneChess(g)
		clone.applyMove(pos, opt)
		if !clone.inCheck(piece.Color) {
			out = append(out, opt)
		}
	}
	return out
}

// applyMove выполняет механику хода (перемещение фигур, взятие на
// проходе, рокировка, превращение) БЕЗ проверки легальности —
// вызывающий код должен был её уже провести.
func (g *ChessGameState) applyMove(from ChessPos, opt chessMoveOption) {
	piece := g.Board[from.Row][from.Col]

	if opt.IsEnPassant {
		g.Board[from.Row][opt.To.Col] = nil // взятая пешка стоит на строке from, столбце назначения
	}
	if opt.IsCastleKingside {
		rook := g.Board[from.Row][7]
		g.Board[from.Row][7] = nil
		g.Board[from.Row][5] = rook
	}
	if opt.IsCastleQueenside {
		rook := g.Board[from.Row][0]
		g.Board[from.Row][0] = nil
		g.Board[from.Row][3] = rook
	}

	g.Board[from.Row][from.Col] = nil
	if opt.IsPromotion {
		g.Board[opt.To.Row][opt.To.Col] = &ChessPiece{Color: piece.Color, Type: ChessQueen}
	} else {
		g.Board[opt.To.Row][opt.To.Col] = piece
	}

	if piece.Type == ChessKing {
		if piece.Color == ChessWhite {
			g.WhiteKingMoved = true
		} else {
			g.BlackKingMoved = true
		}
	}
	markRookMoved := func(pos ChessPos) {
		if pos.Row == 7 && pos.Col == 0 {
			g.WhiteRookQueensideMoved = true
		}
		if pos.Row == 7 && pos.Col == 7 {
			g.WhiteRookKingsideMoved = true
		}
		if pos.Row == 0 && pos.Col == 0 {
			g.BlackRookQueensideMoved = true
		}
		if pos.Row == 0 && pos.Col == 7 {
			g.BlackRookKingsideMoved = true
		}
	}
	markRookMoved(from)
	markRookMoved(opt.To) // на случай взятия ладьи соперника

	if piece.Type == ChessPawn && abs(opt.To.Row-from.Row) == 2 {
		mid := ChessPos{Row: (from.Row + opt.To.Row) / 2, Col: from.Col}
		g.EnPassantTarget = &mid
	} else {
		g.EnPassantTarget = nil
	}
}

func abs(n int) int {
	if n < 0 {
		return -n
	}
	return n
}

// hasAnyLegalMove проверяет, есть ли у стороны color хоть один
// легальный ход где-либо на доске.
func (g *ChessGameState) hasAnyLegalMove(color ChessColor) bool {
	for r := 0; r < 8; r++ {
		for c := 0; c < 8; c++ {
			p := g.Board[r][c]
			if p == nil || p.Color != color {
				continue
			}
			if len(g.LegalMovesFrom(ChessPos{r, c})) > 0 {
				return true
			}
		}
	}
	return false
}

// hasInsufficientMaterial проверяет самые распространённые
// теоретически недостаточные для мата комбинации: King vs King,
// King+Knight vs King, King+Bishop vs King. Более редкие ничейные
// комбинации (например, слоны на клетках одного цвета с обеих
// сторон) не проверяются — упрощение.
func (g *ChessGameState) hasInsufficientMaterial() bool {
	var minor []ChessPieceType
	for r := 0; r < 8; r++ {
		for c := 0; c < 8; c++ {
			p := g.Board[r][c]
			if p == nil || p.Type == ChessKing {
				continue
			}
			if p.Type != ChessKnight && p.Type != ChessBishop {
				return false
			}
			minor = append(minor, p.Type)
		}
	}
	return len(minor) <= 1
}

// Move выполняет ход фигурой из from в to, если он допустим по
// текущим правилам (см. LegalMovesFrom). Превращение пешки всегда
// автоматически происходит в ферзя.
func (g *ChessGameState) Move(from, to ChessPos) error {
	if g.Result != nil {
		return ErrChessGameOver
	}
	piece := g.Board[from.Row][from.Col]
	if piece == nil || piece.Color != g.Turn {
		return ErrChessNoPiece
	}
	options := g.LegalMovesFrom(from)
	var chosen *chessMoveOption
	for i := range options {
		if options[i].To == to {
			chosen = &options[i]
			break
		}
	}
	if chosen == nil {
		return ErrChessInvalidMove
	}

	isCaptureOrPawn := piece.Type == ChessPawn || g.Board[to.Row][to.Col] != nil || chosen.IsEnPassant
	g.applyMove(from, *chosen)

	if isCaptureOrPawn {
		g.HalfmoveClock = 0
	} else {
		g.HalfmoveClock++
	}

	g.Turn = g.Turn.Opponent()
	g.checkGameOver()
	return nil
}

func (g *ChessGameState) checkGameOver() {
	if g.hasInsufficientMaterial() {
		g.Result = &ChessResult{Draw: true, DrawReason: "insufficient_material"}
		return
	}
	if g.HalfmoveClock >= maxChessHalfmovesWithoutProgress {
		g.Result = &ChessResult{Draw: true, DrawReason: "no_progress"}
		return
	}
	if !g.hasAnyLegalMove(g.Turn) {
		if g.inCheck(g.Turn) {
			g.Result = &ChessResult{Winner: g.Turn.Opponent(), Checkmate: true}
		} else {
			g.Result = &ChessResult{Draw: true, DrawReason: "stalemate"}
		}
	}
}