local arco = require "arco"

local cats = {
    "full-house", "four-of-a-kind", "little-straight", "big-straight",
    "choice", "yacht", "1s", "2s", "3s", "4s", "5s", "6s"
}

local function calc_total(sheet)
    local total = 0
    for _, cat in ipairs(cats) do
        total = total + (sheet[cat] or 0)
    end
    return total
end

local function calc_data(dice)
    local sum = 0
    local hist = {0, 0, 0, 0, 0, 0}
    for _, die in ipairs(dice) do
        sum = sum + die
        hist[die] = hist[die] + 1
    end
    local min_count, max_count = math.huge, -math.huge
    for i = 1, 6 do
        local count = hist[i]
        if count > 0 then
            if count < min_count then
                min_count = count
            end
            if count > max_count then
                max_count = count
            end
        end
    end
    return sum, hist, min_count, max_count
end

local function calc_points(dice, cat)
    local sum, hist, min_count, max_count = calc_data(dice)
    local points = 0
    if cat == "full-house" then
        if min_count == 2 and max_count == 3 then
            points = sum
        end
    elseif cat == "four-of-a-kind" then
        if max_count >= 4 then
            for i = 1, 6 do
                if hist[i] >= 4 then
                    points = 4 * i
                    break
                end
            end
        end
    elseif cat == "little-straight" then
        if max_count == 1 and hist[6] == 0 then
            points = 30
        end
    elseif cat == "big-straight" then
        if max_count == 1 and hist[1] == 0 then
            points = 30
        end
    elseif cat == "choice" then
        points = sum
    elseif cat == "yacht" then
        if max_count == 5 then
            points = 50
        end
    else
        local num = tonumber(cat:sub(1, 1))
        points = num * hist[num]
    end
    return points
end

local Yacht = {}
Yacht.__index = Yacht

function Yacht:init()
    self.term:setup()
    self.term:enter_altbuf()
    self.term:hide_cursor()
end

function Yacht:quit()
    self.term:show_cursor()
    self.term:exit_altbuf()
    self.term:restore()
end

function Yacht:setup_game()
    self.term:clear()
    self.nh = tonumber(self.term:input("number of human players: "))
    self.nc = tonumber(self.term:input("number of computer players: "))
    self.np = self.nh + self.nc
    self.players = {}
    for i = 1, self.nh do
        table.insert(self.players, "H-"..i)
    end
    for i = 1, self.nc do
        table.insert(self.players, "C-"..i)
    end
    self.prng:shuffle(self.players)
    self.sheets = {}
    for i = 1, self.np do
        table.insert(self.sheets, {})
    end
end

function Yacht:run_game()
    for i = 1, 12 do
        self:run_round()
    end
    self:show_score()
end

function Yacht:run_round()
    self.term:clear()
    for i, p in ipairs(self.players) do
        self.pi = i
        self:show_sheet()
        local sheet = self.sheets[i]
        local is_bot = p:sub(1, 1) == "C"
        local dice = self:do_rolls(is_bot)
        local sel
        if is_bot then
            sel = self.bot:pick(sheet, dice)
            self.term:delay(1)
            for cur, cat in ipairs(cats) do
                if sheet[cat] == nil then
                    self:putsel(cur)
                    self.term:delay(2)
                    self:remsel(cur)
                    if cur == sel then
                        break
                    end
                end
            end
            self.term:delay(2)
        else
            sel = self:pick()
        end
        local cat = cats[sel]
        local points = calc_points(dice, cat)
        sheet[cat] = points
        self:putpoints(sel, points)
        self.term:goto(2+7*(i-1), 2)
        io.write(("%3d"):format(calc_total(sheet)))
        self.term:get_key()
        self.term:goto(1, 1)
    end
end

function Yacht:show_sheet()
    local sheet = self.sheets[self.pi]
    local fmt
    for i, p in ipairs(self.players) do
        if i == self.pi then
            fmt = "[%s]  "
        else
            fmt = " %s   "
        end
        io.write(fmt:format(p))
    end
    self.term:goto(1, 2)
    for i, p in ipairs(self.players) do
        if i == self.pi then
            fmt = "[%3d]  "
        else
            fmt = " %3d   "
        end
        io.write(fmt:format(calc_total(self.sheets[i])))
    end
    self.term:goto(1, 8)
    for i = 1, #cats/2 do
        local cat_l = cats[i]
        local cat_r = cats[i+6]
        local val_l = sheet[cat_l] or "--"
        if val_l == 0 then
            val_l = "XX"
        end
        local val_r = sheet[cat_r] or "--"
        if val_r == 0 then
            val_r = "XX"
        end
        local fmt = "%16s: %2s"..(" "):rep(4).."%8s: %2s"
        print(fmt:format(cat_l, val_l, cat_r, val_r))
    end
    self.term:goto(1, 4)
end

function Yacht:show_score()
    self.term:clear()
    self.term:down(3)
    print("== final score ==\n")
    for i, p in ipairs(self.players) do
        print(("%s: %d"):format(p, calc_total(self.sheets[i])))
    end
    print("\n=================")
    self.term:get_key()
end

