package main

import (
	"fmt"
	"sort"
)

const aiDepth = 3

var pieceValues = map[int]int{
	Empty:  0, Pawn: 100, Knight: 320, Bishop: 330,
	Rook: 500, Queen: 900, King: 20000,
}

var pawnTable = [8][8]int{
	{0, 0, 0, 0, 0, 0, 0, 0},
	{50, 50, 50, 50, 50, 50, 50, 50},
	{10, 10, 20, 30, 30, 20, 10, 10},
	{5, 5, 10, 25, 25, 10, 5, 5},
	{0, 0, 0, 20, 20, 0, 0, 0},
	{5, -5, -10, 0, 0, -10, -5, 5},
	{5, 10, 10, -20, -20, 10, 10, 5},
	{0, 0, 0, 0, 0, 0, 0, 0},
}

var knightTable = [8][8]int{
	{-50, -40, -30, -30, -30, -30, -40, -50},
	{-40, -20, 0, 0, 0, 0, -20, -40},
	{-30, 0, 10, 15, 15, 10, 0, -30},
	{-30, 5, 15, 20, 20, 15, 5, -30},
	{-30, 0, 15, 20, 20, 15, 0, -30},
	{-30, 5, 10, 15, 15, 10, 5, -30},
	{-40, -20, 0, 5, 5, 0, -20, -40},
	{-50, -40, -30, -30, -30, -30, -40, -50},
}

var bishopTable = [8][8]int{
	{-20, -10, -10, -10, -10, -10, -10, -20},
	{-10, 0, 0, 0, 0, 0, 0, -10},
	{-10, 0, 5, 10, 10, 5, 0, -10},
	{-10, 5, 5, 10, 10, 5, 5, -10},
	{-10, 0, 10, 10, 10, 10, 0, -10},
	{-10, 10, 10, 10, 10, 10, 10, -10},
	{-10, 5, 0, 0, 0, 0, 5, -10},
	{-20, -10, -10, -10, -10, -10, -10, -20},
}

var rookTable = [8][8]int{
	{0, 0, 0, 0, 0, 0, 0, 0},
	{5, 10, 10, 10, 10, 10, 10, 5},
	{-5, 0, 0, 0, 0, 0, 0, -5},
	{-5, 0, 0, 0, 0, 0, 0, -5},
	{-5, 0, 0, 0, 0, 0, 0, -5},
	{-5, 0, 0, 0, 0, 0, 0, -5},
	{-5, 0, 0, 0, 0, 0, 0, -5},
	{0, 0, 0, 5, 5, 0, 0, 0},
}

var queenTable = [8][8]int{
	{-20, -10, -10, -5, -5, -10, -10, -20},
	{-10, 0, 0, 0, 0, 0, 0, -10},
	{-10, 0, 5, 5, 5, 5, 0, -10},
	{-5, 0, 5, 5, 5, 5, 0, -5},
	{0, 0, 5, 5, 5, 5, 0, -5},
	{-10, 5, 5, 5, 5, 5, 0, -10},
	{-10, 0, 5, 0, 0, 0, 0, -10},
	{-20, -10, -10, -5, -5, -10, -10, -20},
}

var kingTable = [8][8]int{
	{-30, -40, -40, -50, -50, -40, -40, -30},
	{-30, -40, -40, -50, -50, -40, -40, -30},
	{-30, -40, -40, -50, -50, -40, -40, -30},
	{-30, -40, -40, -50, -50, -40, -40, -30},
	{-20, -30, -30, -40, -40, -30, -30, -20},
	{-10, -20, -20, -20, -20, -20, -20, -10},
	{20, 20, 0, 0, 0, 0, 20, 20},
	{20, 30, 10, 0, 0, 10, 30, 20},
}

func evaluate(b *Board) int {
	score := 0
	for r := 0; r < 8; r++ {
		for c := 0; c < 8; c++ {
			s := b.Grid[r][c]
			if s.Piece == Empty {
				continue
			}
			val := pieceValues[s.Piece]
			var bonus int
			switch s.Piece {
			case Pawn:
				if s.Color == White {
					bonus = pawnTable[r][c]
				} else {
					bonus = pawnTable[7-r][c]
				}
			case Knight:
				if s.Color == White {
					bonus = knightTable[r][c]
				} else {
					bonus = knightTable[7-r][c]
				}
			case Bishop:
				if s.Color == White {
					bonus = bishopTable[r][c]
				} else {
					bonus = bishopTable[7-r][c]
				}
			case Rook:
				if s.Color == White {
					bonus = rookTable[r][c]
				} else {
					bonus = rookTable[7-r][c]
				}
			case Queen:
				if s.Color == White {
					bonus = queenTable[r][c]
				} else {
					bonus = queenTable[7-r][c]
				}
			case King:
				if s.Color == White {
					bonus = kingTable[r][c]
				} else {
					bonus = kingTable[7-r][c]
				}
			}
			if s.Color == White {
				score += val + bonus
			} else {
				score -= val + bonus
			}
		}
	}
	return score
}

