#include "chess.h"
#include <string.h>
#include <stdio.h>
#include <ctype.h>
#include <stdlib.h>

static const int knight_offsets[8] = {15, 17, -15, -17, 10, -10, 6, -6};
static const int bishop_dirs[4] = {9, 7, -9, -7};
static const int rook_dirs[4] = {8, -8, 1, -1};
static const int king_dirs[8] = {8,-8,1,-1,9,7,-9,-7};

static inline int file_of(int sq){return sq%8;}
static inline int rank_of(int sq){return sq/8;}
static inline bool on_board(int sq){return sq>=0 && sq<64;}
static inline bool same_color(char a, char b){return (isupper(a)&&isupper(b))||(islower(a)&&islower(b));}
static inline bool is_white(char p){return isupper((unsigned char)p);} 

void chess_init_start(Position *pos){
    // a1..h1 (0..7) are white back rank; rank 8 at indexes 56..63 are black back rank
    const char *start = "RNBQKBNRPPPPPPPP................................pppppppprnbqkbnr";
    for (int i=0;i<64;i++) pos->board[i] = start[i];
    pos->white_to_move = true;
    pos->castle_wk = pos->castle_wq = pos->castle_bk = pos->castle_bq = true;
    pos->ep_square = -1; pos->halfmove_clock = 0; pos->fullmove_number = 1;
}

void chess_copy(Position *dst, const Position *src){ *dst = *src; }

static bool is_empty(const Position *p, int sq){ return p->board[sq]=='.'; }

bool chess_square_attacked(const Position *pos, int sq, bool by_white){
    // pawns
    int r = rank_of(sq), f = file_of(sq);
    if (by_white){
        int s1 = (r-1)*8 + (f-1); if (f>0 && r>0 && on_board(s1) && pos->board[s1]=='P') return true;
        int s2 = (r-1)*8 + (f+1); if (f<7 && r>0 && on_board(s2) && pos->board[s2]=='P') return true;
    } else {
        int s1 = (r+1)*8 + (f-1); if (f>0 && r<7 && on_board(s1) && pos->board[s1]=='p') return true;
        int s2 = (r+1)*8 + (f+1); if (f<7 && r<7 && on_board(s2) && pos->board[s2]=='p') return true;
    }
    // knights
    for (int i=0;i<8;i++){
        int t = sq + knight_offsets[i];
        if (!on_board(t)) continue;
        int df = file_of(t)-f; int dr = rank_of(t)-r; if (df< -2||df>2||dr<-2||dr>2) continue; // edge wrap guard
        char pc = pos->board[t];
        if (by_white && pc=='N') return true;
        if (!by_white && pc=='n') return true;
    }
    // bishops/queens
    for (int d=0; d<4; d++){
        int off = bishop_dirs[d]; int t = sq + off;
        while (on_board(t) && abs(file_of(t)-f)==abs(rank_of(t)-r)){
            char pc = pos->board[t];
            if (pc!='.') { if (by_white && (pc=='B'||pc=='Q')) return true; if (!by_white && (pc=='b'||pc=='q')) return true; break; }
            t += off;
        }
    }
    // rooks/queens
    for (int d=0; d<4; d++){
        int off = rook_dirs[d]; int t = sq + off;
        while (on_board(t) && (file_of(t)==f || rank_of(t)==r)){
            char pc = pos->board[t];
            if (pc!='.') { if (by_white && (pc=='R'||pc=='Q')) return true; if (!by_white && (pc=='r'||pc=='q')) return true; break; }
            t += off;
        }
    }
    // king
    for (int i=0;i<8;i++){
        int t=sq+king_dirs[i]; if (!on_board(t)) continue; if (abs(file_of(t)-f)>1||abs(rank_of(t)-r)>1) continue;
        char pc = pos->board[t]; if (by_white && pc=='K') return true; if (!by_white && pc=='k') return true;
    }
    return false;
}

bool chess_is_in_check(const Position *pos, bool white){
    int ks = -1; char k = white? 'K' : 'k';
    for (int i=0;i<64;i++) if (pos->board[i]==k) { ks = i; break; }
    if (ks==-1) return false; // malformed
    return chess_square_attacked(pos, ks, !white);
}

