#include <u.h>
#include <libc.h>
#include <draw.h>
#include <thread.h>
#include <mouse.h>
#include <keyboard.h>

// Missing features
// display stalemate
// arrows to go back and forward in history
// save & load games
// save & load position
// free edit board
// detect repetition
// engine play
// online play (FICS? lichess? custom?)

// the longest possible game with the 75-move rule is 8848.5 moves,
// or 17697 half-moves.
// It's totally overkill but I think we can afford the meg or two of RAM.
#define MAX_MOVES 40000

typedef enum Piece {
	empty,
	bB,
	bK,
	bN,
	bP,
	bQ,
	bR,
	wB,
	wK,
	wN,
	wP,
	wQ,
	wR
} Piece;

char *PIECE_STR[] = {
	"",
	"bB",
	"bK",
	"bN",
	"bP",
	"bQ",
	"bR",
	"wB",
	"wK",
	"wN",
	"wP",
	"wQ",
	"wR",
};

// I would use the raw letters but D is taken by draw.h :(
typedef enum ChessFile {
	fA,
	fB,
	fC,
	fD,
	fE,
	fF,
	fG,
	fH,
} ChessFile;

#define RANK(rank) (8-rank)
#define FILE(file) (file)
#define BRD(file, rank) (board[8-rank][file])
#define BRDPT(pt) (board[pt.y][pt.x])

// originally copied from /sys/src/cmd/page.c
void
resizewin(void)
{
	int wctl;

	if((wctl = open("/dev/wctl", OWRITE)) < 0)
		return;
	// 10px inner border, 8 50px squares, then rio border
	Point size = Pt(20+50*8 + Borderwidth*2 + 200, 20+50*8 + Borderwidth*2);
	fprint(wctl, "resize -dx %d -dy %d\n", size.x, size.y);
	close(wctl);
}

void resetboard(int board[8][8]) {
	// empty it
	for (int i=0; i<8*8; i++) ((int*)board)[i] = empty;
	// set pawns
	for (int x=0; x<8; x++) {
		BRD(x, 2) = wP;
		BRD(x, 7) = bP;
	}

	// set black pieces
	BRD(fA, 8) = BRD(fH, 8) = bR;
	BRD(fB, 8) = BRD(fG, 8) = bN;
	BRD(fC, 8) = BRD(fF, 8) = bB;
	BRD(fD, 8) = bQ;
	BRD(fE, 8) = bK;
	// set white pieces
	BRD(fA, 1) = BRD(fH, 1) = wR;
	BRD(fB, 1) = BRD(fG, 1) = wN;
	BRD(fC, 1) = BRD(fF, 1) = wB;
	BRD(fD, 1) = wQ;
	BRD(fE, 1) = wK;
}

typedef enum Player {
	White = 0,
	Black = 1
} Player;

typedef enum MoveType {
	Invalid,
	Normal,
	EnPassant,
	BlackShortCastle,
	BlackLongCastle,
	WhiteShortCastle,
	WhiteLongCastle,
	PromotionQ,
	PromotionN,
	PromotionB,
	PromotionR,
} MoveType;

typedef struct Move {
	Piece piece;
	Point src;
	Point dst;
} Move;

Player
owner(Piece piece)
{
	if (piece <= bR)
		return Black;
	return White;
}

typedef struct CastleRights {
	int wK_ever_moved;
	int bK_ever_moved;
	int wRA_ever_moved;
	int bRA_ever_moved;
	int wRH_ever_moved;
	int bRH_ever_moved;
} CastleRights;

char dbgstring[50];


