local M = {}

local i_o = require 'i_o'
local pure = require 'pure'

local __string_match = string.match
local __string_gmatch = string.gmatch
local __string_format = string.format
local __tonumber = tonumber

local dirname = function(s)
   return __string_match(s, '(.*)/name')
end

local read_micro = function(path)
   return i_o.read_file(path, nil, '*n') * 0.000001
end

local gmatch_to_table1 = function(pat, s)
   return pure.iter_to_table1(__string_gmatch(s, pat))
end

local gmatch_to_tableN = function(pat, s)
   return pure.iter_to_tableN(__string_gmatch(s, pat))
end

--------------------------------------------------------------------------------
-- memory

local MEMINFO_PATH = '/proc/meminfo'

local fmt_mem_field = function(field)
   return field..':%s+(%d+)'
end

local meminfo_regex = function(read_swap)
   -- ASSUME the order of the meminfo file will never change, but some options
   -- (like swap) might not exist
   local free_fields = {
         'MemFree',
         'Buffers',
         'Cached'
   }
   local swap_field = 'SwapFree'
   local slab_fields = {
      'Shmem',
      'SReclaimable'
   }
   local all_fields = read_swap == true
      and {free_fields, {swap_field}, slab_fields}
      or {free_fields, slab_fields}
   local patterns = pure.map(fmt_mem_field, pure.flatten(all_fields))
   return table.concat(patterns, '.+\n')
end

M.meminfo_updater_swap = function(mem_state, swap_state)
   local regex = meminfo_regex(true)
   return function()
      mem_state.memfree,
         mem_state.buffers,
         mem_state.cached,
         swap_state.free,
         mem_state.shmem,
         mem_state.sreclaimable
         = __string_match(i_o.read_file(MEMINFO_PATH), regex)
   end
end

M.meminfo_updater_noswap = function(mem_state)
   local regex = meminfo_regex(false)
   return function()
      mem_state.memfree,
         mem_state.buffers,
         mem_state.cached,
         mem_state.shmem,
         mem_state.sreclaimable
         = __string_match(i_o.read_file(MEMINFO_PATH), regex)
   end
end

M.meminfo_field_reader = function(field)
   local pattern = fmt_mem_field(field)
   return function()
      return __tonumber(i_o.read_file(MEMINFO_PATH, pattern))
   end
end

--------------------------------------------------------------------------------
-- intel powercap

local SYSFS_RAPL = '/sys/class/powercap'

M.intel_powercap_reader = function(dev)
   local uj = __string_format('%s/%s/energy_uj', SYSFS_RAPL, dev)
   i_o.assert_file_readable(uj)
   return function() return read_micro(uj) end
end

--------------------------------------------------------------------------------
-- battery

local SYSFS_POWER = '/sys/class/power_supply'

local format_power_path = function(battery, property)
   local p = __string_format('%s/%s/%s', SYSFS_POWER, battery, property)
   i_o.assert_file_readable(p)
   return p
end

M.battery_power_reader = function(battery)
   local current = format_power_path(battery, 'current_now')
   local voltage = format_power_path(battery, 'voltage_now')
   return function() return read_micro(current) * read_micro(voltage) end
end

M.battery_status_reader = function(battery)
   local status = format_power_path(battery, 'status')
   return function()
      return i_o.read_file(status, nil, '*l') ~= 'Discharging'
   end
end

--------------------------------------------------------------------------------
-- disk io

M.get_disk_paths = pure.partial(
   pure.map,
   pure.partial(__string_format, '/sys/block/%s/stat', true)
)

-- fields 3 and 7 (sectors read and written)
local RW_REGEX = '%s+%d+%s+%d+%s+(%d+)%s+%d+%s+%d+%s+%d+%s+(%d+)'

-- the sector size of any block device in linux is 512 bytes
-- see https://git.kernel.org/pub/scm/linux/kernel/git/torvalds/linux.git/tree/include/linux/types.h?id=v4.4-rc6#n121
local BLOCK_SIZE_BYTES = 512

M.get_total_disk_io = function(paths)
   local r = 0
   local w = 0
   for i = 1, #paths do
      local _r, _w = __string_match(i_o.read_file(paths[i]), RW_REGEX)
      r = r + __tonumber(_r)
      w = w + __tonumber(_w)
   end
   return r * BLOCK_SIZE_BYTES, w * BLOCK_SIZE_BYTES