func minimax(b *Board, depth int, alpha, beta int, color int) int {
	if depth == 0 {
		return evaluate(b)
	}
	moves := b.generateLegalMoves(color)
	if len(moves) == 0 {
		if b.isInCheck(color) {
			if color == White {
				return -100000 - depth
			}
			return 100000 + depth
		}
		return 0
	}
	oppColor := Black
	if color == White {
		oppColor = Black
	} else {
		oppColor = White
	}
	if color == White {
		maxEval := -999999
		for _, m := range moves {
			clone := b.Clone()
			clone.applyMove(m)
			eval := minimax(clone, depth-1, alpha, beta, oppColor)
			if eval > maxEval {
				maxEval = eval
			}
			if maxEval > alpha {
				alpha = maxEval
			}
			if alpha >= beta {
				break
			}
		}
		return maxEval
	} else {
		minEval := 999999
		for _, m := range moves {
			clone := b.Clone()
			clone.applyMove(m)
			eval := minimax(clone, depth-1, alpha, beta, oppColor)
			if eval < minEval {
				minEval = eval
			}
			if minEval < beta {
				beta = minEval
			}
			if alpha >= beta {
				break
			}
		}
		return minEval
	}
}

func GetBestMove(b *Board, color int) (string, error) {
	moves := b.generateLegalMoves(color)
	if len(moves) == 0 {
		return "", fmt.Errorf("no legal moves")
	}
	oppColor := Black
	if color == White {
		oppColor = Black
	} else {
		oppColor = White
	}
	var bestMove *Move
	var bestScore int
	first := true
	alpha := -999999
	beta := 999999
	sort.Slice(moves, func(i, j int) bool {
		vi := 0
		vj := 0
		if b.Grid[moves[i].To.Row][moves[i].To.Col].Piece != Empty {
			vi = pieceValues[b.Grid[moves[i].To.Row][moves[i].To.Col].Piece]
		}
		if b.Grid[moves[j].To.Row][moves[j].To.Col].Piece != Empty {
			vj = pieceValues[b.Grid[moves[j].To.Row][moves[j].To.Col].Piece]
		}
		return vi > vj
	})
	for _, m := range moves {
		clone := b.Clone()
		clone.applyMove(m)
		score := minimax(clone, aiDepth-1, alpha, beta, oppColor)
		if first {
			bestMove = &m
			bestScore = score
			first = false
		} else if color == White && score > bestScore {
			bestScore = score
			bestMove = &m
			alpha = score
		} else if color == Black && score < bestScore {
			bestScore = score
			bestMove = &m
			beta = score
		}
	}
	if bestMove == nil {
		return "", fmt.Errorf("no move found")
	}
	return MoveToAlgebraic(*bestMove), nil
}

func PlayVsComputer() {
	fmt.Println("=== CHESS VS COMPUTER ===")
	fmt.Println("You play WHITE. Enter moves like: e2e4, g1f3, e7e8q")
	fmt.Println("Type 'quit' to exit.")
	fmt.Println()
	g := NewGame()
	fmt.Print(g.Render())
	for !g.Over {
		fmt.Print("White> ")
		var input string
		fmt.Scanln(&input)
		if input == "quit" || input == "q" {
			fmt.Println("Goodbye!")
			return
		}
		_, err := g.MakeMove(input)
		if err != nil {
			fmt.Printf("  Error: %v\n\n", err)
			continue
		}
		fmt.Print(g.Render())
		if g.Over {
			break
		}
		fmt.Print("  Black thinking... ")
		aiMove, err := GetBestMove(g.Board, Black)
		if err != nil {
			fmt.Printf("AI error: %v\n", err)
			break
		}
		fmt.Printf("Black plays %s\n", aiMove)
		_, err = g.MakeMove(aiMove)
		if err != nil {
			fmt.Printf("  Error: %v\n", err)
			continue
		}
		fmt.Print(g.Render())
	}
	switch g.Result {
	case WhiteWins:
		fmt.Println("  *** YOU WIN! ***")
	case BlackWins:
		fmt.Println("  *** COMPUTER WINS! ***")
	case Draw:
		fmt.Println("  *** DRAW! ***")
	}
}