#!/usr/bin/env ruby
# -*- Mode: Ruby -*- vim: sts=3 sw=3 expandtab

## A simple game reminiscent of Connect Four.
## Author:   Thomas Kirchner, 2005
## Homepage: http://halffull.org/span
## License:  http://creativecommons.org/licenses/by-nc-sa/2.0/
NAME = "span"
VER  = "0.1.0"

require 'ncurses'
require 'optparse'

class Board
   attr_reader :width, :height, :chars, :state, :next_row

   def initialize(game, width = 7, height = 6)
      @game = game
      @width, @height = width, height
      @next_row = []  # lowest available row by column
      @state = []     # board state, i.e. pieces
      (0...width).each {|w| @state[w] = [] }  # 2d
      @chars = %w{_ x o + s e & # @ ~}  # player markers.  0 represents blank
   end

   def clear
      @next_row.fill(0, 0..width)
      @the_win = []
      (0...@width).each {|w| @state[w].fill(@chars[0], 0...@height) }
   end

   def draw
      y = @game.screen.height / 2 - @height / 2
      x = @game.screen.width / 2 - @width
      Ncurses::move(y,x)
      (@height - 1).downto(0) { |h|  # print each row, starting at top
         (0...@width).each { |w|     # print each column from the left
            color = Ncurses::COLOR_PAIR(@chars.index(@state[w][h]) % @game.screen.num_colors)
            color |= Ncurses::A_BOLD if @the_win and @the_win.include?([w,h])
            Ncurses::attrset(color); Ncurses::printw(@state[w][h])
            Ncurses::attrset(Ncurses::A_NORMAL); Ncurses::printw(" ")
         }
         Ncurses::move(y += 1,x)  # move to next screen row
      }
      Ncurses::refresh
   end

   def play(player, column)
      if column.is_a?(Integer) and column >= 0 and column < @width and @next_row[column] < @height
         @state[column][@next_row[column]] = @chars[player.number] || "?"
         @next_row[column] += 1
      else raise ArgumentError, "Error in placement."
      end
   end

   def over?  # If the board has a winner, return the winner, else false.
      span? || tie?
   end

   # req = the count to check for.  players = array of player numbers, 1+
   # for AI: will return [x,y] to play if available
   def span?(req = @game.REQ, players = :all)
      vert_span?(req, players) || horiz_span?(req, players) || diag_span?(req, players)
   end

   def vert_span?(req, players)
      count = 1
      last_char = ""
      @state.each_with_index { |col, c|  # traverse rows
         col.each_with_index { |row, r|  # and up columns
            if row != @chars[0] and last_char == row and (players == :all or players.include?(@chars.index(row)))
               count += 1
            else count = 1
            end

            last_char = row
            if count >= req  # a winner is you
               if req == @game.REQ
                  @the_win = []  # array of winning pieces - [x, y]
                  (0...count).each {|i| @the_win[i] = [c, r - i]}
                  return row  # this was an over? check
               elsif @state[c][r + 1] == @chars[0]  # for AI use
                  return [c, r + 1]
               end
            end
         }
         last_char = ""
      }; false
   end

   def horiz_span?(req, players)
      count = 1
      last_char = ""
      (0...@state.first.size).each { |row|  # traverse up columns
         (0...@state.size).each { |col|     # and across rows
            cur = @state[col][row]
            if cur != @chars[0] and last_char == cur and (players == :all or players.include?(@chars.index(cur)))
               count += 1
            else count = 1
            end

            last_char = cur
            if count >= req
               if req == @game.REQ
                  @the_win = []
                  (0...count).each {|i| @the_win[i] = [col - i, row]}
                  return cur
               # AI: check right + left for available play
               elsif @state[col + 1] and @state[col + 1][row] == @chars[0] and (row == 0 or @state[col + 1][row - 1] != @chars[0])
                  return [col + 1, row]
               elsif col >= req and @state[col - req][row] == @chars[0] and (row == 0 or @state[col - req][row - 1] != @chars[0])
                  return [col - req, row]
               end
            elsif count == req - 1 and req + 1 >= @game.REQ and req != @game.REQ
               # intercept xx_x style patterns
               if @state[col + 2] and @state[col + 1][row] == @chars[0] and @state[col + 2][row] == cur and @next_row[col + 1] == row
                  return [col + 1, row]
               elsif col >= req and @state[col - req + 1][row] == @chars[0] and @state[col - req][row] == cur and @next_row[col - req + 1] == row
                  return [col - req + 1, row]
               end
            end
         }
         last_char = ""
      }; false
   end

   def diag_span?(req, players)
      (0...@state.first.size).each { |row|  # traverse up columns
         (0...@state.size).each { |col|     # and across rows
            if (row - 1 + @game.REQ) < @state.first.size  # too high = can't win
               cur = @state[col][row]
               count = 1
               if (col + 1 - @game.REQ) >= 0  # an up-left match is possible
                  (1...req).each { |i|
                     if cur != @chars[0] and cur == @state[col - i][row + i] and (players == :all or players.include?(@chars.index(cur)))
                        count += 1
                     end
                  }
                  if count >= req
                     if req == @game.REQ
                        @the_win = []
                        (0...count).each {|i| @the_win[i] = [col - i, row + i]}
                        return cur
                     # AI: check up-left for available play
                     elsif col >= req and @state[col - req][row + req] == @chars[0] and @state[col - req][row + req - 1] != @chars[0]
                        return [col - req, row + req]
                     end
                  elsif req != @game.REQ and count == req - 1 and req + 1 >= @game.REQ
                     # intercept xx_x style patterns
                     if @state[col - req + 1][row + req - 1] == @chars[0] and @state[col - req][row + req] == cur and @next_row[col - req + 1] == row + req - 1
                        return [col - req + 1, row + req - 1]
                     elsif @state[col - 1][row + 1] == @chars[0] and @state[col - 2][row + 2] == cur and @next_row[col - 1] == row + 1
                        return [col - 1, row + 1]
                     end
                  end
               end  # up-left

               count = 1  # reset before testing other direction
               if (col - 1 + @game.REQ) < @state.size  # up-right match possible
                  (1...req).each { |i|
                     if cur != @chars[0] and cur == @state[col + i][row + i] and (players == :all or players.include?(@chars.index(cur)))
                        count += 1
                     end
                  }
                  if count >= req
                     if req == @game.REQ
                        @the_win = []
                        (0...count).each {|i| @the_win[i] = [col + i, row + i]}
                        return cur
                     # AI: check up-left for available play
                     elsif @state[col + req] and @state[col + req][row + req] == @chars[0] and @state[col + req][row + req - 1] != @chars[0]
                        return [col + req, row + req]
                     end
                  elsif req != @game.REQ and count == req - 1 and req + 1 >= @game.REQ
                     if @state[col + req - 1][row + req - 1] == @chars[0] and @state[col + req][row + req] == cur and @next_row[col + req - 1] == row + req - 1
                        return [col + req - 1, row + req - 1]
                     elsif @state[col + 1][row + 1] == @chars[0] and @state[col + 2][row + 2] == cur and @next_row[col + 1] == row + 1
                        return [col + 1, row + 1]
                     end
                  end
               end  # up-right
            end  # height test
         }  # end row
      }; false  # sosad
   end

   def tie?  # empty space = game on!
      (0...@state.size).each { |col|  # traverse top row
         return false if @state[col][@state.first.size - 1] == @chars[0]
      }; @chars[0]  # winner = blank (means tie)
   end

   private :vert_span?, :horiz_span?, :diag_span?, :tie?
end

class Game
   attr_reader :REQ, :board, :players, :screen

   def initialize(width, height, connect)
      @board = Board.new(self, width, height)
      @screen = Screen.new(self)
      @REQ = connect  # how many connected pieces to win?
      @players = []
   end

   def add_player(player)
      player.game = self
      @players << player
      player.number = @players.index(player) + 1
   end
end

class Screen
   attr_reader :width, :height, :num_colors

   # Initialize Ncurses and set up screen.
   def initialize(game)
      @game = game
      start_ncurses; get_size
      draw_logo(true)  # true = random color to start
   end

   def start_ncurses
      Ncurses::initscr
      Ncurses::noecho if @game.board.width < 10
      Ncurses::cbreak
      Ncurses::start_color; @num_colors = 0
      [ [0, Ncurses::COLOR_WHITE,   Ncurses::COLOR_BLACK],
                 [1, Ncurses::COLOR_RED,     Ncurses::COLOR_BLACK],
                 [2, Ncurses::COLOR_BLUE,    Ncurses::COLOR_BLACK],
                 [3, Ncurses::COLOR_GREEN,   Ncurses::COLOR_BLACK],
                 [4, Ncurses::COLOR_YELLOW,  Ncurses::COLOR_BLACK],
                 [5, Ncurses::COLOR_CYAN,    Ncurses::COLOR_BLACK],
                 [6, Ncurses::COLOR_MAGENTA, Ncurses::COLOR_BLACK],
      ].each {|color| @num_colors += 1; Ncurses::init_pair(*color)}
   end

   def get_size
      h, w = [], []
      Ncurses::getmaxyx(Ncurses::stdscr, h, w)
      @height, @width = h[0], w[0]  # Ncurses is stupid
      @lines = { :y_base  => (y_base = @height/2 + @game.board.height/2),
                 :y_info1 => y_base + 1,
                 :y_info2 => y_base + 2,
                 :y_input => y_base + 3,
                 :x_base  => (x_base = @width/2),
                 :x_input => x_base - 1,
      }
   end

   def resize
      get_size
      (0..@height).each {|line| Ncurses::mvprintw(line, 0, " "*@width)}  # clear screen
      draw_logo
      @game.board.draw
   end

   def draw_logo(random = false)
      logo = [
      ",---.,---.,---.,---.",
      "`---.|   |,---||   |",
      "`---'|---'`---^`   '",
      "     |"
      ]
      @logo_color = random ? (rand 2) + 1 : @logo_color || 1
      Ncurses::attrset(Ncurses::COLOR_PAIR(@logo_color))
      if @height > logo.size + @game.board.height + 5
         y, x = @height / 4 - logo.size, @width / 2 - logo[0].length / 2
         (0...logo.size).each {|i| Ncurses::mvprintw(y += 1, x, logo[i]) }
      elsif @height > @game.board.height + 3
         Ncurses::mvprintw(@height/4 - 1, @width/2 - 2, NAME)
      end
   end

   def input
      if @game.board.width > 9  # need multi-char input for > 9 columns
         got = ""; Ncurses::mvgetstr(@lines[:y_input], @lines[:x_input], got)
         Ncurses::mvprintw(@lines[:y_input], 0, " "*@width)  # clear old input
      else  # otherwise don't bother with hitting return
         got = Ncurses::getch
         if got > 47 and got < 123: got = got.chr     # ASCII is stupid
         elsif got == Ncurses::KEY_RESIZE: got = "R"  # so is Ncurses
         else got = "-1"
         end
      end
      exit if got == "q"; resize if got == "R"
      got
   end

   # Print a string to the general info area.
   def info(row, string, clear = true)
      Ncurses::attrset(Ncurses::A_NORMAL)
      (@lines[:y_info1]..@height).each {|line| Ncurses::mvprintw(line, 0, " "*@width)} if clear
      Ncurses::mvprintw(@lines[:y_base] + row, @lines[:x_base] - string.length / 2, string)
      Ncurses::refresh
   end

   def stop
      Ncurses::endwin; puts "Thanks for playing!"
   end

   private :start_ncurses, :draw_logo, :get_size
end

class Player
   attr_writer :game
   attr_accessor :number

   def play(column)
      @game.board.play(self, column)
   end

   def move
      play input
      rescue ArgumentError
         @game.screen.info(1, "Out of bounds!")
         sleep 0.5; retry
   end

   private :play
end

class HumanPlayer < Player
   def input
      @game.screen.info(1, "Which column? (1-#{@game.board.width})")
      col = @game.screen.input
      until col =~ /^\d{1,2}$/ and
            col.to_i > 0 and
            col.to_i <= @game.board.width
         @game.screen.info(1, "Invalid entry. Which column? (1-#{@game.board.width})")
         col = @game.screen.input
      end
      col.to_i - 1  # 0-index
   end
end

class DumbAIPlayer < Player
   def input
      @game.screen.info(1, "Thinking..."); sleep 0.5
      if rand(11) == 0
         @game.screen.info(1, "Now I've got you!"); sleep 0.5
      end
      r = rand(@game.board.width)
      until @game.board.next_row[r] < @game.board.height
         r = rand(@game.board.width)
      end; r
   end
end

class AIPlayer < Player
   def input
      @game.screen.info(1, "Thinking..."); sleep 0.5
      players = @game.players.collect {|x| x.number}

      # first, play your imminent wins
      if (span = @game.board.span?(@game.REQ - 1, [self.number]))
         if rand(7) == 0
            @game.screen.info(1, "All too easy."); sleep 0.7
         end
         return span[0]
      end

      # then stop imminent enemy wins
      if (span = @game.board.span?(@game.REQ - 1, players.delete_if {|x| x == self.number}))
         if rand(7) == 0
            @game.screen.info(1, "Surely you jest!"); sleep 0.7
         end
         return span[0]
      end

      # then just increase our longest spans
      if @game.REQ - 2 > 1  # don't check tiny spans
         (@game.REQ - 2).downto(2) { |i|
            if (span = @game.board.span?(i, [self.number]))
               return span[0]
            end
         }
      end

      # random if nothing else
      r = rand(@game.board.width)
      until @game.board.next_row[r] < @game.board.height
         r = rand(@game.board.width)
      end; r
   end
end

if __FILE__ == $0  # Enough setup, let the fun begin!
   begin
      trap('INT') { exit(-1) }

      human, ai, dumb = 1, 1, 0
      width, height, connect = 7, 6, 4
      opts = OptionParser.new  # parse command-line options
      opts.separator ""; opts.separator "Specific options:"
      opts.on("-u=NUM", "--human=NUM", "Number of human players.  (1)", Integer) { |num|
         human = num
      }
      opts.on("-a=NUM", "--ai=NUM", "Number of AI players.  (1)", Integer) { |num|
         ai = num
      }
      opts.on("-d=NUM", "--dumb=NUM", "Number of dumb, random AI players.  (0)", Integer) { |num|
         dumb = num
      }
      opts.on("-w=NUM", "--width=NUM", "Width of the board.  (7)", Integer) { |num|
         width = num
      }
      opts.on("-h=NUM", "--height=NUM", "Height of the board.  (6)", Integer) { |num|
         height = num
      }
      opts.on("-c=NUM", "--connect=NUM", "Number of connected pieces to win.  (4)", Integer) { |num|
         connect = num
      }
      opts.separator ""; opts.separator "General options:"
      opts.on_tail("-H", "--help", "--about", "Prints these usage instructions.") {
         puts opts; exit
      }
      opts.on_tail("-v", "-V", "--version", "Shows #{NAME} version.") {
         puts "#{NAME} version #{VER}"; exit
      }
      opts.parse!

      game = Game.new(width, height, connect)
      human.times { game.add_player(HumanPlayer.new) }
      ai.times { game.add_player(AIPlayer.new) }
      dumb.times { game.add_player(DumbAIPlayer.new) }

      begin
         game.board.clear
         active_player = 0
         until (winner = game.board.over?)  # main game loop
            game.board.draw
            game.players[active_player].move
            active_player = (active_player + 1) % game.players.size
         end

         game.board.draw  # show winning state
         if winner == game.board.chars[0]: game.screen.info(1, "The game is a draw.")
         else game.screen.info(1, "The #{winner}'s have it!")
         end

         game.screen.info(2, "Play again? (y/n)", false)
         response = game.screen.input
         until %w{y Y n N q}.include? response  # only take valid responses
            response = game.screen.input
         end
      end until %w{n N q}.include? response
   ensure
      game.screen.stop if defined?(game) and defined?(game.screen)
   end
end