end

--------------------------------------------------------------------------------
-- network

-- ASSUME realpath exists (part of coreutils)

local NET_DIR = '/sys/class/net'

local get_interfaces = function()
   local cmd = __string_format('realpath %s/* | grep -v virtual', NET_DIR)
   local f = pure.partial(gmatch_to_table1, '/([^/\n]+)\n')
   return pure.maybe({}, f, i_o.execute_cmd(cmd))
end

M.get_net_interface_paths = function()
   return pure.map(
      function(s)
         local dir = __string_format('%s/%s/statistics/', NET_DIR, s)
         return {rx = dir..'rx_bytes', tx = dir..'tx_bytes'}
      end,
      get_interfaces()
   )
end

--------------------------------------------------------------------------------
-- cpu

M.get_cpu_number = function(topology)
   local n = 0
   for g, c in pairs(topology) do
      n = n + g * #c
   end
   return n
end

local get_coretemp_dir = function()
   i_o.assert_exe_exists('grep')
   local s = i_o.execute_cmd('grep -l \'^coretemp$\' /sys/class/hwmon/*/name')
   return pure.fmap_maybe(dirname, s)
end

-- return a table with keys corresponding to physcial core id and values to
-- the number of threads of each core (usually 1 or 2)
M.get_core_threads = function()
   i_o.assert_exe_exists('lscpu')
   local cmd = 'lscpu -y -p=core | grep -v \'^#\' | sort -k1,1n | uniq -c'
   local flip = function(c) return {__tonumber(c[2]), __tonumber(c[1])} end
   local make_indexer = pure.compose(
      pure.array_to_map,
      pure.partial(pure.map, flip),
      pure.partial(gmatch_to_tableN, '(%d+) (%d+)')
   )
   return pure.fmap_maybe(make_indexer, i_o.execute_cmd(cmd))
end

local get_coretemp_mapper = function()
   i_o.assert_exe_exists('grep')
   local d = get_coretemp_dir()
   local get_labels = pure.compose(
      i_o.execute_cmd,
      pure.partial(__string_format, 'grep Core %s/temp*_label', true)
   )
   local to_tuple = function(m)
      return {__tonumber(m[2]), __string_format('%s/%s_input', d, m[1])}
   end
   local to_map = pure.compose(
      pure.array_to_map,
      pure.partial(pure.map, to_tuple),
      pure.partial(gmatch_to_tableN, '/([^/\n]+)_label:Core (%d+)\n')
   )
   return pure.maybe({}, to_map, pure.fmap_maybe(get_labels, d))
end

M.get_core_topology = function()
   i_o.assert_exe_exists('lscpu')
   i_o.assert_exe_exists('grep')
   i_o.assert_exe_exists('sort')
   local coretemp_paths = get_coretemp_mapper()
   local assign_cpu = function(i, x)
      return {
         lgl_cpu_id = i,
         phy_core_id = __tonumber(x[1]),
         phy_cpu_id = __tonumber(x[2])
      }
   end
   local assign_core = function(acc, next)
      local g = acc.grouped
      local max_lgl_core_id = #g
      local new_phy_core_id = next.phy_core_id
      local new_cpu = {phy_cpu_id = next.phy_cpu_id, lgl_cpu_id = next.lgl_cpu_id}
      if acc.prev_phy_core_id == new_phy_core_id then
         local max_thread = #acc.grouped[max_lgl_core_id].cpus
         acc.grouped[max_lgl_core_id].cpus[max_thread + 1] = new_cpu
      else
         local new_lgl_core_id = max_lgl_core_id + 1
         acc.grouped[new_lgl_core_id] = {
            phy_core_id = new_phy_core_id,
            lgl_core_id = new_lgl_core_id,
            coretemp_path = coretemp_paths[new_phy_core_id],
            cpus = {new_cpu}
         }
         acc.prev_phy_core_id = new_phy_core_id
      end
      return acc
   end
   local get_threads = function(x)
      return #x.cpus
   end
   local f = pure.compose(
      pure.partial(pure.group_with, get_threads, pure.id),
      pure.partial(pure.get, 'grouped'),
      pure.partial(pure.reduce, assign_core, {prev_phy_core_id = -1, grouped = {}}),
      pure.partial(pure.imap, assign_cpu),
      pure.partial(gmatch_to_tableN, '(%d+),(%d+)')
   )
   local out =
      i_o.execute_cmd('lscpu -y -p=core,cpu | grep -v \'^#\' | sort -k1,1n')
   return pure.fmap_maybe(f, out)