static void add_move_if_legal(const Position *pos, int from, int to, char promo, Move *out, size_t *n, size_t max){
    if (*n >= max) return;
    Position tmp; chess_copy(&tmp, pos);
    Move m = {0}; m.from=from; m.to=to; m.promo=promo; m.moved=pos->board[from]; m.captured=pos->board[to];
    int prev_ep = tmp.ep_square; m.prev_ep = prev_ep; m.prev_wk=tmp.castle_wk; m.prev_wq=tmp.castle_wq; m.prev_bk=tmp.castle_bk; m.prev_bq=tmp.castle_bq; m.prev_halfmove=tmp.halfmove_clock;
    if (!chess_make_move(&tmp, &m)) return;
    if (chess_is_in_check(&tmp, !tmp.white_to_move)) return; // after make, side switched
    out[(*n)++] = m; // store pseudo move with added flags from make_move
}

size_t chess_generate_legal_moves(const Position *pos, Move *out, size_t max){
    size_t n=0;
    bool white = pos->white_to_move;
    for (int sq=0; sq<64; sq++){
        char p = pos->board[sq]; if (p=='.') continue; if (white != is_white(p)) continue;
        int f=file_of(sq), r=rank_of(sq);
        switch (tolower(p)){
            case 'p': {
                int dir = white? 8 : -8; int start_rank = white? 1 : 6; int prom_rank = white? 6 : 1;
                int one = sq + dir;
                if (on_board(one) && is_empty(pos, one)){
                    if (r==prom_rank){
                        const char *pr = white? "QRBN" : "qrbn";
                        for (int i=0;i<4;i++) add_move_if_legal(pos, sq, one, pr[i], out, &n, max);
                    } else add_move_if_legal(pos, sq, one, 0, out, &n, max);
                    // two
                    int two = sq + 2*dir; if (r==start_rank && is_empty(pos, two)) add_move_if_legal(pos, sq, two, 0, out, &n, max);
                }
                // captures
                int caps[2] = { dir+1, dir-1 };
                for (int i=0;i<2;i++){
                    int t = sq + caps[i]; if (!on_board(t)) continue; if (abs(file_of(t)-f)!=1) continue;
                    if (!is_empty(pos, t) && !same_color(pos->board[sq], pos->board[t])){
                        if (r==prom_rank){ const char *pr = white? "QRBN" : "qrbn"; for (int j=0;j<4;j++) add_move_if_legal(pos, sq, t, pr[j], out, &n, max); }
                        else add_move_if_legal(pos, sq, t, 0, out, &n, max);
                    }
                }
                // en passant
                if (pos->ep_square!=-1){
                    int ep = pos->ep_square; if (abs(file_of(ep)-f)==1 && (ep == sq+dir+1 || ep==sq+dir-1)) add_move_if_legal(pos, sq, ep, 0, out, &n, max);
                }
            } break;
            case 'n': {
                for (int i=0;i<8;i++){
                    int t = sq + knight_offsets[i]; if (!on_board(t)) continue; if (abs(file_of(t)-f)>2 || abs(rank_of(t)-r)>2) continue; if (!is_empty(pos, t) && same_color(p, pos->board[t])) continue; add_move_if_legal(pos, sq, t, 0, out, &n, max);
                }
            } break;
            case 'b': {
                for (int d=0; d<4; d++){
                    int off = bishop_dirs[d]; int t = sq + off; while (on_board(t) && abs(file_of(t)-f)==abs(rank_of(t)-r)){
                        if (!is_empty(pos, t)){ if (!same_color(p, pos->board[t])) add_move_if_legal(pos, sq, t, 0, out, &n, max); break; }
                        add_move_if_legal(pos, sq, t, 0, out, &n, max); t += off;
                    }
                }
            } break;
            case 'r': {
                for (int d=0; d<4; d++){
                    int off = rook_dirs[d]; int t = sq + off; while (on_board(t) && (file_of(t)==f || rank_of(t)==r)){
                        if (!is_empty(pos, t)){ if (!same_color(p, pos->board[t])) add_move_if_legal(pos, sq, t, 0, out, &n, max); break; }
                        add_move_if_legal(pos, sq, t, 0, out, &n, max); t += off;
                    }
                }
            } break;
            case 'q': {
                for (int d=0; d<4; d++){
                    int off = bishop_dirs[d]; int t = sq + off; while (on_board(t) && abs(file_of(t)-f)==abs(rank_of(t)-r)){
                        if (!is_empty(pos, t)){ if (!same_color(p, pos->board[t])) add_move_if_legal(pos, sq, t, 0, out, &n, max); break; }
                        add_move_if_legal(pos, sq, t, 0, out, &n, max); t += off;
                    }
                }
                for (int d=0; d<4; d++){
                    int off = rook_dirs[d]; int t = sq + off; while (on_board(t) && (file_of(t)==f || rank_of(t)==r)){
                        if (!is_empty(pos, t)){ if (!same_color(p, pos->board[t])) add_move_if_legal(pos, sq, t, 0, out, &n, max); break; }
                        add_move_if_legal(pos, sq, t, 0, out, &n, max); t += off;
                    }
                }
            } break;
            case 'k': {
                for (int i=0;i<8;i++){
                    int t = sq + king_dirs[i]; if (!on_board(t)) continue; if (abs(file_of(t)-f)>1||abs(rank_of(t)-r)>1) continue; if (!is_empty(pos, t) && same_color(p, pos->board[t])) continue; add_move_if_legal(pos, sq, t, 0, out, &n, max);
                }
                // castling
                if (white){
                    if (pos->castle_wk && pos->board[5]=='.' && pos->board[6]=='.' && !chess_square_attacked(pos,4,false) && !chess_square_attacked(pos,5,false) && !chess_square_attacked(pos,6,false)) add_move_if_legal(pos, 4, 6, 0, out, &n, max);
                    if (pos->castle_wq && pos->board[3]=='.' && pos->board[2]=='.' && pos->board[1]=='.' && !chess_square_attacked(pos,4,false) && !chess_square_attacked(pos,3,false) && !chess_square_attacked(pos,2,false)) add_move_if_legal(pos, 4, 2, 0, out, &n, max);
                } else {
                    if (pos->castle_bk && pos->board[61]=='.' && pos->board[62]=='.' && !chess_square_attacked(pos,60,true) && !chess_square_attacked(pos,61,true) && !chess_square_attacked(pos,62,true)) add_move_if_legal(pos, 60, 62, 0, out, &n, max);
                    if (pos->castle_bq && pos->board[59]=='.' && pos->board[58]=='.' && pos->board[57]=='.' && !chess_square_attacked(pos,60,true) && !chess_square_attacked(pos,59,true) && !chess_square_attacked(pos,58,true)) add_move_if_legal(pos, 60, 58, 0, out, &n, max);
                }
            } break;
        }
    }
    return n;
}

