#!/usr/bin/python3

import glob
import json
import os
import shutil
import subprocess
import tarfile


def generate_cache():
    print('Creating pickles')

    arcadia_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), '../../..')

    cfgs = [p.split('/')[-2] for p in glob.glob(os.path.join(arcadia_path, 'balancer/config/build/*/ya.make'))]
    bin_dir_path = os.path.join(arcadia_path, 'gencfg/custom_generators/balancer_gencfg/utils')
    bin_path = os.path.join(bin_dir_path, 'utils')

    subprocess.check_call(['ya', 'm', '-r', '--checkout', bin_dir_path])

    generators = []

    shutil.rmtree('temp', ignore_errors=True)
    os.mkdir('temp')

    for cfg in cfgs:
        cfg_prj = os.path.basename(cfg).split('_')[0]
        cfg_path = os.path.join('temp', cfg + '.cfg')
        pickle_path = os.path.join('temp', cfg + '_gencfg_cache.pickle')
        if cfg_prj == "any":
            any_dc = os.path.basename(cfg).split('_')[1]
            cmd = [
                bin_path, '-b', 'any',
                '-p', any_dc, '-t', 'api', '-o', cfg_path, '--disable-trunk-fallback',
                '--enable-cache', '--save-cache', pickle_path
            ]
        else:
            cmd = [
                bin_path, '-b', 'l7heavy',
                '-p', cfg, '-t', 'api', '-o', cfg_path, '--disable-trunk-fallback',
                '--enable-cache', '--save-cache', pickle_path
            ]
        generators += [subprocess.Popen(cmd, stdout=open(os.devnull, 'w'))]

    for proc in generators:
        assert proc.wait() == 0

    print('Packing pickles')

    pickles_file = os.path.join('temp', 'pickles.tgz')
    tar = tarfile.open(pickles_file, 'w:gz')

    for p in glob.glob('temp/*.pickle'):
        tar.add(p, p.split('/')[-1])
    tar.close()

    print('Uploading cache to sandbox')
    proc = subprocess.Popen(['ya', 'upload', '--ttl=inf', '--json-output', pickles_file], stdout=subprocess.PIPE)
    sbr_id = json.loads(proc.communicate()[0])['resource_id']
    print('Uploaded sandbox resource with id {}'.format(sbr_id))
    assert proc.wait() == 0

    with open('SbrId.inc', 'w') as f:
        f.write('SET(SBR_ID {})\n'.format(sbr_id))

    shutil.rmtree('temp')


generate_cache()