MoveType
ismove(int board[8][8], Point src, Point dst, Move prev_move)
{
	Piece actor = BRDPT(src);
	Piece target = BRDPT(dst);
	if (target && owner(actor) == owner(target)) {
		return Invalid; // can't take your own pieces
	}
	int x, y, dx, dy;
	switch (actor) {
	case bP:
		if (src.y == RANK(7) && src.x == dst.x && dst.y == RANK(5)
			&& !target && BRD(src.x, 6) == empty) { // double
			return Normal;
		}
		if (dst.y != src.y+1) return Invalid;
		if (dst.x == src.x && !target) { // forward
			return Normal;
		}
		if (abs(dst.x - src.x) == 1 && target) { // taking
			return Normal;
		}
		// TODO: en passant
		if (dst.y == RANK(3) && abs(dst.x - src.x) == 1
			&& prev_move.piece == wP && eqpt(prev_move.src, Pt(dst.x, RANK(2)))
			&& prev_move.dst.y == RANK(4))
		{
			return EnPassant;
		}
		return Invalid;
	case wP:
		if (src.y == RANK(2) && src.x == dst.x && dst.y == RANK(4)
			&& !target && BRD(src.x, 3) == empty) { // double
			return Normal;
		}
		if (dst.y != src.y-1) return Invalid;
		if (dst.x == src.x && !target) { // forward
			return Normal;
		}
		if (abs(dst.x - src.x) == 1 && target) { // taking
			return Normal;
		}
		// en passant
		if (dst.y == RANK(6) && abs(dst.x - src.x) == 1
			&& prev_move.piece == bP && eqpt(prev_move.src, Pt(dst.x, RANK(7)))
			&& prev_move.dst.y == RANK(5))
		{
			return EnPassant;
		}
		return 0;
	case bK:
	case wK:
		if (abs(dst.x - src.x) <= 1 && abs(dst.y - src.y) <= 1) {
			return Normal;
		}
		// castling
		if (actor == wK) {
			if (eqpt(src, Pt(fE, RANK(1))) && dst.y == RANK(1)) {
				if (dst.x == fG) return WhiteShortCastle;
				if (dst.x == fC) return WhiteLongCastle;
			}
		} else {
			if (eqpt(src, Pt(fE, RANK(8))) && dst.y == RANK(8)) {
				if (dst.x == fG) return BlackShortCastle;
				if (dst.x == fC) return BlackLongCastle;
			}
		}		
		return Invalid;
	case bR:
	case wR:
		if (dst.x != src.x && dst.y != src.y) return 0;
		dx = 0;
		dy = 0;
		if (dst.x > src.x) dx = 1;
		else if (dst.x < src.x) dx = -1;
		else if (dst.y > src.y) dy = 1;
		else dy = -1;

		for (x=src.x+dx, y=src.y+dy; x!=dst.x || y!=dst.y; x+=dx, y+=dy)
			if (board[y][x]) return Invalid;
		return Normal;
	case bB:
	case wB:
		if (abs(dst.x - src.x) != abs(dst.y - src.y)) return 0;
		dx = dst.x > src.x ? 1 : -1;
		dy = dst.y > src.y ? 1 : -1;
		for (x=src.x+dx, y=src.y+dy; x!=dst.x; x+=dx, y+=dy)
			if (board[y][x]) return Invalid;
		return Normal;
	case bN:
	case wN:
		return ((abs(dst.x - src.x)==1 && abs(dst.y - src.y)==2) ||
			    (abs(dst.x - src.x)==2 && abs(dst.y - src.y)==1))
				? Normal : Invalid;
	case bQ:
	case wQ:
		if (dst.x == src.x || dst.y == src.y) { // moving like a rook
			dx = 0;
			dy = 0;
			if (dst.x > src.x) dx = 1;
			else if (dst.x < src.x) dx = -1;
			else if (dst.y > src.y) dy = 1;
			else dy = -1;
		} else { // moving like a bishop
			if (abs(dst.x - src.x) != abs(dst.y - src.y)) return 0;
			dx = dst.x > src.x ? 1 : -1;
			dy = dst.y > src.y ? 1 : -1;
		}
		for (x=src.x+dx, y=src.y+dy; x!=dst.x || y!=dst.y; x+=dx, y+=dy)
			if (board[y][x]) return Invalid;
		return Normal;
	}
	return Normal;
}