bool chess_make_move(Position *pos, Move *m){
    m->is_castle=false; m->is_enpassant=false;
    char p = pos->board[m->from]; char tgt = pos->board[m->to];
    if (p=='.') return false;
    // handle special: en passant capture
    if (tolower(p)=='p' && m->to == pos->ep_square && file_of(m->to)!=file_of(m->from) && tgt=='.'){
        m->is_enpassant = true;
        int cap_sq = pos->white_to_move ? (m->to - 8) : (m->to + 8);
        m->captured = pos->board[cap_sq]; pos->board[cap_sq]='.';
    }

    // move piece
    pos->board[m->to] = p;
    pos->board[m->from] = '.';

    // promotion
    if (tolower(p)=='p' && m->promo){ pos->board[m->to] = m->promo; }

    // castling rook move
    if (tolower(p)=='k'){
        int from = m->from, to = m->to;
        if (from==4 && to==6){ pos->board[5]='R'; pos->board[7]='.'; m->is_castle=true; }
        else if (from==4 && to==2){ pos->board[3]='R'; pos->board[0]='.'; m->is_castle=true; }
        else if (from==60 && to==62){ pos->board[61]='r'; pos->board[63]='.'; m->is_castle=true; }
        else if (from==60 && to==58){ pos->board[59]='r'; pos->board[56]='.'; m->is_castle=true; }
    }

    // update castling rights
    if (m->from==0||m->to==0) pos->castle_wq=false;
    if (m->from==7||m->to==7) pos->castle_wk=false;
    if (m->from==56||m->to==56) pos->castle_bq=false;
    if (m->from==63||m->to==63) pos->castle_bk=false;
    if (tolower(p)=='k'){ if (is_white(p)) { pos->castle_wk=pos->castle_wq=false;} else {pos->castle_bk=pos->castle_bq=false;} }

    // update ep square
    pos->ep_square = -1;
    if (tolower(p)=='p'){
        int df = rank_of(m->to) - rank_of(m->from);
        if (df==2 || df==-2){ pos->ep_square = (m->from + m->to)/2; }
    }

    // halfmove clock
    if (tolower(p)=='p' || tgt!='.') pos->halfmove_clock = 0; else pos->halfmove_clock++;

    // side to move
    pos->white_to_move = !pos->white_to_move; if (pos->white_to_move) pos->fullmove_number++;

    return true;
}

