# -*- coding: utf-8 -*-
import sys
import json
import pylua
import argparse


__TEMPLATE = '''\
function string.starts(String,Start)
   return string.sub(String,1,string.len(Start))==Start
end

local function extract_backends()
    {data}
    local res = {{}};

    for key, value in pairs(_G) do
        if key and string.starts(key, 'backends_') then
            res[key] = value;
        end;
    end;
    return res;
end

return require('json').encode(extract_backends())
'''


__ALL_TEMPLATE = '''\
function string.starts(String,Start)
   return string.sub(String,1,string.len(Start))==Start
end

local function extract_backends()
    {data}
    local res = {{}};

    for key, value in pairs(_G) do
        if key and (string.starts(key, 'backends_') or string.starts(key, 'ipdispatch_')) then
            res[key] = value;
        end;
    end;
    return res;
end

return require('json').encode(extract_backends())
'''


def main():
    parser = argparse.ArgumentParser(description='Convert lua balancer config to json')
    parser.add_argument('infile', type=argparse.FileType('r'), help='original lua config')
    parser.add_argument('outfile', nargs='?', type=argparse.FileType('w'), default=sys.stdout,
                        help='output file')
    parser.add_argument('--pretty', action='store_true', help='enable pretty printing')
    parser.add_argument('--loadtesting', action='store_true', help='replace backends for load_testing')
    parser.add_argument('--functesting', action='store_true', help='replace backends for finc_testing')
    args = parser.parse_args()

    with args.infile as f:
        tmpl = __TEMPLATE
        if args.functesting:
            tmpl = __ALL_TEMPLATE
        code = tmpl.format(data=f.read())
        json_data = json.loads(pylua.eval_raw(code))

    indent = 4 if args.pretty else None

    if args.loadtesting:
        lite = json_data.get('backends_VLA_WEB_RUS_YALITE_SAS_WEB_RUS_YALITE_MAN_WEB_RUS_YALITE', [])
        for k in list(json_data.keys()):
            if not k.startswith('backends'):
                continue
            kl = k.lower()
            if 'laas' in kl:
                json_data[k] = [{'host': 'localhost', 'cached_ip': '127.0.0.1', 'port': 31137, 'weight': 1}]
            elif 'antirobot' in kl:
                json_data[k] = [{'host': 'localhost', 'cached_ip': '127.0.0.1', 'port': 31140, 'weight': 1}]
            elif 'uaas' in kl:
                json_data[k] = [{'host': 'localhost', 'cached_ip': '127.0.0.1', 'port': 31138, 'weight': 1}]
            elif 'click' in kl:
                json_data[k] = [{'host': 'localhost', 'cached_ip': '127.0.0.1', 'port': 31139, 'weight': 1}]
            elif 'wfront' in kl or 'yaru' in kl or 'report' in kl:
                json_data[k] = lite
            else:
                del json_data[k]

    json_data['_'] = [{'host': 'localhost', 'cached_ip': '127.0.0.1', 'port': 1, 'weight': 1}]

    with args.outfile as f:
        json.dump(json_data, f, indent=indent)


if __name__ == '__main__':
    main()