function Yacht:do_rolls(is_bot)
    local pad = (" "):rep(15)
    print(pad.."roll: 1")
    self.term:clear_line()
    self.term.down()
    local dice, sel = {}, {}
    io.write((" "):rep(12).."[ - - - - - ]")
    self:roll(dice)
    for i = 2, 3 do
        if is_bot then
            local sheet = self.sheets[self.pi]
            sel = self.bot:select(sheet, dice, i-1)
            self.term:delay(3)
            self.term:goto_col(15)
            for cur = 1, 5 do
                io.write("^")
                self.term:left()
                self.term:delay(2)
                self.term:up(2)
                io.write(sel[cur] and "= " or "  ")
                self.term:delay(1)
                self.term:down(2)
                self.term:clear_line()
            end
            self.term:delay(3)
            self.term:up(3)
            self.term:goto_col(1)
        else
            self:select(sel)
        end
        for j = 1, 5 do
            if not sel[j] then
                dice[j] = nil
            end
        end
        print(pad.."roll: "..i.."\n")
        self:roll(dice)
    end
    self.term:up(2)
    self.term:clear_line()
    return dice
end

function Yacht:roll(dice)
    self.term:goto_col(15)
    for i = 1, 5 do
        if dice[i] == nil then
            io.write("- ")
        else
            self.term:right(2)
        end
    end
    self.term:goto_col(15)
    for i = 1, 5 do
        if dice[i] == nil then
            dice[i] = self.prng:randint(1, 6)
            self.term:delay(1)
            io.write(("%d "):format(dice[i]))
        else
            self.term:right(2)
        end
    end
    self.term:down()
end

function Yacht:select(sel)
    local cur = 1
    for i = 1, 5 do
        if not sel[i] then
            cur = i
            break
        end
    end
    self.term:goto_col(cur * 2 + 13)
    io.write("^")
    while true do
        local key = self.term:get_key()
        local update = true
        if key == "\n" then
            break
        elseif key == "a" or key == "h" then
            cur = (cur > 1) and (cur - 1) or 5
        elseif key == "d" or key == "l" then
            cur = (cur < 5) and (cur + 1) or 1
        elseif key == " " then
            sel[cur] = not sel[cur]
        else
            update = false
        end
        if update then
            self.term:up(2)
            self.term:clear_line()
            self.term:goto_col(15)
            for i = 1, 5 do
                io.write(sel[i] and "= " or "  ")
            end
            self.term:down(2)
            self.term:clear_line()
            self.term:goto_col(cur * 2 + 13)
            io.write("^")
        end
    end
    self.term:clear_line()
    self.term:up(3)
    self.term:goto_col(1)
end

function Yacht:getpos(s)
    if s <= #cats/2 then
        return 18, s+7
    else
        return 34, s+1
    end
end

function Yacht:putsel(s)
    self.term:goto(self:getpos(s))
    io.write("[")
    self.term:right(2)
    io.write("]")
end

function Yacht:remsel(s)
    self.term:goto(self:getpos(s))
    io.write(" -- ")
end

function Yacht:putpoints(s, p)
    self.term:goto(self:getpos(s))
    if p == 0 then
        io.write(" XX ")
    else
        io.write((" %2d "):format(p))
    end
end

function Yacht:pick()
    local sheet = self.sheets[self.pi]
    local sel
    for i, cat in ipairs(cats) do
        if sheet[cat] == nil then
            sel = i
            break
        end
    end
    self:putsel(sel)
    while true do
        local key = self.term:get_key()
        local update = true
        local newsel = sel
        if key == "\n" then
            break
        elseif key == "w" or key == "k" then
            repeat
                newsel = (newsel > 1) and newsel - 1 or #cats
            until sheet[cats[newsel]] == nil
        elseif key == "s" or key == "j" then
            repeat
                newsel = (newsel < #cats) and newsel + 1 or 1
            until sheet[cats[newsel]] == nil
        elseif key == "a" or key == "d" or key == "h" or key == "l" then
            newsel = (sel <= #cats/2) and sel + 6 or sel - 6
            if sheet[cats[newsel]] ~= nil then
                update = false
            end
        else
            update = false
        end
        if update then
            self:remsel(sel)
            sel = newsel
            self:putsel(sel)
        end
    end
    return sel
end

local function new_yacht(Bot)
    local term = arco.tui.new_term()
    local prng = arco.rand.new_prng()
    local bot = setmetatable({prng=prng}, Bot)
    return setmetatable({term=term, prng=prng, bot=bot}, Yacht)
end

local Bob = {}
Bob.__index = Bob

function Bob:select(sheet, dice, roll)
    local sum, hist, min_count, max_count = calc_data(dice)
    local sel = {}
    if max_count > 1 then
        for i = 1, 5 do
            sel[i] = hist[dice[i]] == max_count
        end
    end
    return sel
end

function Bob:pick(sheet, dice, roll)
    local min_points, max_points = 99, 0
    local sel, sel_min, sel_max
    for cur, cat in ipairs(cats) do
        if sheet[cat] == nil then
            local points = calc_points(dice, cat)
            if points > max_points then
                sel_max = cur
                max_points = points
            end
            if points > 0 and points < min_points then
                sel_min = cur
                min_points = points
            end
        end
    end
    if max_points > 20 or (sel_max ~= nil and sel_max > 6) then
        sel = sel_max
    else
        sel = sel_min
    end
    if sel == nil then
        repeat
            sel = self.prng:randint(1, #cats)
        until sheet[cats[sel]] == nil
    end
    return sel
end

local yacht = new_yacht(Bob)
yacht:init()
yacht:setup_game()
yacht:run_game()
yacht:quit()