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 cat_index(cat)
    for i, c in ipairs(cats) do
        if c == cat then
            return i
        end
    end
end

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 cat
        if is_bot then
            cat = self.bot:pick(sheet, dice)
            assert(sheet[cat] == nil and cat_index(cat) ~= nil)
            self.term:delay(1)
            for cur, curcat in ipairs(cats) do
                if sheet[curcat] == nil then
                    self:putsel(cur)
                    self.term:delay(1)
                    self:remsel(cur)
                    if curcat == cat then
                        break
                    end
                end
            end
            self.term:delay(2)
        else
            cat = cats[self:pick()]
        end
        local points = calc_points(dice, cat)
        sheet[cat] = points
        self:putpoints(cat_index(cat), 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(1)
                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" or key == " " 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

local function log(msg)
    --~ io.write("\x1B7")       -- save cursor position
    --~ io.write("\x1B[15;1H")  -- goto row 15, col 1
    --~ io.write("\x1B[2K")     -- clear line
    --~ io.write(msg)           -- print msg
    --~ io.write("\x1B8")       -- restore cursor position
end

function Bob:select(sheet, dice, roll)
    local sum, hist, min_count, max_count = calc_data(dice)
    local round = 1
    for _, cat in ipairs(cats) do
        if sheet[cat] ~= nil then
            round = round + 1
        end
    end
    local min_fh, min_st, min_yt, min_rp, min_4k
    if round <= 4 then
        min_fh, min_st, min_yt, min_rp, min_4k = 3, 4, 4, 3, 3
    elseif round <= 8 then
        min_fh, min_st, min_yt, min_rp, min_4k = 2, 3, 3, 2, 2
    else  -- round > 8
        min_fh, min_st, min_yt, min_rp, min_4k = 1, 3, 3, 1, 1
    end
    local sel = {}
    if sheet["full-house"] == nil then
        local rep = {}
        for i = min_fh, 6 do
            if hist[i] >= 2 then
                table.insert(rep, i)
            end
        end
        if #rep == 2 then
            for i = 1, 5 do
                sel[i] = hist[dice[i]] >= 2
            end
            log("full-house")
            return sel
        end
    end
    if sheet["little-straight"] == nil or sheet["big-straight"] == nil then
        local little_match, big_match
        little_match = (hist[1] > 0) and 1 or 0
        big_match = (hist[6] > 0) and 1 or 0
        for i = 2, 5 do
            if hist[i] > 0 then
                little_match = little_match + 1
                big_match = big_match + 1
            end
        end
        local cat
        if sheet["little-straight"] ~= nil then
            if big_match >= min_st then
                cat = "big-straight"
            end
        elseif sheet["big-straight"] ~= nil then
            if little_match >= min_st then
                cat = "little-straight"
            end
        elseif math.max(big_match, little_match) >= min_st then
            if big_match > little_match then
                cat = "big-straight"
            else
                cat = "little-straight"
            end
        end
        local tosel
        if cat == "little-straight" then
            tosel = {1, 1, 1, 1, 1, 0}
        elseif cat == "big-straight" then
            tosel = {0, 1, 1, 1, 1, 1}
        end
        if cat ~= nil then
            for i = 1, 5 do
                local d = dice[i]
                if tosel[d] > 0 then
                    sel[i] = true
                    tosel[d] = tosel[d] - 1
                end
            end
            log(cat)
            return sel
        end
    end
    if sheet["yacht"] == nil and max_count >= min_yt then
        for i = 1, 5 do
            sel[i] = hist[dice[i]] == max_count
        end
        log("yacht")
        return sel
    end
    if max_count >= min_rp or sheet["choice"] ~= nil or sum < 12 then
        local d
        for i = 1, 6 do
            if hist[i] == max_count then
                d = i
            end
        end
        local cat = cats[d+6]
        if sheet[cat] == nil or (sheet["four-of-a-kind"] == nil and d >= min_4k) then
            for i = 1, 5 do
                sel[i] = dice[i] == d
            end
            log("repeat "..d)
            return sel
        end
    end
    if sheet["choice"] == nil then
        for i = 1, 5 do
            sel[i] = dice[i] > 3
        end
        log("choice")
        return sel
    end
    log("discard roll")
    return {}
end

function Bob:pick(sheet, dice, roll)
    local sum, hist, min_count, max_count = calc_data(dice)
    local good_cats = {"yacht", "four-of-a-kind", "full-house", "little-straight", "big-straight"}
    local good_mins = {1, 16, 17, 1, 1}
    for i, cat in ipairs(good_cats) do
        local thresh = good_mins[i]
        if sheet[cat] == nil and calc_points(dice, cat) >= thresh then
            log("good match: "..cat)
            return cat
        end
    end
    if sheet["choice"] == nil and sum >= 20 then
        log("ok match: choice")
        return "choice"
    end
    for i = 1, 6 do
        local cat = i.."s"
        if sheet[cat] == nil and hist[i] >= 2 then
            log("ok match: "..cat)
            return cat
        end
    end
    -- no good pick, discard a sane item
    local bad_cats = {"1s", "2s", "3s", "little-straight", "big-straight"}
    for _, cat in ipairs(bad_cats) do
        if sheet[cat] == nil then
            log("discard bad")
            return cat
        end
    end
    -- discard anything with points (useful on last few rounds)
    for _, cat in ipairs(bad_cats) do
        if sheet[cat] == nil  and calc_points(dice, cat) > 0 then
            log("discard with points")
            return cat
        end
    end
    -- no sane discard possible, discard anything
    log("insane discard")
    local cat
    repeat
        cat = cats[self.prng:randint(1, 12)]
    until sheet[cat] == nil
    return cat
end

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