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

import time
import os.path as osp
import shutil
import unittest
import re

from common import dict_conf, execcmd, tmpdir
from data   import server_export, available_hosts

from asrun.run    import AsRunFactory
from asrun.job    import parse_actu_result, parse_tail_result
from asrun.profil import AsterProfil, MODES
from asrun.client import ClientConfig, AsterCalcHandler
from asrun.client.config import serv_infos_prof
from asrun.calcul import parse_submission_result
from asrun.common.utils import YES_VALUES, NO_VALUES
from asrun.common.rcfile import read_rcfile
from asrun.common.sysutils import same_hosts


def client():
    # local astkrc directory for test
    rcdir = osp.join(tmpdir, 'astkrc')
    if osp.exists(rcdir):
        shutil.rmtree(rcdir)

    run = AsRunFactory(rcdir=rcdir)
    client = ClientConfig(rcdir)

    # test the default configuration (rcdir does not exist)
    prefs = client.get_user_preferences()
    assert type(prefs) is dict
    assert len(prefs) == 11, prefs  # number of fields in default prefs file
    assert 'nom_domaine' in prefs

    # before the refresh the values are not set
    client.init_server_config()
    assert client.get_server_rcinfo() is None
    servers = client.get_server_list()
    assert len(servers) == 3
    cfg = client.get_server_config(servers[0])
    if cfg['nom_complet'].find('?') > -1:
        # not configured: '?FULL_SERVER_NAME?'
        cfg['nom_complet'] = ''
    assert cfg.get('batch') is None
    assert cfg.get('interactif_memmax') is None

    # ask for a refresh
    client.refresh_server_config()

    servers = client.get_server_list()
    assert len(servers) == 3
    cfg = client.get_server_config(servers[0])
    assert 'batch' in cfg and cfg['batch'] == "non", cfg
    assert type(cfg['interactif_memmax']) in (int, int)
    assert cfg['interactif_memmax'] > 1
    byip = client.get_server_config(cfg["nom_complet"], use_ip=True)
    assert cfg == byip, (cfg, byip)

    # test a real configuration
    shutil.copy(osp.join(dict_conf['DATA'], "config_serveurs"), rcdir)
    client.init_server_config(force=True)
    assert client.get_server_rcinfo() == '1.1'

    servers = client.get_server_list()
    assert len(servers) == 2, servers
    for serv in servers:
        cfg = client.get_server_config(serv)
        assert cfg['etat'] in ('on', 'off')
    assert "localhost" in servers
    cfg = client.get_server_config("localhost")
    assert cfg["rep_serv"].find("ASTK_SERV") < 0, cfg["rep_serv"]
    all_cfg = client.get_dict_server_config()
    assert all_cfg["localhost"]["interactif"] in (YES_VALUES + NO_VALUES)

    # fails outside EDF network
    iret, output = run.Shell("ping -c 1 -W 2 mailhost.der.edf.fr")
    if iret == 0:
        client.init_server_config(refresh=True)

    # test refresh of available hosts
    servers = client.get_server_list()
    for sname in list(available_hosts.keys()):
        found = False
        for label in servers:
            cfg = client.get_server_config(label)
            if same_hosts(sname, cfg['nom_complet']):
                client.refresh_server_config([label])
                cfg = client.get_server_config(label)
                found = True
                break
        assert found, "%s not found" % sname
    # check save
    client.save_server_config()

def _get_protocol(main_dict, type, default=None, optional=None):
    """Return the protocol found in config dicts."""
    assert type in ('shell', 'copy'), type
    key = 'remote_%s_protocol' % type.lower()
    main = main_dict.get(key)
    optv = main
    if optional is not None:
        optv = optional.get(key)
    #if main != optv:
        #print "'%s' differs '%s' (prefs) != '%s' (config)" % (key, main, optv)
    proto = main or optv or default
    print("shell protocol: '%s' (prefs), '%s' (config), '%s' (default) => '%s'" \
        % (main, optv, default, proto))
    return proto

def module_aster():
        # Test the usage done by the Code_Aster module :
        #   - compare remote_xxxx_protocol
        from pprint import pprint
        rcdir = osp.join(tmpdir, 'astkrc')
        #rcdir = ".astkrc"
        if osp.exists(rcdir):
            shutil.rmtree(rcdir)

        run = AsRunFactory(rcdir=rcdir)
        client = ClientConfig(rcdir)
        prefs = client.get_user_preferences()
        #TODO should be removed when config/prefs will be merged or proto imposed by server
        dconf = {}
        read_rcfile(client.rcfile('prefs'), dconf)
        for type, defv in (('shell', 'SSH'),
                           ('copy', 'SCP')):
            key = 'remote_%s_protocol' % type
            proto = _get_protocol(prefs, type, defv, optional=dconf)
            client.set_user_preferences(key, proto)
        client.save_user_preferences()
        # test saving
        prefs = client.get_user_preferences(force=True)
        assert prefs.get(key) is not None and prefs[key] == proto, prefs.get(key)

