#!/usr/bin/env python
# -*- coding: utf-8 -*-

import argparse
import math
import psutil
import io

CONF_FILE = "/etc/postgresql/9.5/main/conf.artix.d/conf.artix.pgtune.conf"

memInfo = psutil.virtual_memory()

parser = argparse.ArgumentParser()
parser.add_argument("--memory", dest="memory", type=int, help='RAM size in GB, int value, if not set, get system memory info')
parser.set_defaults(memory=memInfo.total / 1024 / 1024 / 1024)


def convertSize(size):
    if (size == 0):
        return '0B'
    size_name = ("KB", "MB", "GB", "TB", "PB", "EB", "ZB", "YB")
    i = int(math.floor(math.log(size, 1024)))
    p = math.pow(1024, i)
    s = int(round(size / p, 2))
    return '%s%s' % (s, size_name[i])


args = parser.parse_args()
# half of memory in our case get mysql
memoryInKb = args.memory * 1024 * 1024 / 2
if memoryInKb < 1 * 1024 * 1024:
    memoryInKb = 1 * 1024 * 1024
print "Set %s of memory" % convertSize(memoryInKb)

pgconfParameters = {
    "max_connections": 100,
    "shared_buffers": convertSize(memoryInKb / 4),
    "effective_cache_size": convertSize(memoryInKb * 3 / 4),
    "work_mem": 1,
    "maintenance_work_mem": 1,
    "min_wal_size": 1,
    "max_wal_size": 1,
    "checkpoint_completion_target": 1,
    "wal_buffers": 1,
    "default_statistics_target": 1
}

gConfigRaw = {}
gConfigRaw["max_connections"] = 100
gConfigRaw["shared_buffers"] = memoryInKb / 4
gConfigRaw["effective_cache_size"] = memoryInKb * 3 / 4
workMem = (memoryInKb - gConfigRaw["shared_buffers"]) / (gConfigRaw["max_connections"] * 3)
gConfigRaw["work_mem"] = workMem / 2
gConfigRaw["maintenance_work_mem"] = memoryInKb / 16
# Cap maintenance RAM at 2GB on servers with lots of memory
if gConfigRaw["maintenance_work_mem"] > 2 * 1024 * 1024:
    gConfigRaw["maintenance_work_mem"] = 2 * 1024 * 1024
gConfigRaw["min_wal_size"] = "1GB"
gConfigRaw["max_wal_size"] = "2GB"
gConfigRaw["checkpoint_completion_target"] = 0.9
gConfigRaw["wal_buffers"] = 3 * gConfigRaw["shared_buffers"] / 100
# Follow auto-tuning guideline for wal_buffers added in 9.1, where it's
# set to 3% of shared_buffers up to a maximum of 16MB.
# It's nice of wal_buffers is an even 16MB if it's near that number.  Since
# that is a common case on Windows, where shared_buffers is clipped to 512MB,
# round upwards in that situation
if gConfigRaw["wal_buffers"] > 16 * 1024 or (14 * 1024 < gConfigRaw["wal_buffers"] > 16 * 1024 and gConfigRaw["wal_buffers"] > 16 * 1024 < 16 * 1024):
    gConfigRaw["wal_buffers"] = 16 * 1024
gConfigRaw["default_statistics_target"] = 100

gConfig = dict(gConfigRaw)
for key in ["shared_buffers", "effective_cache_size", "work_mem", "maintenance_work_mem", "wal_buffers"]:
    gConfig[key] = convertSize(gConfigRaw[key])

with io.open(CONF_FILE, "w", encoding="UTF-8") as conffile:
    conffile.write(u"#     DO NOT EDIT THIS FILE BY HAND -- YOUR CHANGES WILL BE OVERWRITTEN AFTER NEXT artix-posgtresql-config PACKAGE INSTALL\n")
    for key in sorted(gConfig.keys()):
        conffile.write(u"%s = %s\n" % (key, gConfig[key]))