void
do_move(int board[8][8], Point src, Point dst, MoveType move_type) {
	BRDPT(dst) = BRDPT(src);
	BRDPT(src) = empty;
	if (move_type == EnPassant) {
		Point victim = addpt(dst, Pt(0, 1));
 		if (BRDPT(dst) == bP) {
			victim = subpt(dst, Pt(0, 1));
		}
		BRDPT(victim) = empty;
	}
	// promotoion
	if (BRDPT(dst) == bP) {
		if (move_type == PromotionQ) {
			BRDPT(dst) = bQ;
		} else if (move_type == PromotionN) {
			BRDPT(dst) = bN;
		} else if (move_type == PromotionB) {
			BRDPT(dst) = bB;
		} else if (move_type == PromotionR) {
			BRDPT(dst) = bR;
		}
	}
 	if (BRDPT(dst) == wP) {
		if (move_type == PromotionQ) {
			BRDPT(dst) = wQ;
		} else if (move_type == PromotionN) {
			BRDPT(dst) = wN;
		} else if (move_type == PromotionB) {
			BRDPT(dst) = wB;
		} else if (move_type == PromotionR) {
			BRDPT(dst) = wR;
		}
	}
	Point rook_src, rook_dst;
	switch (move_type) {
		case BlackLongCastle:
		case WhiteLongCastle:
			rook_src = Pt(fA, src.y);
			rook_dst = Pt(fD, src.y);
			BRDPT(rook_dst) = BRDPT(rook_src);
			BRDPT(rook_src) = empty;
			break;
		case BlackShortCastle:
		case WhiteShortCastle:
			rook_src = Pt(fH, src.y);
			rook_dst = Pt(fF, src.y);
			BRDPT(rook_dst) = BRDPT(rook_src);
			BRDPT(rook_src) = empty;
			break;
	}	
}

int
is_in_check(int board[8][8], Player current) {
	Piece our_king = bK;
	if (current == White) our_king = wK;
	int x, y;
	Point king_pos = Pt(-1, -1);
	for (x=0; x<8; x++) for (y=0; y<8; y++) {
		if (board[y][x] == our_king) {
			king_pos = Pt(x, y);
			break;
		}
	}
	if (eqpt(king_pos, Pt(-1, -1))) return 0; // no king found, whatever
	Move prev_move = {0}; // doesn't matter
	for (x=0; x<8; x++) for (y=0; y<8; y++) {
		if (board[y][x] && ismove(board, Pt(x, y), king_pos, prev_move)) {
			// this piece could take the king so we are in check
			return 1;
		}
	}
	return 0;
}

MoveType
islegal(int board[8][8], Point src, Point dst, Move prev_move, CastleRights cr)
{
	int tmpboard[8][8];
	memcpy(tmpboard, board, sizeof(int)*8*8);
	MoveType move_type = ismove(board, src, dst, prev_move);
	if (move_type == Invalid) {
		return Invalid;
	}
	Point s1, s2;
	if (move_type == BlackShortCastle) {
		if (cr.bK_ever_moved || cr.bRH_ever_moved) return Invalid;
		s1 = Pt(fF, RANK(8));
		s2 = Pt(fG, RANK(8));
	} else if (move_type == BlackLongCastle) {
		if (cr.bK_ever_moved || cr.bRA_ever_moved) return Invalid;
		s1 = Pt(fD, RANK(8));
		s2 = Pt(fC, RANK(8));
	} else if (move_type == WhiteShortCastle) {
		if (cr.wK_ever_moved || cr.wRH_ever_moved) return Invalid;
		s1 = Pt(fF, RANK(1));
		s2 = Pt(fG, RANK(1));
	} else if (move_type == WhiteLongCastle) {
		if (cr.wK_ever_moved || cr.wRA_ever_moved) return Invalid;
		s1 = Pt(fD, RANK(1));
		s2 = Pt(fC, RANK(1));
	}
	if (move_type >= BlackShortCastle) { // we check all castling the same
		if (islegal(tmpboard, src, s1, prev_move, cr)) {
			do_move(tmpboard, src, s1, Normal);
			if (islegal(tmpboard, s1, s2, prev_move, cr)) {
				// if the king could move there in two steps,
				// without going through check, then it's ok
				return move_type;
			}
		}
		return Invalid;
	}

	if (move_type == Normal && BRDPT(src) == bP && dst.y == RANK(1)) {
		move_type = PromotionQ;
	}
	if (move_type == Normal && BRDPT(src) == wP && dst.y == RANK(8)) {
		move_type = PromotionQ;
	}

	Player current = owner(BRDPT(src));
	// pretend to play the move, test if we are in check now
	do_move(tmpboard, src, dst, move_type);
	if (is_in_check(tmpboard, current)) {
		return Invalid;
	}
	return move_type;
}