def use_proxy(prof):
        nomjob = prof['nomjob'][0]
        export = prof.get_filename()
        prof.WriteExportTo(export)

        cmd = dict_conf["as_run"] + ["--proxy", "--serv", export]
        iret, out = execcmd(cmd, "%s.1" % nomjob, return_output=True)
        assert iret == 0
        jobid, queue, studyid = parse_submission_result(out)
        assert jobid.strip() != ""
        assert studyid.strip() != ""

        # here, the client could build a minimal export with user, server, protocol_exec
        with open(export, "a") as f:
            f.write("P jobid %s\n" % jobid)
        # test api : call_plugin("actu", ...)

        _keyword = 'Iteration'
        rekw = re.compile(_keyword, re.M)
        with open(export, "a") as f:
            f.write("P tail_nbline 100\n")
        with open(export, "a") as f:
            f.write("P tail_regexp %s\n" % _keyword)
        etat = "RUN"
        i = 0
        found = 0
        while etat in ("PEND", "RUN"):
            cmd = dict_conf["as_run"] + ["--proxy", "--actu", export]
            iret, out = execcmd(cmd, "%s.2" % nomjob, return_output=True)
            assert iret == 0
            etat, diag, node, tcpu, wrk, queue = parse_actu_result(out)
            if etat == "RUN":
                i += 1
                cmd = dict_conf["as_run"] + ["--proxy", "--tail", export]
                iret, out = execcmd(cmd, "%s.3.%d" % (nomjob, i), return_output=True)
                job, jobid_, etat, diag = parse_tail_result(out)
                found += len(rekw.findall(out))

        assert etat == "ENDED", etat
        assert diag == "<A>_ALARM", diag
        assert found > 0

        time.sleep(4)
        cmd = dict_conf["as_run"] + ["--proxy", "--del", export]
        iret = execcmd(cmd, "%s.5" % nomjob)
        assert iret == 0

        cmd = dict_conf["as_run"] + ["--proxy", "--info", export]
        iret = execcmd(cmd, "%s.6" % nomjob)
        assert iret == 0

def use_proxy_calcul():
        # same as test02_proxy using a AsterCalcHandler object.
        assert len(available_hosts) > 0, "no remote host available."
        run = AsRunFactory(rcdir=osp.join(dict_conf['TMPDIR'], 'astkrc'))
        dico = dict_conf.copy()
        dico['nomjob'] = 'proxy_calcul'
        prof = AsterProfil(run=run)
        prof.parse(server_export % dico)
        pserv = init_prof_aster5()
        prof['mclient'] = 'clienthost'
        prof.update(pserv)
        calcul = AsterCalcHandler(prof)
        #calcul.on_host(dict_conf['localhost'], None)
        host = dict_conf['remhost']
        root = available_hosts[host]
        calcul.on_host(host, None)
        prof['aster_root'] = root
        assert prof['serveur'][0], prof['serveur']
        req = calcul.request('mem')
        assert req == 500.0, prof
        req = calcul.request('cpu')
        assert req == 1, req
        jret, out = calcul.start()
        if jret != 0:
            print(out)
        assert jret == 0, (jret, out, prof)
        time.sleep(0.2)
        res = calcul.get_state()
        calcul.wait(1.1)
        res = calcul.get_diag()
        assert res[0] == '<A>_ALARM', "res=%s, more details on %s in %s" \
            % (res, host, calcul.flash('output'))
        calcul.kill()

def init_prof_aster5():
    host = "aster5"
    assert host in [h.split('.')[0] for h in list(available_hosts.keys())], host
    rcdir = osp.join(tmpdir, 'astkrc')
    if osp.exists(rcdir):
        shutil.rmtree(rcdir)
    run = AsRunFactory(rcdir=rcdir)
    client = ClientConfig(rcdir)
    shutil.copy(osp.join(dict_conf['DATA'], "config_serveurs"), rcdir)
    client.init_server_config()
    cfg = client.get_server_config(host)
    assert cfg, "Can not read configuration of aster5"
    pserv = serv_infos_prof(cfg)
    pserv['aster_root'] = available_hosts[dict_conf['remhost']]
    return pserv

class TestClientFunctions(unittest.TestCase):

    def test01_client(self):
        client()

    def test02_module_aster(self):
        module_aster()

    def test03_proxy(self):
        # the same test from server side is 110_server.test01_from_cmdline
        export = osp.join(tmpdir, "proxy.export")
        dico = dict_conf.copy()
        #dico['localhost'] = '127.0.0.1'
        dico['nomjob'] = 'proxy_cmdline'
        with open(export, "w") as f:
            f.write(server_export % dico)
        prof = AsterProfil(export)
        prof.set_running_mode(MODES.INTERACTIF)
        use_proxy(prof)

    def test03_proxy_aster5(self):
        assert len(available_hosts) > 0, "no remote host available."
        pserv = init_prof_aster5()
        # the same test from server side is 110_server.test01_from_cmdline
        export = osp.join(tmpdir, "proxy_batch.export")
        dico = dict_conf.copy()
        #dico['localhost'] = '127.0.0.1'
        dico['nomjob'] = 'proxy_aster5'
        with open(export, "w") as f:
            f.write(server_export % dico)
        prof = AsterProfil(export)
        prof.set_running_mode(MODES.BATCH)
        prof.update(pserv)
        print(prof)

        use_proxy(prof)

    def test04_proxy_calcul(self):
        use_proxy_calcul()


if __name__ == "__main__":
    #unittest.main()
    #available_hosts = { 'aster4.cla.edfgdf.fr' : '/home/courtois/src/astk' }
    #module_aster()
    if True:
        from pprint import pprint
        rcdir = osp.join(tmpdir, 'astkrc')
        #rcdir = ".astkrc"
        if osp.exists(rcdir):
            shutil.rmtree(rcdir)

        run = AsRunFactory(rcdir=rcdir)
        client = ClientConfig(rcdir)
        prefs = client.get_user_preferences()
        client.save_user_preferences()
