# -*- coding: utf-8 -*-
"""Base for drivers

:copyright: Copyright (c) 2019 RadiaSoft LLC.  All Rights Reserved.
:license: http://www.apache.org/licenses/LICENSE-2.0.html
"""
from pykern import pkconfig, pkio, pkinspect, pkcollections, pkconfig, pkjson
from pykern.pkcollections import PKDict
from pykern.pkdebug import pkdp, pkdlog, pkdc, pkdexc, pkdformat
from sirepo import job
import asyncio
import importlib
import pykern.pkio
import re
import sirepo.auth
import sirepo.events
import sirepo.sim_db_file
import sirepo.simulation_db
import sirepo.srcontext
import sirepo.srdb
import sirepo.tornado
import time
import tornado.gen
import tornado.ioloop
import tornado.locks
import tornado.queues


KILL_TIMEOUT_SECS = 3

#: map of driver names to class
_CLASSES = None

#: default class when not determined by op
_DEFAULT_CLASS = None

_DEFAULT_MODULE = 'local'

cfg = None

OPS_THAT_NEED_SLOTS = frozenset((job.OP_ANALYSIS, job.OP_RUN))

_UNTIMED_OPS = frozenset((job.OP_ALIVE, job.OP_CANCEL, job.OP_ERROR, job.OP_KILL, job.OP_OK))


class AgentMsg(PKDict):

    async def receive(self):
        # Agent messages do not block the supervisor
        DriverBase.receive(self)


def assign_instance_op(op):
    m = op.msg
    if m.jobRunMode == job.SBATCH:
        res = _CLASSES[job.SBATCH].get_instance(op)
    else:
        res = _DEFAULT_CLASS.get_instance(op)
    assert m.uid == res.uid, \
        'op.msg.uid={} is not same as db.uid={} for jid={}'.format(
            m.uid,
            res.uid,
            m.computeJid,
        )
    res.ops[op.opId] = op
    return res


class DriverBase(PKDict):

    __instances = PKDict()

    _AGENT_STARTING_SECS_DEFAULT = 5

    def __init__(self, op):
        super().__init__(
            driver_details=PKDict({'type': self.__class__.__name__}),
            kind=op.kind,
            ops=PKDict(),
            #TODO(robnagler) sbatch could override OP_RUN, but not OP_ANALYSIS
            # because OP_ANALYSIS touches the directory sometimes. Reasonably
            # there should only be one OP_ANALYSIS running on an agent at one time.
            op_slot_q=PKDict({k: job_supervisor.SlotQueue() for k in OPS_THAT_NEED_SLOTS}),
            uid=op.msg.uid,
            _agentId=job.unique_key(),
            _agent_start_lock=tornado.locks.Lock(),
            _agent_starting_timeout=None,
            _idle_timer=None,
            _websocket=None,
            _websocket_ready=sirepo.tornado.Event(),
        )
        self._sim_db_file_token = sirepo.sim_db_file.token_for_user(self.uid)
        # Drivers persist for the life of the program so they are never removed
        self.__instances[self._agentId] = self
        pkdlog('{}', self)

    def destroy_op(self, op):
        """Clear our op and (possibly) free cpu slot"""
        self.ops.pkdel(op.opId)
        op.cpu_slot.free()
        if op.op_slot:
            op.op_slot.free()
        if 'lib_dir_symlink' in op:
            # lib_dir_symlink is unique_key so not dangerous to remove
            pykern.pkio.unchecked_remove(op.pkdel('lib_dir_symlink'))

    def free_resources(self, internal_error=None):
        """Remove holds on all resources and remove self from data structures"""
        pkdlog('{} internal_error={}', self, internal_error)
        try:
            self._agent_starting_done()
            self._websocket_ready.clear()
            w = self._websocket
            self._websocket = None
            if w:
                # Will not call websocket_on_close()
                w.sr_close()
            for o in list(self.ops.values()):
                o.destroy(internal_error=internal_error)
            self._websocket_free()
        except Exception as e:
            pkdlog('{} error={} stack={}', self, e, pkdexc())

    async def kill(self):
        raise NotImplementedError(
            'DriverBase subclasses need to implement their own kill',
        )

    def make_lib_dir_symlink(self, op):
        import sirepo.auth_db
        m = op.msg
        with sirepo.auth_db.session(), \
             sirepo.auth.set_user_outside_of_http_request(m.uid):
            d = sirepo.simulation_db.simulation_lib_dir(m.simulationType)
            op.lib_dir_symlink = job.LIB_FILE_ROOT.join(
                job.unique_key()
            )
            op.lib_dir_symlink.mksymlinkto(d, absolute=True)
            m.pkupdate(