int can_move(int board[8][8], Player player, Move prev_move, CastleRights cr) {
	for (int srcx=0; srcx<8; srcx++) for (int srcy=0; srcy<8; srcy++) {
		if (owner(board[srcy][srcx]) == player) {
			for (int dstx=0; dstx<8; dstx++) for (int dsty=0; dsty<8; dsty++) {
				if (islegal(board, Pt(srcx, srcy), Pt(dstx, dsty), prev_move, cr)) {
					return 1;
				}
			}
		}
	}
	return 0;
}

// longest possible move would be something like Qa1xe4+,
// so with null byte 8 characters is enough
#define MOVE_STR_LEN 8
void print_move(char* movestring, int board[8][8], Move move, Move prev_move, MoveType type, int was_capture, CastleRights cr) {
	char capturestr[2] = "x";
	char piece_id[3] = "";
	char promotion[3] = "";
	char check[2] = "";

	if (!was_capture) {
		capturestr[0] = 0;
	}
	if (move.piece == bP || move.piece == wP) {
		if (was_capture) {
			sprint(piece_id, "%c", 'A'+move.src.x);
		}
	} else {
		// TODO: check the board and remove redundant info when possible
		char src_file[2] = "";
		char src_rank[2] = "";
		int must_add_file = 0;
		int must_add_rank = 0;
		int others = 0;
		// if other piece of the same type which could also go there
		for (int x=0; x<8; x++) for (int y=0; y<8; y++) {
			if (x == move.src.x && y == move.src.y || board[y][x] != move.piece) continue;
			if (islegal(board, Pt(x, y), move.dst, prev_move, cr)) {
				others = 1;
				// if the rank is the same we must add the file, and vice versa
				if (y == move.src.y) must_add_file = 1;
				if (x == move.src.x) must_add_rank = 1;
			}
		}
		if (others && !must_add_file && !must_add_rank) must_add_file = 1;
		if (must_add_file) {
			sprint(src_file, "%c", 'a'+move.src.x);
		}
		if (must_add_rank) {
			sprint(src_rank, "%d", 8-move.src.y);
		}
		sprint(piece_id, "%c%s%s", PIECE_STR[move.piece][1], src_file, src_rank);
	}
	if (type == PromotionQ) {
		sprint(promotion, "=Q");
	} else if (type == PromotionN) {
		sprint(promotion, "=N");
	} else if (type == PromotionB) {
		sprint(promotion, "=B");
	} else if (type == PromotionR) {
		sprint(promotion, "=R");
	}
	// test if it will be check
	int tmpboard[8][8];
	memcpy(tmpboard, board, sizeof(int)*8*8);
	do_move(tmpboard, move.src, move.dst, type);
	if (is_in_check(tmpboard, !owner(move.piece))) {
		sprint(check, "+");
		if (!can_move(tmpboard, !owner(move.piece), move, cr))
			sprint(check, "#");
	}
	if (type == WhiteLongCastle || type == BlackLongCastle) {
		sprint(movestring, "O-O-O");
	} else if (type == WhiteShortCastle || type == BlackShortCastle) {
		sprint(movestring, "O-O");
	} else {
		sprint(movestring, "%s%s%c%d%s%s", piece_id,
			   capturestr, 'a'+move.dst.x, 8-move.dst.y, promotion, check);
	}
}

