#!/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.2.0"

require 'ncurses'
require 'optparse'
require 'socket'

ERR = {:input => "err_in", :quit => "err_qu"}

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) }
      Ncurses::reset_prog_mode; Ncurses::refresh
   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
      }
      (1..@width).each {|i| Ncurses::printw("#{i}#{i<10?' ':''}")}  # col numbers
      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 + 1] || "?"
         @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
   # 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) - 1))
               count += 1
            else count = 1
            end

            last_char = row
            if count >= req  # a winner is you
               if req == @game.REQ  # this was an over? check
                  @the_win = []  # array of winning pieces - [x, y]
                  (0...count).each {|i| @the_win[i] = [c, r - i]}
                  return row
               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) - 1))
               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) - 1))
                        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) - 1))
                        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 Board

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)
   end

   def renumber(start)  # align players with server, start = where locals begin
      server_players = @players.find_all {|p| p.class == ServerPlayer}
      (0...start).each { |i|
         last_p = @players[i].dup
         @players[i] = server_players[i]
         ((i + 1)..@players[i].number).each { |j|
            new_p = @players[j].dup
            @players[j] = last_p
            last_p = new_p
         }
      }
      @players.each_with_index {|p, i| p.number = i}
   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
      Ncurses::def_prog_mode; Ncurses::endwin  # pause Ncurses til we need it
   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 + 1),
                 :y_info1 => y_base + 1,
                 :y_info2 => y_base + 2,
                 :y_input => y_base + 3,
                 :x_base  => (x_base = @width/2 - 1),
                 :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 - 1
         (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 - NAME.length/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 Screen

class Player
   attr_writer :game
   attr_accessor :number

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

   def move
      play (col = input)
      [@number, col]
   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, player #{@number + 1}?")
      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, [@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 == @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, [@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

class NetPlayer < Player
   attr_writer :server
   attr_accessor :client

   def input
      @game.screen.info(1, "Waiting for player #{@number+1}...")
      ret = @server.receive(@number)
      until ret =~ /^\d{1,2} \d{1,2}$/ and ret.split[0].to_i == @number
         @server.send(ERR[:input], @number)
         ret = @server.receive(@number)
      end
      ret.split[1].to_i - 1  # 0-index
   end
end

class ServerPlayer < Player
   attr_writer :server

   def input
      @game.screen.info(1, "Waiting for player #{@number+1}...")
      ret = @server.receive
      if ret.split[0].to_i == @number: return ret.split[1].to_i
      else at_exit {puts "Server has unexpectedly quit!"}; exit(-1)
      end
   end
end

class GameServer
   def initialize(game, host = nil, port = 2202, clients = 1)
      @game, @host = game, host
      port ||= 2202
      if host: @server = TCPSocket.new(host, port)  # we're a client
      else  # we're a server
         @server = TCPServer.new('localhost', port)
         @clients = []  # array of TCPSockets to players
         STDOUT.sync = true  # Ncurses is stupid
         net_start = @game.players.find {|p| p.class == NetPlayer}.number - 1
         (0...clients).each { |i|
            puts "Waiting for player #{@clients.compact.size + 1}..."
            @clients[net_start += 1] = @server.accept
            size = receive(net_start).to_i
            (1...size).each { |i|  # adjust @clients for multiples
               @game.add_player NetPlayer.new
               net_start += 1
               @clients[net_start] = @clients[net_start - 1]
            }
         }
         sent = []
         @clients.each_with_index { |c, i|  # tell remotes their position
            @game.players[i].client = (c == @clients[i-1] ? @game.players[i-1].client : i) if @game.players[i].class == NetPlayer
            unless c == nil or sent.include? c
               sent << c
               send("#{i} of #{@game.players.size - 1}", i)
            end
         }
      end
   end

   def send(message, client = nil)
      if client: @clients[client].send(message, 0)
      else @server.send(message, 0)
      end
   rescue Errno::EPIPE  # player quit
      at_exit {puts "A player has unexpectedly quit!"}; exit(-1)
   end

   def receive(client = nil)
      if client: @clients[client].recv(8).chomp
      else @server.recv(8).chomp
      end
   end

   def update(move)
      unless @host  # server: send recent plays
         @game.players.find_all {|p| p.class == NetPlayer}.collect {|p| p.client}.uniq.each { |c|
            unless @game.players[move[0]].class == NetPlayer and c == @game.players[move[0]].client
               send(move.join(" "), c)
            end
         }
      else  # client: send your plays to server
         local_players = @game.players.reject {|p| p.class == ServerPlayer}.collect {|p| p.number}
         if local_players.include? move[0]
            send "#{move[0]} #{move[1] + 1}"
         end
      end
   end
end  # class GameServer

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

      width, height, connect = 7, 6, 4
      human, ai, dumb, clients = -1, -1, -1, 1

      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 "Network options:"
      host = net_game = port = nil
      opts.on("-s [HOST]", "--server [HOST]", "Create a server, or connect to one at HOST.", String) { |h|
         host = h; net_game = true
         ai = 0 if ai < 0
      }
      opts.on("-p=PORT", "--port=PORT", "Port to serve on/connect to.  (2202)", Integer) { |p|
         port = p
      }
      opts.on("-C=NUM", "--clients=NUM", "Number of clients, if serving.  (1)", Integer) { |c|
         clients = c
      }

      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 = 1 if human < 0; ai = 1 if ai < 0; dumb = 0 if dumb < 0
      human.times { game.add_player HumanPlayer.new }
      ai.times    { game.add_player AIPlayer.new }
      dumb.times  { game.add_player DumbAIPlayer.new }
      if net_game
         clients.times { game.add_player NetPlayer.new } unless host
         server = GameServer.new(game, host, port, clients)
         if host  # align server player numbers to local, if client
            server.send((local = human + ai + dumb).to_s)
            server_players = server.receive.split(' of ')
            server_start, total = server_players.collect {|s| s.to_i}
            (total - local + 1).times { game.add_player ServerPlayer.new }
            game.renumber(server_start)
         end
         game.players.each {|p| p.server = server if p.class == NetPlayer or p.class == ServerPlayer}
      end

      begin
         game.board.clear
         active_player = 0
         until (winner = game.board.over?)  # main game loop
            game.board.draw
            move = game.players[active_player].move
            server.update(move) if net_game
            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  # main loop