void chess_unmake_move(Position *pos, const Move *m){
    pos->white_to_move = !pos->white_to_move; if (!pos->white_to_move) pos->fullmove_number--;
    // restore halfmove/flags
    pos->ep_square = m->prev_ep;
    pos->castle_wk = m->prev_wk; pos->castle_wq = m->prev_wq; pos->castle_bk = m->prev_bk; pos->castle_bq = m->prev_bq;
    pos->halfmove_clock = m->prev_halfmove;

    char p = m->moved;
    // undo promotions
    if (tolower(p)=='p' && m->promo){ p = is_white(p)? 'P':'p'; }

    pos->board[m->from] = p;
    pos->board[m->to] = m->captured? m->captured : '.';
    if (m->is_enpassant){ int cap_sq = pos->white_to_move ? (m->to - 8) : (m->to + 8); pos->board[m->to]='.'; pos->board[cap_sq]= m->captured; }
    if (m->is_castle){
        if (m->from==4 && m->to==6){ pos->board[7]='R'; pos->board[5]='.'; }
        else if (m->from==4 && m->to==2){ pos->board[0]='R'; pos->board[3]='.'; }
        else if (m->from==60 && m->to==62){ pos->board[63]='r'; pos->board[61]='.'; }
        else if (m->from==60 && m->to==58){ pos->board[56]='r'; pos->board[59]='.'; }
    }
}

void sq_to_coord(int sq, int *file, int *rank){ if (file) *file=file_of(sq); if (rank) *rank=rank_of(sq); }
int coord_to_sq(int file, int rank){ return rank*8 + file; }

void move_to_uci(const Move *m, char buf[8]){
    int f1=file_of(m->from), r1=rank_of(m->from), f2=file_of(m->to), r2=rank_of(m->to);
    buf[0]='a'+f1; buf[1]='1'+r1; buf[2]='a'+f2; buf[3]='1'+r2; int i=4;
    if (m->promo){ buf[i++]=tolower(m->promo); }
    buf[i]='\0';
}

bool parse_uci_move(const char *s, const Position *pos, Move *out){
    if (!s || strlen(s)<4) return false;
    int f1=s[0]-'a', r1=s[1]-'1', f2=s[2]-'a', r2=s[3]-'1'; if (f1<0||f1>7||f2<0||f2>7||r1<0||r1>7||r2<0||r2>7) return false;
    int from = r1*8+f1, to=r2*8+f2; char promo = s[4]? s[4]:0; if (promo) promo = pos->white_to_move? toupper((unsigned char)promo):tolower((unsigned char)promo);
    Move list[MAX_MOVES]; size_t n = chess_generate_legal_moves(pos, list, MAX_MOVES);
    for (size_t i=0;i<n;i++){ if (list[i].from==from && list[i].to==to){ *out = list[i]; out->promo = promo? promo : list[i].promo; return true; } }
    return false;
}

bool chess_to_fen(const Position *pos, char *out, size_t outsz){
    char buf[256]; int idx=0;
    for (int r=7;r>=0;r--){
        int empty=0;
        for (int f=0;f<8;f++){
            char p = pos->board[r*8+f];
            if (p=='.') empty++; else { if (empty){ buf[idx++]= '0'+empty; empty=0; } buf[idx++]=p; }
        }
        if (empty) buf[idx++]= '0'+empty;
        if (r) buf[idx++]='/';
    }
    buf[idx++]=' ';
    buf[idx++]= pos->white_to_move ? 'w':'b';
    buf[idx++]=' ';
    int start=idx;
    if (pos->castle_wk) buf[idx++]='K';
    if (pos->castle_wq) buf[idx++]='Q';
    if (pos->castle_bk) buf[idx++]='k';
    if (pos->castle_bq) buf[idx++]='q';
    if (idx==start) buf[idx++]='-';
    buf[idx++]=' ';
    if (pos->ep_square==-1){ buf[idx++]='-'; }
    else { int f=file_of(pos->ep_square), r=rank_of(pos->ep_square); buf[idx++]='a'+f; buf[idx++]='1'+r; }
    idx += snprintf(buf+idx, sizeof(buf)-idx, " %d %d", pos->halfmove_clock, pos->fullmove_number);
    buf[idx]='\0';
    snprintf(out, outsz, "%s", buf);
    return true;
}