void
threadmain(int, char **)
{
	print("starting chess...\n");

	resizewin();

	if (initdraw(nil, nil, "chess") < 0)
		sysfatal("init failed: %r");

	resizewin();

	// LOAD ASSETS

	Image* pieceimg[12];
	for (int i=bB; i<=wR; i++) {
		char* fn = "AA.image"; // build the filename
		fn[0] = PIECE_STR[i][0];
		fn[1] = PIECE_STR[i][1];
		int piecefd = open(fn, OREAD);
		if (piecefd < 0)
			sysfatal("cannot open piece image: %r");
		pieceimg[i] = readimage(display, piecefd, 0);
		if (pieceimg[i] == nil)
			sysfatal("cannot read piece image: %r");
		close(piecefd);
	}
	Image* black = allocimagemix(display, DBlack, DBlack);
	Image* white = allocimagemix(display, DWhite, DWhite);
	Image* dark = allocimagemix(display, DRed, DDarkyellow);
	Image* light = allocimagemix(display, DWhite, DWhite);
	Image* bordercolor = allocimagemix(display, DWhite, DBlack);
	Image* selcolor = allocimagemix(display, DWhite, DRed);
	Image* poscolor = allocimagemix(display, DBlue, DWhite);

	// SETUP EVENTS

	Mousectl *mousectl;
	if ((mousectl = initmouse(nil, screen)) == nil)
		sysfatal("cannot init mouse: %r");

	Keyboardctl *keyboardctl;
	if ((keyboardctl = initkeyboard(nil)) == nil)
		sysfatal("cannot init keyboard: %r");

	enum {
		Emouse,
		Eresize,
		Ekbd,
	};
	Rune key;
	Mouse mouse;
	Alt a[] = {
		{ nil, &mouse, CHANRCV },
		{ nil, nil, CHANRCV },
		{ nil, &key, CHANRCV },
		{ nil, nil, CHANEND }
	};
	a[Emouse].c = mousectl->c;
	a[Eresize].c = mousectl->resizec;
	a[Ekbd].c = keyboardctl->c;

	dbgstring[0] = 0;

	int click = 0;
	int prev_click = 0;

	// SETUP GAME STATE

	int board[8][8];
	resetboard(board);
	Player turn = White;

	Point selected = Pt(-1, -1);
	Move prev_move = {0};
	CastleRights cr;
	memset(&cr, 0, sizeof(cr));

	Move* history = calloc(MAX_MOVES, sizeof(Move));
	char* movestrings = calloc(MAX_MOVES, MOVE_STR_LEN);
	int move_num = 0;

	// MAIN LOOP

	for(;;) {
		// RENDER
		draw(screen, screen->r, white, nil, ZP);
		Rectangle board_r = Rpt(addpt(screen->r.min, Pt(10, 10)),
							    addpt(screen->r.min, Pt(10+50*8, 10+50*8)));
		Rectangle border = Rpt(screen->r.min,
							   addpt(screen->r.min, Pt(20+50*8, 20+50*8)));
		draw(screen, border, bordercolor, nil, ZP);
		Rectangle r;
		for (int y=0; y<8; y++) for (int x=0; x<8; x++) {
			r = Rpt(addpt(board_r.min, Pt(50*x, 50*y)),
					addpt(board_r.min, Pt(50*x+50, 50*y+50)));
			if ((x+y) & 1)
				draw(screen, r, dark, nil, ZP);
			else
				draw(screen, r, light, nil, ZP);

			if (eqpt(Pt(x, y), selected)) {
				draw(screen, r, selcolor, nil, ZP);
			}
			if (!eqpt(selected, Pt(-1, -1)) && islegal(board, selected, Pt(x, y), prev_move, cr)) {
				draw(screen, r, poscolor, nil, ZP);
			}

			r = Rpt(addpt(r.min, Pt(5, 5)), r.max);
			int piece_n = board[y][x];
			if (piece_n) {
				Image* piece_i = pieceimg[piece_n];
				draw(screen, r, piece_i, 0, ZP);
			}
		}
		for (int i=0; i<8; i++) {
			// draw the rank numbers
			char buf;
			//sprint(buf, "%d", 8-i);
			buf = '8'-i;
			stringn(screen, addpt(board_r.min, Pt(2, 50*i+2)),
				    black, ZP, font, &buf, 1);
			// and the file letters
			buf = 'A' + i;
			stringn(screen, addpt(board_r.min, Pt(50*i+2, 50*7+35)),
				    black, ZP, font, &buf, 1);
		}
		string(screen, addpt(screen->r.min, Pt(50*8+25, 0)),
			   black, ZP, font, dbgstring);

		Point history_base =  addpt(screen->r.min, Pt(50*8+25, 5));
		for (int move = 0; move < move_num; move++) {
			char fullstr[20];
			char numstr[6] = "";
			if (move % 2 == 0) {
				sprint(numstr, "%d. ", move / 2 + 1);
			}
			char* movestring = &movestrings[move*MOVE_STR_LEN];
			sprint(fullstr, "%s%s", numstr, movestring);
			string(screen, addpt(history_base, Pt(move & 1 ? 90 : 0, move / 2 * 20)),
				   black, ZP, font, fullstr);
		}
		flushimage(display, 1);

		// PROCESS EVENTS
		char* menu_option_q = "Queen";
		char* menu_option_n = "Knight";
		char* menu_option_b = "Bishop";
		char* menu_option_r = "Rook";
		char* menu_options[5] = {
			menu_option_q, menu_option_n, menu_option_b, menu_option_r, nil
		};
		Menu promo_menu = {
			menu_options,
			nil,
			-1,
		};
		Point hovered;
		switch (alt(a)) {
			case Eresize:
				if (getwindow(display, Refnone) < 0)
					sysfatal("resize failed: %r");
				break;
			case Ekbd:
				switch (key) {
					case 'q':
						print("bye!\n");
						threadexitsall(nil);
						break;
					case 'r':
						// reset game state
						resetboard(board);
						turn = White;
						selected = Pt(-1, -1);
						prev_move.piece = empty;
						memset(&cr, 0, sizeof(cr));
						move_num = 0;
						break;
				}
				break;
			case Emouse:
				//dbgstring[0] = 0;
				if (ptinrect(mouse.xy, board_r)) {
					hovered = divpt(subpt(mouse.xy, board_r.min), 50);
					//sprint(dbgstring, "%d, %d, %d", hovered.x, hovered.y, BRDPT(hovered));
					int src_piece = empty;
					int dst_piece = empty;
					if (!eqpt(selected, Pt(-1, -1))) {
						src_piece = BRDPT(selected);
					}
					if (!eqpt(hovered, Pt(-1, -1))) {
						dst_piece = BRDPT(hovered);
					}
					click = mouse.buttons & 1;
					if (click && !prev_click) {
						MoveType move_type;
						if (eqpt(hovered, selected)) { // unselect
							selected = Pt(-1, -1);
						} else if (dst_piece && owner(dst_piece) == turn
								   && (!src_piece || owner(src_piece) == owner(dst_piece))) {
							selected = hovered; // select
						} else if (src_piece && (move_type = islegal(board, selected, hovered, prev_move, cr))) {
							 // move
							Move new_move;
							new_move.piece = src_piece;
							new_move.src = selected;
							new_move.dst = hovered;

							if (move_type == PromotionQ) {
								int promo_type = menuhit(1, mousectl, &promo_menu, nil);
								print("%d\n", promo_type);
								sprint(dbgstring, "%d", promo_type);
								if (promo_type >= 1)
									move_type = PromotionQ + promo_type;
							}

							if (move_num < MAX_MOVES) {
								history[move_num] = new_move;
								print_move(&movestrings[MOVE_STR_LEN*move_num], board,
										   new_move, prev_move, move_type, dst_piece, cr);
								move_num++;
							}

							do_move(board, selected, hovered, move_type);
							selected = Pt(-1, -1); // unselect
							if (src_piece == bK) cr.bK_ever_moved = 1;
							if (src_piece == wK) cr.wK_ever_moved = 1;
							if (eqpt(new_move.src, Pt(0, 0)) || eqpt(new_move.dst, Pt(0, 0)))
								cr.bRA_ever_moved = 1;
							if (eqpt(new_move.src, Pt(0, 7)) || eqpt(new_move.dst, Pt(0, 7)))
								cr.wRA_ever_moved = 1;
							if (eqpt(new_move.src, Pt(7, 0)) || eqpt(new_move.dst, Pt(7, 0)))
								cr.bRH_ever_moved = 1;
							if (eqpt(new_move.src, Pt(7, 7)) || eqpt(new_move.dst, Pt(7, 7)))
								cr.wRH_ever_moved = 1;
							prev_move = new_move;
							turn = !turn;
						}
					}
					prev_click = click;
				}
				break;
		}
	}
}