#lang racket
(require debug)
(define stdout (current-output-port))
(define ((appls f) l) (apply f l))
(define falses (stream-cons #f falses))
(define member? (compose1 not not member))
(define (char+ . args)
  (integer->char (apply + (map (λ(x)
                                 (cond
                                   [(char? x) (char->integer x)]
                                   [(integer? x) x])) args))))
(define (char- . args)
  (integer->char (apply - (map (λ(x)
                                 (cond
                                   [(char? x) (char->integer x)]
                                   [(integer? x) x])) args))))

#| velocities of all playable characters except pawn |#
(define knight-vel (append-map permutations(cartesian-product'(-2 2)'(-1 1))))
(define bishop-vel (cartesian-product '(-1 1) '(-1 1)))
(define rook-vel (append-map permutations '((1 0) (-1 0))))
(define queen-vel (append bishop-vel rook-vel))
(define king-vel queen-vel)

(define pretty-assoc
  (map list (map (compose car string->list symbol->string)
                 (append '(K Q R B N P k q r b n p | | _)))
       (map integer->char (append (range 9812 9824)
                                  '(#x3000 #x25a0)))))

;; possible black squares: #x25a0 ,

(define (empty-piece . args)
  (match-define (or `(,x ,y) `((,x ,y))) args)
  (if (= 0 (modulo (- x y) 2)) #\space #\_))

(define empty-board
  (build-list 8 (λ(y) (build-list 8 (λ(x) (empty-piece x y))))))

(define black-piece? char-lower-case?)
(define white-piece? char-upper-case?)

(define (update-cell brd f . args)
  (match-define (or `(,x ,y) `((,x ,y))) args)
  (list-update brd y (curryr list-update x f)))

(define (get-cell brd . args)
  (match-define (or `(,x ,y) `((,x ,y))) args)
  (list-ref (list-ref brd y) x))

(define (set-row brd y to-set)
  (list-set brd y to-set))

(define start-board
  (foldl
   (λ (args b)
     (set-row b (car args)
              (map (compose car string->list symbol->string)
               (cadr args))))
   empty-board
   `((0 (r n b q k b n r))
     (1 (p p p p p p p p))
     (6 (P P P P P P P P))
     (7 (R N B Q K B N R)))))

(define (place-piece brd p . args)
  (match-define (or `(,x ,y) `((,x ,y)) `((,x . ,y))) args)
  (update-cell brd (const (if p p (empty-piece x y))) x y))

(define test-board
  (foldl (λ (args b) (place-piece b (car args) (cdr args)))
   empty-board '((#\Q 3 3) (#\b 5 5) (#\B 5 3) (#\K 0 0)
                           (#\k 7 7) (#\r 1 7) (#\r 7 1))))

(define (valid-coord? . args)
  (match-define (or `(,x ,y) `((,x ,y))) args)
  (and (< -1 x 8) (< -1 y 8)))

(define (safe-take l n)
  (if (or (null? l) (<= n 0)) '()
      (cons (car l) (safe-take (cdr l) (sub1 n)))))
(define (safe-last ls) (if (null? ls) #f (last ls)))

(define (get-coords brd char)
  (for/list {[row brd]  [y (in-naturals)] #:when #t
             [cell row] [x (in-naturals)] #:when (char=? char cell)}
    `(,x ,y)))

(define (in-2D-range start vel)
  (let loop {[cur start]}
    (stream-cons cur (loop (map + cur vel)))))

;; you're only allowed to promote to queen
(define (do-move brd . args)
  (match-define (or `(,x,y,x2,y2) `((,x,y)(,x2,y2))) args)
  (define p (get-cell brd x y))
  (place-piece
   (place-piece brd #f x y)
   (match `(,p,y2)
     ['(#\p 7) #\q]
     ['(#\P 0) #\Q]
     [else p]) x2 y2))

;; requires: x and y in args are valid-coord?
(define (valid-moves brd . args)
  (match-define (or `(,x ,y) `((,x ,y))) args)
  (define start `(,x,y))
  (define p (get-cell brd start))
  (define black? (black-piece? p))
  (define ally? (if black? black-piece? white-piece?))
  (define opponent? (if black? white-piece? black-piece?))

  (define (travel-until brd start vel [black? black?])
    (define ally? (if black? black-piece? white-piece?))
    (define opponent? (if black? white-piece? black-piece?))
    (for/list
        {[pos (stream-rest (in-2D-range start vel))]
         #:break (or (not (valid-coord? pos))
                     (ally? (get-cell brd pos)))
         #:final (opponent? (get-cell brd pos))}
      pos))

  ;; pawn moves are hard... en-poisson not implemented
  (define (calc-pawn-moves)
    (define vert-dir (if black? 1 -1))
    (define vert-scalar
      (if (or (and black? (= y 1))
              (and (not black?) (= y 6)))
          2 1))
    (append
     (safe-take (travel-until brd start `(0 ,vert-dir)) vert-scalar)
     (filter-map
      (λ(x)
        (and (valid-coord? (map + start `(,x ,vert-dir)))
             (opponent?
              (get-cell brd (map + start `(,x ,vert-dir))))
             `(,x ,vert-dir)))
      '(-1 1))))

  (define (valid+free? pos)
    (and (valid-coord? pos) (not (ally? (get-cell brd pos)))))

  (define (raw-valid-moves)
    (match (char-downcase p)
      [#\r (append-map (curry travel-until brd start)   rook-vel)]
      [#\b (append-map (curry travel-until brd start) bishop-vel)]
      [#\q (append-map (curry travel-until brd start)  queen-vel)]
      [#\n (filter valid+free?(map(curry map + start) knight-vel))]
      [#\k (filter valid+free?(map(curry map + start)   king-vel))]
      [#\p (calc-pawn-moves)]))

  (define (in-check? brd)
    (define king (if black? #\k #\K))
    (define king-pos (car (get-coords brd king)))
    (define opponent? (if black? white-piece? black-piece?))
    (define pawn-vert (if black? 1 -1))

    (define ((offending-pieces-here? bad-pieces) opp-pos)
      (define opp (get-cell brd opp-pos))
      (and (opponent? opp) (memv (char-downcase opp) bad-pieces)))

    (or
     (ormap
      (offending-pieces-here? '(#\r #\q))
      (filter-map (λ (vel)
                    (safe-last (travel-until brd king-pos vel)))
                  rook-vel))
     (ormap
      (offending-pieces-here? '(#\b #\q))
      (filter-map (λ (vel)
                    (safe-last (travel-until brd king-pos vel)))
                  bishop-vel))
     (ormap
      (offending-pieces-here? '(#\n))
      (filter valid-coord?
              (map (curry map + king-pos) knight-vel)))
     (ormap
      (offending-pieces-here? '(#\p))
      (filter valid-coord?
              (map
               (λ(x) (map + king-pos `(,x ,pawn-vert)))
               '(-1 1))))))

  (define/debug (try-move . args2)
    (match-define (or `(,x2 ,y2) `((,x2 ,y2))) args2)
    (place-piece
     (place-piece brd #f x y)
     (match `(,p,y2)
       ['(#\p 7) #\q]
       ['(#\P 0) #\Q]
       [else p])
     x2 y2))

  (filter (compose not in-check? try-move) (raw-valid-moves)))

(define (print-stars brd ls)
  (begin
    (printf "\n\u3000~a\n"
            (string-join (map (λ(i) (string (char+ #\uff21 i))) (range 8))
                         "\u3000" #:before-first "\u3000" #:after-last "\u3000"))
    (printf "\u3000┏~a┓\n" (string-join
             (build-list 8 (λ(x) (if (member `(,x 0) ls) "┳" "━"))) "┳"))
    (for {[row brd]
          [y (in-naturals)]}
      (printf "~a" (char- #\uFF18 y))
      (for {[cell row]
            [x (in-naturals)]}
        (match (list (member? `(,x,y) ls) (member? `(,(- x 1),y) ls))
          ['(#t #t)
           (printf "╋~a" (cadr (assoc cell pretty-assoc)))]
          ['(#t #f)
           (printf "┣~a" (cadr (assoc cell pretty-assoc)))]
          ['(#f #t)
           (printf "┫~a" (cadr (assoc cell pretty-assoc)))]
          [else
           (printf "┃~a" (cadr (assoc cell pretty-assoc)))]))
      (if (member `(7,y) ls)
          (printf "┫~a\n\u3000" (char- #\uFF18 y))
          (printf "┃~a\n\u3000" (char- #\uFF18 y)))
      (if (< y 7)
          (printf "┣~a┫\n" (string-join (build-list
                    8 (λ(x)
                        (match (list(member? `(,x,y) ls)(member?`(,x,(+ y 1)) ls))
                          ['(#t #t) "╋"]
                          ['(#t #f) "┻"]
                          ['(#f #t) "┳"]
                          [else "━"]))) "╋"))
          (printf "┗~a┛\n" (string-join
                   (build-list 8 (λ(x) (if (member `(,x,y) ls) "┻" "━"))) "┻"))))
    (printf "\u3000~a\n\n"
            (string-join (map (λ(i) (string (char+ #\uff21 i))) (range 8)) "\u3000"
                         #:before-first "\u3000" #:after-last "\u3000"))))

(define (symbols->strings lol)
  (cond
    [(pair? lol) (map symbols->strings lol)]
    [(symbol? lol) (symbol->string lol)]
    [else lol]))

(define (parse-move in)
  (match (symbols->strings in)
    [`(,x,y) #:when (and (integer? x) (integer? y)) in]
    [(regexp "([a-hA-h])([1-8])" (list _ x y))
     (list (char->integer (char- (char-upcase (car (string->list x))) #\A))
           (- 8 (string->number y)))]
    [_ '(-1 -1)]))

(define (all-valid-moves brd [black? #f])
  (define ally? (if black? black-piece? white-piece?))
  (for/fold
   {[acc '()]}
   {[row brd]
    [y (in-naturals)] #:when #t
    [cell row]
    [x (in-naturals)] #:when (ally? cell)}
    (append (valid-moves brd x y) acc)))

; #|
;; plays a game. returns void if successful, and 'quit if the user quit.
(define {game [brdls `(,start-board)] [black? #f]}
  (define ally? (if black? black-piece? white-piece?))
  (define brd (car brdls))
  (print-stars brd '())
  (define all-moves (all-valid-moves brd black?))
  (cond
    [(null? all-moves)
     (printf "~a is in checkmate! ~a wins!\nRestart? (y/n)\n"
             (if black? 'Black 'White) (if black? 'White 'Black))
     (match (read)
       [`y (printf "restarting.\n") (game `(,start-board) #f)]
       [`n (printf "quiting.\n")])]
    [else
     (printf "~a's turn!\n" (if black? 'Black 'White))
     (let prompt {}
       (match (read)
         [`move
          (define start (parse-move (read)))
          (cond
            [(or (not (valid-coord? start))
                 (not (ally? (get-cell brd start))))
             (printf "Please input a valid coordinate.\n") (prompt)]
            [(let {[valid-ls (valid-moves brd start)]}
               (print-stars brd (cons start valid-ls))
               (define end (parse-move (read)))
               (and (member end valid-ls) end))
             => (λ (end) (game (cons (do-move brd start end) brdls) (not black?)))]
            [else
             (printf "Please input a valid move.\n") (prompt)])]
         [`undo (printf "Undoing. Please wait 3-5 minutes.\n") (time (sleep 10))
          (game (cdr brdls) (not black?))]
         [`resign (printf "Restarting.\n")
          (game `(,start-board) #f)]
         [(or (? eof-object?) `quit) 'quit]
         [_ (printf "Didn't understand. Please enter another command.\n")
          (prompt)]))])) ; |#

(define (setup [brd empty-board])
  3)

(define (main)
  (define (quit) (printf "Quiting.\n"))
  (printf "
Welcome to Chess!
Enter \"setup\" if you'd like to setup the board.
Enter \"play\" if you'd like to play a game.
Enter \"quit\" at any time to quit the program.
")

  (define first-setup? (box #t))
  (define first-game? (box #t))

  (let menu {[brd empty-board]}
    (match (read)
      [`setup (menu (setup brd))]
      [`play (define res (game (list (if (equal? brd empty-board) start-board brd))))
             (if (eq? res 'quit) (quit) (menu brd))]
      [`quit (quit)]
      [_ (printf "Didn't understand. Please enter another command.\n")
         (menu brd)])))