end

M.topology_to_cpu_map = function(topology)
   local r = {}
   for group_id, group in pairs(topology) do
      for _, core in pairs(group) do
         for _, cpu in pairs(core.cpus) do
            r[cpu.lgl_cpu_id] = group_id
         end
      end
   end
   return r
end

M.read_ave_freqs = function(topology, cpu_group_map)
   -- NOTE: Using the builtin conky functions for getting cpu freq seems to make
   -- the entire loop jittery due to high variance latency. Querying
   -- scaling_cur_freq in sysfs seems to do the same thing. It appears lscpu
   -- (which queries /proc/cpuinfo) is much faster and doesn't have this jittery
   -- problem.
   i_o.assert_exe_exists('lscpu')
   local out = i_o.execute_cmd('lscpu -p=MHZ')
   local init_freqs = function(v)
      local r = {}
      for group_id, _ in pairs(topology) do
         r[group_id] = v
      end
      return r
   end
   if out == nil then
      return init_freqs('N/A')
   else
      local ave_freqs = init_freqs(0)
      local cpu_id = 1
      for s in __string_gmatch(out, '(%d+%.%d+)') do
         local group_id = cpu_group_map[cpu_id]
         ave_freqs[group_id] = ave_freqs[group_id] + __tonumber(s)
         cpu_id = cpu_id + 1
      end
      for group_id, _ in pairs(ave_freqs) do
         ave_freqs[group_id] =
            __string_format(
               '%.0f Mhz',
               ave_freqs[group_id] / (group_id * #topology[group_id])
            )
      end
      return ave_freqs
   end
end

M.get_hwp_paths = function(topology)
   -- ASSUME this will never fail
   return pure.map_n(
      function(i)
         return '/sys/devices/system/cpu/cpu'
            .. (i - 1)
            .. '/cpufreq/energy_performance_preference'
      end,
      M.get_cpu_number(topology)
   )
end

local HWP_MAP = {
   power = 'Power',
   balance_power = 'Bal. Power',
   balance_performance = 'Bal. Performance',
   performance = 'Performance',
   default = 'Default',
}

local read_hwp_path = function(path)
   return i_o.read_file(path, nil, "*l")
end

M.read_hwp = function(hwp_paths)
   -- read HWP of first cpu, then test all others to see if they match
   local hwp_pref = read_hwp_path(hwp_paths[1])
   local mixed = false
   local i = 2
   while not mixed and i <= #hwp_paths do
      mixed = hwp_pref ~= read_hwp_path(hwp_paths[i])
      i = i + 1
   end
   return mixed and 'Mixed' or (HWP_MAP[hwp_pref] or 'Unknown')
end

M.init_cpu_loads = function(topo)
   local ncpus = M.get_cpu_number(topo)
   local cpu_loads = {}
   for lgl_cpu_id = 1, ncpus do
      cpu_loads[lgl_cpu_id] = {
         active_prev = 0,
         total_prev = 0,
         percent_active = 0,
      }
   end
   return cpu_loads
end

M.read_cpu_loads = function(cpu_loads)
   local iter = io.lines('/proc/stat')
   iter() -- ignore first line
   for lgl_cpu_id = 1, #cpu_loads do
      local ln = iter()
      local user, system, idle =
         __string_match(ln, '%d+ (%d+) %d+ (%d+) (%d+)', 4)
      local active = user + system
      local total = active + idle
      local cpu = cpu_loads[lgl_cpu_id]
      if total > cpu.total_prev then -- guard against 1/0 errors
         cpu.percent_active = (active - cpu.active_prev) / (total - cpu.total_prev)
         cpu.active_prev = active
         cpu.total_prev = total
      end
   end
   return cpu_loads
end

return M