-- Traitor credits
CreateConVar("ttt_max_startweight", "7", FCVAR_NOTIFY)

local GetTraitorCount = nil
local GetDetectiveCount = nil
local GetAvoidDetective = nil


function Initialize()
   print("I RAN INITIALIZE")

   for i = 1, math.huge do

      local k, v = debug.getupvalue( SelectRoles, i )

      if k == "GetTraitorCount" then
         GetTraitorCount = v
      end
      if k == "GetDetectiveCount" then
         GetDetectiveCount = v
      end
      if k == "GetAvoidDetective" then
         GetAvoidDetective = v
      end

      if GetTraitorCount ~= nil and GetDetectiveCount ~= nil and GetAvoidDetective ~= nil or k == nil then
         break;
      end

   end
end
hook.Add( "Initialize", "InitializeExternal", Initialize)


function SetDefaultWeight()
   -- Set players weight to 10 by default
   for k,v in pairs(player.GetAll()) do
      if IsValid(v) and (not v:IsSpec()) and (v.Weight == nil) then
         v.Weight = StartWeight()
      end
   end
end

-- Adds all players weight together.
function TotalPlayerWeight()
   local totalWeight = 0
   for k,v in pairs(player.GetAll()) do
      if IsValid(v) and (not v:IsSpec()) and (v.Weight ~= nil) then
         totalWeight = totalWeight + v.Weight
      end
   end
   return totalWeight
end

-- Chooses a starting weight between 1 and max number
function StartWeight()
   return math.random(1, GetConVar("ttt_round_limit"):GetInt())
end

math.randomseed( os.time() )
local function shuffleTable( t )
    local rand = math.random
    local iterations = #t
    local j

    for i = iterations, 2, -1 do
        j = rand(i)
        t[i], t[j] = t[j], t[i]
    end
end

function SelectRoles()
   local choices = {}
   local prev_roles = {
      [ROLE_INNOCENT] = {},
      [ROLE_TRAITOR] = {},
      [ROLE_DETECTIVE] = {}
   };

   SetDefaultWeight()

   if not GAMEMODE.LastRole then GAMEMODE.LastRole = {} end

   for k,v in pairs(player.GetAll()) do
      -- everyone on the spec team is in specmode
      if IsValid(v) and (not v:IsSpec()) then
         -- save previous role and sign up as possible traitor/detective

         local r = GAMEMODE.LastRole[v:UniqueID()] or v:GetRole() or ROLE_INNOCENT

         table.insert(prev_roles[r], v)

         table.insert(choices, v)
      end
      v:SetRole(ROLE_INNOCENT)
   end

   -- determine how many of each role we want
   local choice_count = #choices
   local traitor_count = GetTraitorCount(choice_count)
   local det_count = GetDetectiveCount(choice_count)
   local total_weight = TotalPlayerWeight()

   if choice_count == 0 then return end

   print("Traitor Count: " .. traitor_count)

   -- first select traitors
   local ts = 0
   while ts < traitor_count do

      -- Shuffle all the choices before selecting first traitor
      -- Makes it more random.
      shuffleTable(choices)

      -- select random number between 1 and the total weight.
      local r = math.random(1, total_weight)

      -- For all the players who are an available choice.
      for k,v in pairs(choices) do

         -- If the random number selected between 1 and total weight minus the users weight
         -- is less then or equal to 0 the person will be selected as traitor
         if (r - v.Weight) <= 0 then
            -- make this guy traitor if he was not a traitor last time, or if he makes
            -- a roll
            if IsValid(v) and
               ((not table.HasValue(prev_roles[ROLE_TRAITOR], v)) or (math.random(1, 3) == 2)) then
               v:SetRole(ROLE_TRAITOR)

               -- Set the players weight who just turned to traitor back to a starting weight.
               v.Weight = StartWeight()

               -- Remove the player from being a choice for a traitor if more then 1 is available.
               table.remove(choices, choices[v])
               ts = ts + 1
               break
            end
         end
         -- Subtract weight from random number to keep number getting closer to 0
         r = r - v.Weight
      end

      -- Recalculate weight after selecting a traitor.
      -- Should no longer count the person that isn't involved in the selections.
      total_weight = 0
      for k,v in pairs(choices) do
         total_weight = total_weight + v.Weight
      end
   end

   -- now select detectives, explicitly choosing from players who did not get
   -- traitor, so becoming detective does not mean you lost a chance to be
   -- traitor
   local ds = 0
   local min_karma = GetConVarNumber("ttt_detective_karma_min") or 0
   while (ds < det_count) and (#choices >= 1) do

      -- sometimes we need all remaining choices to be detective to fill the
      -- roles up, this happens more often with a lot of detective-deniers
      if #choices <= (det_count - ds) then
         for k, pply in pairs(choices) do
            if IsValid(pply) then
               pply:SetRole(ROLE_DETECTIVE)
            end
         end

         break -- out of while
      end


      local pick = math.random(1, #choices)
      local pply = choices[pick]

      -- we are less likely to be a detective unless we were innocent last round
      if (IsValid(pply) and
          ((pply:GetBaseKarma() > min_karma and
           table.HasValue(prev_roles[ROLE_INNOCENT], pply)) or
           math.random(1,3) == 2)) then

         -- if a player has specified he does not want to be detective, we skip
         -- him here (he might still get it if we don't have enough
         -- alternatives)
         if not pply:GetAvoidDetective() then
            pply:SetRole(ROLE_DETECTIVE)
            ds = ds + 1
         end

         table.remove(choices, pick)
      end
   end

   -- Debug,
   for k,v in pairs(player.GetAll()) do
      if IsValid(v) and v:GetRole() == ROLE_INNOCENT then
         v.Weight = v.Weight + math.random(5, 15)
      end
   end

   GAMEMODE.LastRole = {}

   for _, ply in pairs(player.GetAll()) do
      -- initialize credit count for everyone based on their role
      ply:SetDefaultCredits()

      -- store a uid -> role map
      GAMEMODE.LastRole[ply:UniqueID()] = ply:GetRole()
   end
end
hook.Add( "PostGamemodeLoaded", "OverrideDefaultSelectRolesFunction", SelectRoles)


function PlayerWeight(ply)
      if (not IsValid(ply)) or ply:IsAdmin() or ply:IsSuperAdmin() or cvars.Bool("sv_cheats", 0) then
         PrintMessage( HUD_PRINTTALK, ply:GetName() .. " has accessed traitor weight logs.")

         ply:PrintMessage(HUD_PRINTCONSOLE, "  *** Weight log:")
         for k, v in pairs(player.GetAll()) do
            if IsValid(v) and v.Weight ~= nil then
               ply:PrintMessage(HUD_PRINTCONSOLE, v:GetName() .. " = " .. v.Weight)
            end
         end
         ply:PrintMessage(HUD_PRINTCONSOLE, "  *** Weight log end.")


      else
         ply:PrintMessage(HUD_PRINTCONSOLE, "You must be a GMod Admin or SuperAdmin on the server to use this command, or sv_cheats must be enabled.")
      end
end
concommand.Add("ttt_print_playerweight", PlayerWeight)