"""
This Module defines the AbipyWorker object
"""
from __future__ import annotations

import sys
import os
import threading
import json
import socket
import time
import subprocess
import tempfile
import getpass
import functools
import tornado
import tornado.web
import tornado.escape
import pandas as pd
import requests
import panel as pn

from datetime import datetime
from uuid import uuid4
#from queue import Queue, Empty
from typing import Union, Optional  # , Type  # Awaitable
from enum import Enum
from monty.json import MSONable
from monty.termcolor import colored
#from monty.functools import lazy_property
from tqdm import tqdm

from abipy.tools.serialization import pmg_serialize
from abipy.tools.printing import print_dataframe
from abipy.tools import duck
from abipy.tools.iotools import yaml_safe_load_path
from abipy.flowtk.flows import Flow
from abipy.flowtk.tasks import TaskManager, set_user_config_taskmanager
from abipy.flowtk.launcher import MultiFlowScheduler
from .tools import pid_exists, port_is_open, find_free_port
from .base_models import AbipyModel, Field, MongoConnector, AbipyEncoder


# This should become a global variable used in all the other modules to facilitate unit tests.
ABIPY_DIRPATH = os.path.join(os.path.expanduser("~"), ".abinit", "abipy")


class BaseHandler(tornado.web.RequestHandler):
    """
    Our base class for RequestHandler used by the AbipyWorker bokeh server.
    Each handler has a reference to the worker set in `initialize`.
    """

    #def data_received(self, chunk: bytes) -> Optional[Awaitable[None]]:
    #    pass

    def initialize(self, worker: AbipyWorker):
        self.worker = worker

    def get_json_from_body(self):
        try:
            data = tornado.escape.json_decode(self.request.body)
            return data
        except ValueError:
            # TODO: handle the error
            raise

    #def write_error(self, status_code, **kwargs):
    #    # https://stackoverflow.com/questions/26371051/better-way-to-handle-errors-in-tornado-request-handler
    #   import traceback, logging
    #   if status_code == 500:
    #      excp = kwargs['exc_info'][1]
    #      tb   = kwargs['exc_info'][2]
    #      stack = traceback.extract_tb(tb)
    #      clean_stack = [i for i in stack if i[0][-6:] != 'gen.py' and i[0][-13:] != 'concurrent.py']
    #      error_msg = '{}\n  Exception: {}'.format(''.join(traceback.format_list(clean_stack)),excp)

    #      # do something with this error now... e.g., send it to yourself
    #      # on slack, or log it.
    #      logging.error(error_msg)  # do something with your error...

    #   # don't forget to show a user friendly error page!
    #   self.render("oops.html")


class ActionHandler(BaseHandler):

    SUPPORTED_METHODS = ("POST",)

    def post(self):
        data = self.get_json_from_body()
        action = data.get("action", None)
        if action is None: return
        reply = {}

        if action == "kill":
            reply["message"] = "Received kill action. Aborting now!"
            self.write(reply)
            # TODO: Should lock and give a change to the secondary thread to complete operations
            time.sleep(4)
            sys.exit(0)

        #elif action == "rm_errored_flows":
        #    self.worker.scheduler.rm_flows_with_status("Error")
        #    reply["message"] = f"Removed: {nflows} Errored Flows"

        else:
            reply["message"] = f"Unknown action: {action}"

        self.write(reply)


class PostFlowHandler(BaseHandler):
    """
    This endpoint allows the client to send a python script to the server.
    """

    SUPPORTED_METHODS = ("POST",)

    def post(self):
        #if self.worker.flow_model is not None: raise RuntimeError("Cannot send stuff in flow_model is ON")
        data = self.get_json_from_body()

        pyscript_basename = data["pyscript_basename"]
        pyscript_string = data["pyscript_string"]
        user_message = data["user_message"]

        scratch_dir = self.worker.scratch_dir
        filepath = os.path.join(scratch_dir, pyscript_basename)
        with open(filepath, "wt") as fh:
            fh.write(pyscript_string)

        now = datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
        workdir = tempfile.mkdtemp(prefix=f"Flow-{now}", dir=scratch_dir)

        # FIXME: problem with manager and possible collisions in filepath
        p = subprocess.Popen(f"python {filepath} -w {workdir}",
                             stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
        out, err = p.communicate()

        reply = dict(
            returncode=p.returncode,
            out=out.decode("utf-8") if out is not None else None,
            err=err.decode("utf-8") if err is not None else None,
            workdir=workdir,
        )

        if p.returncode == 0:
            print(f"Running new Flow generated by {pyscript_basename} in workdir: {workdir}")
            flow = Flow.from_file(workdir)
            self.worker.flow_scheduler.add_flow(flow, user_message)

        self.write(reply)


class _JsonHandler(BaseHandler):

    def prepare(self):
        self.set_header(name="Content-Type", value="application/json")


class JsonStatusHandler(_JsonHandler):
    """
    This endpoint allows the client to get a json document with the status of the server.
    """

    SUPPORTED_METHODS = ("GET",)

    def get(self):
        # FIXME: This is broken now
        json_status = self.worker.flow_scheduler.get_json_status()
        print("json_status", json_status)
        self.write(json_status)


class FlowDirsPostHandler(BaseHandler):
    """
    This endpoint allows the client to send to the server a list of directories
    containing pregenerated Flows. It is assumed that both the client and the server
    are running on the same host machine.
    """

    SUPPORTED_METHODS = ("POST",)

    def post(self):
        data = self.get_json_from_body()
        user_message = data["user_message"]
        flow_dir = data["flow_dir"]

        error = ""
        try:
            flow = Flow.from_file(flow_dir)
            self.worker.flow_scheduler.add_flow(flow, user_message)
        except Exception as exc:
            error = str(exc)

        reply = dict(
            returncode=0 if not error else 1,
            error=error,
        )

        self.write(reply)


class WorkerStatusEnum(str, Enum):
    """
    Possible status of an AbipyWorker instance.
    """
    init = "init"
    serving = "serving"
    dead = "dead"

    def colored_str(self) -> str:
        s = self.value
        color = {"init": "blue", "serving": "green",  "dead": "red"}[s]
        return colored(s, color=color)


def _str_uuid() -> str:
    return str(uuid4())


class WorkerState(AbipyModel):

    name: str = Field(..., description="Name of the AbiPy worker. It is assumed to be unique.")

    status: WorkerStatusEnum = Field(WorkerStatusEnum.init, description="Status of the worker.")

    pid: int = Field(default_factory=os.getpid, description="Process identifier")

    address: str = Field("localhost", description="Server address.")

    port: int = Field(0, description="Server port. 0 means find a free port at runtime.")

    scratch_dir: str = Field(..., description="Scratch directory used by the Worker to generate Flows.")

    username: str = Field(default_factory=getpass.getuser, description="User name at the host")

    hostname: str = Field(default_factory=socket.gethostname, description="Host name.")

    #uuid: UUID = Field(default_factory=uuid4)
    uuid: str = Field(default_factory=_str_uuid,
                      description="AbipyWorker ID. Used by the client to detect if the server has restarted.")

    mng_connector: MongoConnector = Field(None, description="MongoDB connection parameters and name of default collection")



#class FlowContext(BaseModel):
#    workdir: str
#    manager: TaskManager
#    mng_connector: MongoConnector


class AbipyWorker:
    """
    An AbipyWorker is a bokeh web-server running on the same machine where Abipy Flows are executed.
    An AbipyWorker has a name that is supposed to be unique across different machines.
    Flows are scheduled according to the `sched_options` and executed inside the `scratch_dir` directory.

    A command line interface to create/start/restart Workers is provided by the `abiw.py` script.
    """
    MAX_NUM_EXCS_IN_BUILDFLOW = 30

    #MONGO_PERIODIC_CALLBACK_TIME_MS = 10000
    MONGO_PERIODIC_CALLBACK_TIME_MS = 5000

    def __init__(self, name: str, sched_options: dict, scratch_dir: str,
                 address: str = "localhost", port: int = 0,
                 manager: Optional[TaskManager] = None,
                 mng_connector: Optional[MongoConnector] = None):
        """
        Args:
            name: The name of the AbipyWorker. Must be unique across machines.
            sched_options:
            scratch_dir:
            address:
                The address the server should listen on for HTTP requests.
            port:
                 Allows specifying a specific port
                 Port 0 means to select an arbitrary unused port
            manager:
            mng_connector:
        """
        self.name = name
        self.address, self.port = address, port
        self.pid = os.getpid()
        self.manager = manager
        self.mng_connector = mng_connector
        self.flow_model_cls = None
        self.num_exceptions_in_build_flow_from_mongodb = 0
        self.flowid2_oid_model = {}

        if self.mng_connector:
            # Get the FlowMode subclass from the DB collection.
            collection = self.mng_connector.get_collection()
            from .flow_models import FlowModel
            self.flow_model_cls = FlowModel.get_subclass_from_collection(collection)

        # url --> callables returning panel objects.
        self.routes = {
            "/": self.serve_worker_homepage,
            r"/flow/\d+": self.serve_panel_flow,
        }

        # url --> tornado handler
        self.extra_patterns = [
            ("/json_status", JsonStatusHandler, dict(worker=self)),
            ("/action", ActionHandler, dict(worker=self)),
        ]

        if self.flow_model_cls is None:
            # These endpoints are activated only if we are not using a MongoDB collection.
            self.extra_patterns.extend([
                ("/post_flow_script", PostFlowHandler, dict(worker=self)),
                ("/post_flow_dirs", FlowDirsPostHandler, dict(worker=self)),
            ])

        self.config_dir = os.path.join(ABIPY_DIRPATH, f"worker_{self.name}")
        if not os.path.exists(self.config_dir):
            os.mkdir(self.config_dir)

        self.scratch_dir = scratch_dir
        if not os.path.isdir(scratch_dir):
            raise ValueError(f"Scratch directory: `{scratch_dir}` does not exist!")

        sqldb_path = os.path.join(self.config_dir, "flows.db")
        self.flow_scheduler = MultiFlowScheduler(sqldb_path, **sched_options)

        state_file = os.path.join(self.config_dir, "state.json")
        if os.path.exists(state_file):
            state = WorkerState.from_json_file(state_file)
            if state.status == state.status.serving:
                pid = state.pid
                if pid_exists(pid):
                    raise RuntimeError(f"There's already an AbiPyWorker serving on this machine with pid: {pid}")
                # FIXME This does not work in demon mode!
                print(f"Status 'serving' but pid {pid} does not exist!")

        #self.write_state_file(status="init")

    # TODO
    #@lazy_property
    #def full_name(self)
    #    return f"{self.name}@{socket.gethostname()}"

    def __repr__(self) -> str:
        return f"{self.name} at {self.address}:{self.port}. pid: {self.pid}"

    #def __str__(self) -> str:
    #    lines = []
    #    app = lines.append
    #    return "\n".join(lines)

    #def mkworkdir(self):
    #    now = datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
    #    workdir = tempfile.mkdtemp(prefix=f"Flow-{now}", dir=self.scratch_dir)

    def build_flow_from_mongodb(self) -> None:
        # Stop the server if the number of consecutive exceptions becomes too large
        # This may be due to e.g. a quota limit or other hardware-related problems.
        #print("in build_flow")
        try:
            self._build_flow_from_mongodb()
            if self.num_exceptions_in_build_flow_from_mongodb > 0:
                self.num_exceptions_in_build_flow_from_mongodb -= 1

        except Exception as exc:
            print(exc)
            self.num_exceptions_in_build_flow_from_mongodb += 1

        if self.num_exceptions_in_build_flow_from_mongodb > self.MAX_NUM_EXCS_IN_BUILDFLOW:
            self.shutdown_server("Too many exceptions in build_flow_from_mongodb")

    def _build_flow_from_mongodb(self) -> None:
        """
        Periodic callback executed by the primary thread to fetch new calculations from
        the MongoDB collection when then AbipyWorker is running with a `flow_model_cls`
        """
        if self.flow_model_cls is None:
            raise ValueError("The AbipyWorker should have a flow_model_cls when running with a MongDB database.")

        # TODO: Improve error handling.
        collection = self.mng_connector.get_collection()

        if self.flowid2_oid_model:
            node_ids = list(self.flowid2_oid_model.keys())
            rows = self.flow_scheduler.get_sql_rows_with_node_ids(node_ids)
            #print("Found rows", rows)

            for row in rows:
                flow_id = row["flow_id"]
                if row["status"] == "Completed":
                    print("In completed")
                    flow = Flow.from_file(row["workdir"])
                    oid, model = self.flowid2_oid_model[flow_id]
                    model.postprocess_flow_and_update_db(flow, oid, self)
                    self.flowid2_oid_model.pop(flow_id)
                else:
                    pass
                    #print("Should handle row with status", row["status"])

        # TODO: Implement Pipelines?
        # Use atomic update of the models so that we can have multiple workers possibly on
        # different machines fetching and updating documents in the same collection.
        #print("Finding runnable_oid_models")
        oid_flowmodel_list = self.flow_model_cls.find_runnable_oid_models(collection, limit=5)
        if not oid_flowmodel_list: return
        print("in build_flow_from_mongodb, found", len(oid_flowmodel_list), "flows in database")

        for oid, model in oid_flowmodel_list:
            #print("Running model:", repr(model))
            now = datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
            workdir = tempfile.mkdtemp(prefix=f"Flow-{now}", dir=self.scratch_dir)

            flow = model.build_flow_and_update_collection(workdir, oid, self)
            if flow is None: continue
            self.flow_scheduler.add_flow(flow, user_message="")  # m.user_message ?
            self.flowid2_oid_model[flow.node_id] = (oid, model)

    @classmethod
    def _get_state_path(cls, name: str) -> tuple[WorkerState, str]:
        config_dir = os.path.join(ABIPY_DIRPATH, f"worker_{name}")
        state_filepath = os.path.join(config_dir, "state.json")

        return WorkerState.from_json_file(state_filepath), state_filepath

    @classmethod
    def init_from_dirname(cls, dirname: str) -> AbipyWorker:
        """
        Initialize the AbipyWorker from a directory basename.
        """
        state, path = cls._get_state_path(dirname)

        if state.status != state.status.init:
            raise RuntimeError(f"`status` entry in json file: `{path}`\n"
                               f"should be set to `init` while it is: `{state.status}`")

        config_dir = os.path.dirname(path)
        sched_options = yaml_safe_load_path(os.path.join(config_dir, "scheduler.yml"))
        manager = TaskManager.from_file(os.path.join(config_dir, "manager.yml"))

        # TODO: mng_connector, flow_model_cls
        return cls(state.name, sched_options, state.scratch_dir,
                   address=state.address, port=state.port, manager=manager,
                   mng_connector=state.mng_connector)

    @classmethod
    def new_with_name(cls, worker_name: str, scratch_dir: str,
                      scheduler_path: Optional[str] = None,
                      manager_path: Optional[str] = None,
                      mng_connector=None, verbose: int = 1) -> AbipyWorker:
        """
        Create a new AbiPyWorker.

        Args:
            worker_name:
            scratch_dir:
            scheduler_path:
            manager_path:
            mng_connector:
            verbose:
        """
        config_dir = os.path.join(ABIPY_DIRPATH, f"worker_{worker_name}")
        errors = []
        eapp = errors.append

        if os.path.exists(config_dir):
            eapp(f"Directory `{config_dir}` already exists!")

        scheduler_path = scheduler_path or os.path.join(ABIPY_DIRPATH, "scheduler.yml")
        manager_path = manager_path or os.path.join(ABIPY_DIRPATH, "manager.yml")

        if not os.path.exists(scheduler_path):
            eapp(f"Cannot find scheduler file at `{scheduler_path}` to initialize the worker!")
        if not os.path.exists(manager_path):
            eapp(f"Cannot find manager file at `{manager_path}` to initialize the worker!")
        if not os.path.isdir(scratch_dir):
            eapp(f"scratch_dir: `{scratch_dir}` is not a directory!")

        if errors:
            lines = ["The following problems have been detected:"]
            for i, err in enumerate(errors):
                lines.append(f"\t [{i+1}] {err}")
            raise RuntimeError("\n".join(lines))

        if verbose:
            print(f"""
Creating new worker directory `{config_dir}`
Copying manager.yml and scheduler.yml files into it.
These are the configuration files used by the AbipyWorker to submit and schedule tasks.
You may want to customize these files before running calculations.

After this step, use:

    abiw.py start {worker_name}

to start the AbiPy worker, and:

    abiw.py lscan

to update the list of local clients.
""")

        os.mkdir(config_dir)
        from shutil import copy
        copy(scheduler_path, config_dir)
        copy(manager_path, config_dir)

        worker_state = WorkerState(
                name=worker_name,
                status=WorkerStatusEnum.init,
                scratch_dir=scratch_dir,
                mng_connector=mng_connector,
        )

        worker_state.json_write(os.path.join(config_dir, "state.json"), indent=2)

        return cls.init_from_dirname(worker_name)

    @classmethod
    def restart_from_dirname(cls, name: str, force: bool = False) -> AbipyWorker:
        """
        Restart the AbipyWorker given its name. Must be executed on the same machine
        where the Worker resides.
        """
        try:
            state, path = cls._get_state_path(name)
        except FileNotFoundError as exc:
            print(exc)
            raise ValueError(f"{name} is a remote Worker.\n"
                             "To restart it, execute the command on the same machine where the worker resides.")

        pid = state.pid

        #if state.status != "dead":
        # TODO: status can be set to serving if the worker has been killed by kill -9
        # should implement a restart_hard option and perhaps make a bkp of the database!
        if state.status == state.status.serving:
            e_pid = pid_exists(pid)

            msg = f"It seems there's already a worker serving on this machine with pid: {pid}"
            if not force:
                if not e_pid:
                    print(f"Status is set to `serving` but pid {pid} does not exist!")
                else:
                    raise RuntimeError(msg)
            else:
                print(msg)
                print("but as `force` is specified by the user....")
                if e_pid:
                    raise RuntimeError("But pid exists so it's dangereous")

        # TODO: Problem with the default manager when creating the flow.
        # Replace node_id with uuid4
        config_dir = os.path.dirname(path)
        sched_options = yaml_safe_load_path(os.path.join(config_dir, "scheduler.yml"))
        manager = TaskManager.from_file(os.path.join(config_dir, "manager.yml"))

        # TODO: manager, mng_connector, flow_model_cls
        # Should be stored in state.json so that we can restart.
        new = cls(state.name, sched_options, state.scratch_dir, address=state.address, port=state.port,
                  manager=manager, mng_connector=state.mng_connector)

        print("Remember to execute lscan (rscan) to update the list of local (remote) clients...")
        return new

    def write_state_file(self, status: str = "dead", filepath: Optional[str] = None) -> str:
        """
        Write/update the state.json file of the worker.
        """
        if filepath is None:
            filepath = os.path.join(self.config_dir, "state.json")

        state = WorkerState(
            name=self.name,
            status=status,
            pid=self.pid,
            address=self.address,
            port=self.port,
            scratch_dir=self.scratch_dir,
            mng_connector=self.mng_connector,
        )

        state.json_write(filepath)
        return filepath

    #def remove_flows(self):
    #    """This requires locking the SQLite database."""

    def serve_worker_homepage(self):
        """
        Serve the landing page of the AbipyWorker website
        """
        import platform
        system, node, release, version, machine, processor = platform.uname()

        header = pn.pane.Markdown(f"""
# Homepage of **{self.name}** {self.__class__.__name__}

Running on {socket.gethostname()} -- system {system} -- Python {platform.python_version()}
""", sizing_mode="stretch_width")

        d = self.flow_scheduler.groupby_status()
        md_lines = ["Refresh the page to update the list of Abipy Flows."]
        if d:
            for flow_status, values in d.items():
                if flow_status == "error": flow_status = "Errored"
                md_lines.append(f"## {flow_status.capitalize()} Flows:\n")
                for row in values:
                    workdir, flow_id = row["workdir"], row["flow_id"]
                    md_lines.append(f"- [flow_id: {flow_id}, workdir: {workdir}](/flow/{flow_id})\n")
        else:
            md_lines.append("## Empty Flow list!")

        with open(os.path.join(self.config_dir, "state.json"), "rt") as fp:
            state = json.load(fp)

        from abipy.panels.viewers import JSONViewer
        main = pn.Column(
                header,
                pn.Row("## Server state", JSONViewer(state)),
                #pn.layout.Divider(),
                pn.pane.Markdown("\n".join(md_lines), sizing_mode="stretch_width"),
                sizing_mode="stretch_width",
        )

        from abipy.panels.core import get_abinit_template_cls_kwds
        tmpl_cls, tmpl_kwds = get_abinit_template_cls_kwds()
        return tmpl_cls(main=main, **tmpl_kwds)

    def serve_panel_flow(self):
        # URL example: /flow/123 where 123 is the flow id.
        tokens = pn.state.app_url.split("/")
        #print("In serve_panel_flow with app_url", pn.state.app_url, "tokens:", tokens)
        flow_id = int(tokens[2])
        flow, flow_status = self.flow_scheduler.get_flow_and_status_by_nodeid(flow_id)
        if flow is not None:
            return flow.get_panel(template=self.panel_template)

        return pn.pane.Alert(f"Cannot find Flow with node id: {flow_id}", alert_type="danger")

    def shutdown_server(self, message: str) -> None:
        """
        Shutdown the AbipyWorker.
        """
        self.write_state_file(status="dead")
        #self.flow_scheduler.stop()

        # If we are connected to a MongoDB collection, we try to log the problem in the database.
        #if self.flow_model_cls is not None:
        #    try:
        #        collection = self.mng_connector.get_collection()
        #        self.flow_model_cls.register_worker_exception(self, collection)
        #    except:
        #        pass

        print("Shuting the server down.")
        print(message)
        sys.exit(1)

    def serve(self, **serve_kwargs):
        """
        Start the webserver.
        """
        # Change the default manager returned by TaskManager.from_user_config.
        # so that we always use the manager in the AbipyWorker directory when costructing new Flows.
        assert self.manager is not None, "TaskManager should not be None"
        set_user_config_taskmanager(self.manager)

        # Register function atexit to update the state.json file if the worker gets killed.
        # Note that the callback won't be executed if we receive SIGKILL (kill -9)
        # or if an exception is raised in write_state_file or in any other registered callback.
        import atexit
        atexit.register(self.write_state_file)

        from abipy.panels.core import abipanel #, get_abinit_template_cls_kwds, AbipyParameterized
        self.panel_template = serve_kwargs.pop("panel_template", "FastList")
        abipanel(panel_template=self.panel_template)

        import abipy.panels as mod
        assets_path = os.path.join(os.path.dirname(mod.__file__), "assets")
        serve_kwargs.update(
            show=False,
            static_dirs={"/assets": assets_path},
            websocket_origin="*",
        )

        #print("serve_kwargs:", serve_kwargs)
        self.address = serve_kwargs.pop("address", self.address)
        self.port = serve_kwargs.pop("port", self.port)
        if self.port == 0:
            self.port = find_free_port(address=self.address)

        # Now write state.json with the actual port.
        self.write_state_file(status="serving")

        # Start the secondary thread for the MultiFlowScheduler.
        thread = threading.Thread(target=self.flow_scheduler.start, name="flow_scheduler", daemon=True)
        thread.start()
        #termcolor.enable(False)

        if self.flow_model_cls is not None:
            print("Abipy Worker will try to connect to the MongoDB server using:")
            print(self.mng_connector)
            print("Using FlowModel class:", self.flow_model_cls)
            print("Connecting every", self.MONGO_PERIODIC_CALLBACK_TIME_MS, "ms")
            from tornado.ioloop import PeriodicCallback
            periodic = PeriodicCallback(self.build_flow_from_mongodb, self.MONGO_PERIODIC_CALLBACK_TIME_MS)
            periodic.start()

        server = pn.serve(self.routes, port=self.port, address=self.address,
                          extra_patterns=self.extra_patterns, start=True, **serve_kwargs)

        #return server.start()


def _get_all_local_worker_states(dirpath: Optional[str] = None) -> list[WorkerState]:
    dirpath = ABIPY_DIRPATH if not dirpath else dirpath

    worker_dirs = [dirname for dirname in os.listdir(dirpath) if dirname.startswith("worker_")]
    states = []
    for workdir in worker_dirs:
        state = WorkerState.from_json_file(os.path.join(dirpath, workdir, "state.json"))
        states.append(state)

    return states


def print_local_workers(dirpath: Optional[str] = None) -> None:
    df = pd.DataFrame(_get_all_local_worker_states(dirpath=dirpath))
    print_dataframe(df, title="\nLocal AbiPy Workers:\n")



_SSH_CONFIG = None


def get_ssh_config():
    global _SSH_CONFIG

    if _SSH_CONFIG is None:
        # For the API, see http://docs.paramiko.org/en/stable/api/config.html#config-module-api-documentation
        from paramiko import SSHConfig
        user_config = os.path.expanduser("~/.ssh/config")
        if not os.path.exists(user_config):
            raise RuntimeError(f"Cannot find ssh configuration file at {user_config}")
        _SSH_CONFIG = SSHConfig.from_path(user_config)
        #print(_SSH_CONFIG)

    return _SSH_CONFIG

    # Return the set of literal hostnames defined in the SSH config
    # (both explicit hostnames and wildcard entries).
    #ssh_config.get_hostnames()
    #host_conf = ssh_config.lookup(self.hostname)


class ClientError(Exception):
    """Exceptions raised by WorkerClient"""

    @classmethod
    def from_response(cls, response):
        """Builds the object from a requests response object."""
        raise cls(f"\nServer returned status code: {response.status_code}\nwith text:\n{response.text}\n")


def need_serving_worker(method):
    """
    Decorator for methods of the WorkerClient that require a running AbipyWorker.
    """

    @functools.wraps(method)
    def decorated(*args, **kwargs):
        self: WorkerClient = args[0]
        if self.worker_state.status != self.worker_state.status.serving:
            raise ClientError(
                f"Status of {self.worker_state.name} worker is `{self.worker_state.status}`\n"
                f"Use `abiw.py restart {self.worker_state.name}` to restart the worker"
            )

        return method(*args, **kwargs)

    return decorated


class WorkerClient(MSONable):
    """
    The client used by the user to connect to an AbipyWorker, send python scripts or directories for execution.
    This object also provides an API to query the status of the Worker.
    """

    def __init__(self, worker_state, is_default_worker=False, timeout=240,
                 is_local_worker=None, ssh_destination=None, local_port=None, saved_uuid=None):

        if not isinstance(worker_state, WorkerState):
            self.worker_state = WorkerState(**worker_state)
        else:
            self.worker_state = worker_state

        self.is_default_worker = bool(is_default_worker)
        self.timeout = timeout

        if is_local_worker is None:
            # Well, this is not perfect as one may have a laptop with the same hostname as the remote cluster.
            self.is_local_worker = self.worker_state.hostname == socket.gethostname()
        else:
            self.is_local_worker = bool(is_local_worker)

        if ssh_destination is None:
            if self.is_local_worker:
                self.ssh_destination = ""
            else:
                hostnames = get_ssh_config().get_hostnames()
                if self.worker_state.hostname not in hostnames:
                    print(f"WARNING: Cannot find `{self.worker_state.hostname}` in `{hostnames}`")

                self.ssh_destination = f"{self.worker_state.username}@{self.worker_state.hostname}"
        else:
            self.ssh_destination = str(ssh_destination)

        if local_port is not None:
            self.local_port = int(local_port)
            self.server_url = f"http://localhost:{self.local_port}"

        else:
            if self.is_local_worker:
                self.local_port = self.worker_state.port
                if self.local_port is None:
                    raise ValueError("If local worker, local_port cannot be None")
                self.server_url = f"http://localhost:{self.local_port}"
            else:
                self.local_port = local_port
                if self.local_port is not None:
                    self.server_url = f"http://localhost:{self.local_port}"
                else:
                    self.server_url = None

        self.saved_uuid = saved_uuid if saved_uuid else self.worker_state.uuid

    @property
    def is_remote_worker(self) -> bool:
        """True if this is a remote worker."""
        return not self.is_local_worker

    #def lworker_dirpath(self) -> bool:
    #    dirname = "worker_{self.worker_state.name}"
    #    return os.path.expanduser(o.path.join("~", ".abinit". "abipy", dirname))

    @pmg_serialize
    def as_dict(self) -> dict:
        return dict(
                worker_state=self.worker_state.dict(),
                is_default_worker=self.is_default_worker,
                timeout=self.timeout,
                is_local_worker=self.is_local_worker,
                ssh_destination=self.ssh_destination,
                local_port=self.local_port,
                saved_uuid=self.saved_uuid,
        )

    #def from_dict(self, d):
    #    return self(**d)

    def setup_ssh_port_forwarding(self) -> None:
        # TODO Rationalize this part. Make sure hostname is listed in ~/.ssh/config.
        info = []

        # Return immediately if port forwarding is not needed.
        if self.is_local_worker or self.worker_state.status != self.worker_state.status.serving:
            return

        print("in setup_ssh_port_forwarding")
        # Need ssh port-forwarding.
        if self.local_port is None:
            self.local_port = find_free_port()
        else:
            if not port_is_open(self.local_port):
                self.local_port = find_free_port()
            else:
                info.append("port self.local_port is already opened")
                return

        ssh_cmd = f"ssh -N -f -L localhost:{self.local_port}:localhost:{self.worker_state.port} {self.ssh_destination}"
        print(f"executing: `{ssh_cmd}`")
        subprocess.run(ssh_cmd, shell=True)

        #return info

    def update_state(self, worker_state: WorkerState) -> None:
        """
        Update the state of the Worker stored by the client.
        """
        prev_hostname = self.worker_state.hostname
        self.worker_state = worker_state.copy()
        # This is needed to preserve the changes done by the user in clients.json
        # if the hostname differs from the one used in the ~/.ssh/config file.
        self.worker_state.hostname = prev_hostname

        if self.saved_uuid != worker_state.uuid:
            print("Detected change in worker_state.uuid!")
            self.local_port = None
            if self.is_default_worker:
                self.local_port = worker_state.port

            self.saved_uuid = worker_state.uuid

    def __repr__(self) -> str:
        return repr(self.worker_state)

    def __str__(self) -> str:
        return self.to_json()

    def to_json(self) -> str:
        """
        Returns a string with the json representation of the client
        """
        # TODO: monty.to_json should accept **kwargs
        return json.dumps(self, cls=AbipyEncoder, indent=2)

    def check_server_url(self) -> None:
        if self.server_url is None:
            raise ClientError("server_url cannot be None.")
        if self.local_port is None:
            raise ClientError("localport cannot be None.")
        if not port_is_open(self.local_port):
            raise ClientError("local port {local_port} is not open.")

    def _get_url(self, end_point: str, action: str) -> str:
        url = f"{self.server_url}/{end_point}"
        print(f"{action} to {self.worker_state.name} worker at url: {url} ...\n")
        self.check_server_url()
        return url

    ##################################################
    # HTTP interface implemented via requests package
    ##################################################

    @need_serving_worker
    def send_pyscript(self, filepath, user_message="", end_point="post_flow_script"):
        """
        Send a python script `filepath` to the AbipyWorker in order to build a new Flow.
        Returns the json-encoded content of the response, if any.
        """
        filepath = os.path.expanduser(filepath)
        with open(filepath, "rt") as fp:
            data = dict(
                pyscript_basename=os.path.basename(filepath),
                pyscript_string=fp.read(),
                user_message=user_message,
            )

        url = self._get_url(end_point, f"Sending `{filepath}` script")
        r = requests.post(url=url, json=data, timeout=self.timeout)
        if r.status_code == 200:
            return r.json()

        raise ClientError.from_response(r)

    @need_serving_worker
    def send_flow_dir(self, flow_dir: str, user_message: str = "", end_point="post_flow_dirs"):
        """
        Send a directory containing a Flow to the AbipyWorker.
        Returns the json-encoded content of the response, if any.

        Args:
            flow_dir: Directory with the Flow (already built by the user)
            user_message:
            end_point:
        """
        if not self.is_local_worker:
            raise RuntimeError(f"You cannot add a Flow to the remote worker: {self.worker_state.name}")

        if not os.path.isdir(flow_dir):
            raise RuntimeError(f"Cannot find flow directory: {flow_dir}")

        url = self._get_url(end_point, f"Sending `{flow_dir}` directory")
        data = dict(flow_dir=flow_dir, user_message=user_message)

        r = requests.post(url=url, json=data, timeout=self.timeout)
        if r.status_code == 200:
            return r.json()

        raise ClientError.from_response(r)

    @need_serving_worker
    def send_kill_message(self, end_point="action") -> str:
        """
        Ask the server to terminate execution.
        """
        url = self._get_url(end_point, "Sending kill message")
        r = requests.post(url=url, json=dict(action="kill"), timeout=self.timeout)
        if r.status_code == 200:
            return r.json()

        raise ClientError.from_response(r)

    @need_serving_worker
    def get_json_status(self, end_point="json_status") -> str:
        """
        Inquire server status.
        """
        url = self._get_url(end_point, "Sending status request")
        r = requests.get(url=url, timeout=self.timeout)

        if r.status_code == 200:
            return r.json()
            #print(r.text)
            #json_status = r.json()
            #from pandas.io.json import read_json
            #json_status["dataframe"] = read_json(json_status["dataframe"])
            #print_dataframe(json_status["dataframe"], title="\nWorker Status:\n")
            #return json_status

        raise ClientError.from_response(r)

    @need_serving_worker
    def open_webgui(self):
        """
        Open a new tab in the browser with a panel dashboard to interact with the Worker.
        """
        end_point = ""
        url = self._get_url(end_point, f"Contacting {self.worker_state.name} worker")
        #print(self)
        import webbrowser
        webbrowser.open_new_tab(url)


class WorkerClients(list, MSONable):
    """
    A list of Clients.
    """

    @classmethod
    def lscan(cls, dirpath: Optional[str] = None) -> WorkerClients:
        """
        Discover all the AbipyWorkers created on the same machine by inspecting
        the configuration directories in the AbiPy directory.

        Return: list of clients associated to local AbiPy workers.
        """
        worker_states = _get_all_local_worker_states(dirpath)
        all_clients = cls.from_json_file(empty_if_not_file=True)
        all_clients = all_clients.remove_stale_local_clients()
        all_clients.update_worker_states(worker_states, write_json=True)

        return cls(c for c in all_clients if c.is_local_worker)

    @classmethod
    def from_json_file(cls, empty_if_not_file: bool = False, filepath: Optional[str] = None) -> WorkerClients:
        """
        Build object from JSON file.

        Args:
            empty_if_not_file:
            filepath:
        """
        filepath = os.path.expanduser(filepath) if filepath else os.path.join(ABIPY_DIRPATH, "clients.json")

        new = cls()
        if empty_if_not_file and not os.path.exists(filepath):
            return new

        if not os.path.exists(filepath):
            raise RuntimeError(f"""
Could not find `{filepath}`
Please execute:

    abiw.py lscan

or

    abiw.py rscan

to create it""")

        with open(filepath, "rt") as fp:
            data = json.load(fp)
            for d in data["clients"]:
                client = WorkerClient.from_dict(d)
                new.append(client)

        new._validate()
        return new

    @classmethod
    def rscan(cls, hostnames: list[str], printout: bool = True) -> WorkerClients:
        """
        Uses fabric to discover AbipyWorkers running on remote machines.

        Args:
            hostnames: List of hostnames as declared in ~/.ssh/config.
            printout:
        """
        #
        # From https://docs.fabfile.org/en/2.6/api/transfer.html#fabric.transfer.Transfer
        #
        # Most SFTP servers set the remote working directory to the connecting user’s home directory,
        # and (unlike most shells) do not expand tildes (~).
        #
        # For example, instead of saying get("~/tmp/archive.tgz"), say get("tmp/archive.tgz").

        try:
            from fabric import Connection
        except ImportError:
            raise ImportError("rscan requires fabric package. Install it with `pip install fabric`.")

        from io import BytesIO
        print("Make sure the hosts are listed in your ~/.ssh/config file and that ProxyJump is activated if needed")
        worker_states = []
        for host in tqdm(hostnames):
            print(f"Contacting {host}. It may take some time ...")
            try:
                c = Connection(host)
                result = c.run("ls ~/.abinit/abipy")
                files = result.stdout.split()
                worker_dirs = [f for f in files if f.startswith("worker_")]
                worker_dirs = [os.path.join(".abinit/abipy", d) for d in worker_dirs]
            except Exception as exc:
                print(exc)
                continue

            # Now get the state.json via sftp.
            for w in worker_dirs:
                state_path = os.path.join(w, "state.json")
                try:
                    strio = BytesIO()
                    c.get(state_path, local=strio)
                    json_bstring = strio.getvalue()
                    strio.close()
                    worker_states.append(WorkerState.from_json(json_bstring))

                except IOError as exc:
                    print(f"Cannot find `state.json` file: {host}@{state_path}. Ignoring exception.")
                    print(exc)

        all_clients = cls.from_json_file(empty_if_not_file=True)
        all_clients = all_clients.remove_stale_local_clients()
        all_clients.update_worker_states(worker_states, write_json=False)

        #for client in all_clients:
        #    client.setup_ssh_port_forwarding()

        all_clients.write_json_file()

        return cls(c for c in all_clients if c.is_remote_worker)

    def _validate(self) -> None:
        err_lines = []
        app = err_lines.append
        from collections import Counter

        # worker_name must be unique.
        for worker_name, count in Counter((w.worker_state.name for w in self)).items():
            if count > 1:
                app(f"Server name: `{worker_name}` appears: {count} times. This is forbidden!")

        # Only zero or one default server is allowed.
        count = sum(1 if w.is_default_worker == True else 0 for w in self)
        if count not in (0, 1):
            app(f"is_default_worker=True appears `{count}` times." +
                "This is forbidden as only one default worker is allowed!")

        #self.check_local_port()

        for client in self:
            #print(client.local_port, "\n", client)
            if duck.is_intlike(client.local_port):
                address = "localhost"
                if not port_is_open(client.local_port, address=address):
                    #print(f"WARNING: Nobody is listening to `{address}:{client.local_port}`")
                    client.local_port = None
                else:
                    print(f"{client.worker_state.name} listening to `{address}:{client.local_port}`")

        if err_lines:
            raise RuntimeError("\n".join(err_lines))

    def __str__(self):
        lines = [str(w) for w in self]
        return (2 * "\n").join(lines)

    #@pmg_serialize
    def as_dict(self) -> dict:
        return {"clients": [w.as_dict() for w in self]}

    #def dict(self) -> dict:
    #    return {"clients": [w.dict() for w in self]}

    def write_json_file(self, filepath: Optional[str] = None) -> None:
        """
        Write json file with the list of clients
        """
        filepath = os.path.join(ABIPY_DIRPATH, "clients.json") if not filepath else os.path.expanduser(filepath)

        # TODO: Compare old with new. return diff
        #old_clients = self.__class__.from_json_file(empty_if_not_file=True)
        #if old_clients is None:

        with open(filepath, "wt") as fp:
            json.dump(self.as_dict(), fp, indent=2)

    def update_worker_states(self, worker_states: list[WorkerState], write_json: bool = True) -> None:
        """
        Update clients from worker_states
        """
        for state in worker_states:
            client = self.select_from_worker_name(state.name, allow_none=True)
            if client is not None:
                client.update_state(state)
            else:
                self.append(WorkerClient(state))

        self._validate()
        if write_json: self.write_json_file()

    def remove_stale_local_clients(self) -> WorkerClients:
        """
        Return new list of clients where all entries associated to workers whose
        directory is not found have been removed.
        """
        new_clients = []
        for client in self:
            worker_dirpath = os.path.join(ABIPY_DIRPATH, f"worker_{client.worker_state.name}")
            if client.is_local_worker and not os.path.exists(worker_dirpath):
                print("Removing stale section with worker:", client.worker_state.name)
            else:
                new_clients.append(client)

        return self.__class__(new_clients)

    def get_dataframe(self) -> pd.DataFrame:
        """
        Build and return pandas Dataframe with the AbipyWorker states.
        """
        rows = []
        for client in self:
            d = client.as_dict()
            d = {k: v for k, v in d.items() if not k.startswith("@")}
            state = d.pop("worker_state", )
            if "status" in state:
                state["status"] = state["status"].value  #.colored_str()

            mng_connector_dict = state.pop("mng_connector", None)
            if mng_connector_dict is not None:
                # Add the most important info from the mng_connector.
                # Prepend some keys with `mongo_prefix`.
                remap = {"host": "mongo_host", "dbname": "mongo_dbname", "collection_name": None}
                mongo_info = {k: mng_connector_dict[k] for k in remap}
                for k, new_key in remap.items():
                    if new_key is None: continue
                    mongo_info[new_key] = mongo_info.pop(k)

                state.update(**mongo_info)

            d.update(**state)
            rows.append(d)

        df = pd.DataFrame(rows)
        return df

    def print_dataframe(self) -> None:
        #if select_keys == "short":
        #    selected_keys = ["name", "status", "pid", "scratch_dir", "username", "hostname", ]

        #all_keys = ["is_default_worker", "timeout", "is_local_worker", "ssh_destination",
        #            "local_port", "name",  "status", "pid", "address",  "port",  scratch_dir",
        #            "username", "hostname" "uuid version"
        #        ]

        df = self.get_dataframe().set_index("name")
        print_dataframe(df, title="\nAbiPyWorker Clients:\n")

    def summarize(self) -> pd.DataFrame:
        """
        Return dataframe with the most important columns.
        """
        select_keys = ["name", "status", "pid", "scratch_dir"]
        if any(c.is_remote_worker for c in self):
            select_keys += ["username", "hostname"]
        df = self.get_dataframe()

        if df.empty: return df
        return df[select_keys].set_index("name")

    def get_all_worker_names(self) -> list[str]:
        """Return the list of worker names stored in self."""
        return [client.worker_state.name for client in self]

    def select_from_worker_name(self, worker_name: str,
                                allow_none=False) -> Union[WorkerClient, None]:
        """
        Select and return a WorkerClient from its name `worker_name`.

        Return: None if `worker_name` is not found and `allow_none` is True else raises ValueError.
        """
        if worker_name is None:
            # If there are n > 1 workers, a default must be set by the user.
            for client in self:
                if client.is_default_worker: return client

            if len(self) == 1: return self[0]

            raise ValueError("Cannot find default worker in list:\n" +
                             f"\t{self.get_all_worker_names()}\n" +
                             "Use:\n\tabiw.py set_default WORKER_NAME")

        for client in self:
            if client.worker_state.name == worker_name: return client

        if allow_none: return None
        raise ValueError(f"Cannot find worker with name `{worker_name}`. Choose among:\n" +
                         f"\t{self.get_all_worker_names()}\n")

    def set_default(self, worker_name: str) -> None:
        """
        Set the Client with `worker_name` as the default one.
        """
        the_one = self.select_from_worker_name(worker_name, allow_none=False)
        for client in self:
            client.is_default_worker = False
        the_one.is_default_worker = True

        self._validate()
        self.write_json_file()





import subprocess
import random
import tempfile

class SSHTunnel:
    """
    https://stackoverflow.com/questions/4364355/how-to-open-an-ssh-tunnel-using-python

    And here's how you could use it, for example with MySQL (port 3306 usually):

    with SSHTunnel('database.server.com', 'you', '22', '/path/to/private_key', '3306') as tunnel:
        print("Connected on port {} at {}".format(tunnel.local_port, tunnel.local_host))
    """

    def __init__(self, host, user, port, key, remote_port):
        self.host = host
        self.user = user
        self.port = port
        self.key = key
        self.remote_port = remote_port
        # Get a temporary file name
        tmpfile = tempfile.NamedTemporaryFile()
        tmpfile.close()
        self.socket = tmpfile.name
        self.local_port = random.randint(10000, 65535)
        self.local_host = '127.0.0.1'
        self.open = False

    def start(self):
        # I had to use -o ExitOnForwardFailure=yes instead of -o ExitOnForwardFailure=True.
        exit_status = subprocess.call(['ssh', '-MfN',
            '-S', self.socket,
            '-i', self.key,
            '-p', self.port,
            '-l', self.user,
            '-L', '{}:{}:{}'.format(self.local_port, self.local_host, self.remote_port),
            '-o', 'ExitOnForwardFailure=yes',
            self.host
        ])
        if exit_status != 0:
            raise Exception('SSH tunnel failed with status: {}'.format(exit_status))
        if self.send_control_command('check') != 0:
            raise Exception('SSH tunnel failed to check')
        self.open = True

    def stop(self):
        if self.open:
            if self.send_control_command('exit') != 0:
                raise Exception('SSH tunnel failed to exit')
            self.open = False

    def send_control_command(self, cmd):
        return subprocess.check_call(['ssh', '-S', self.socket, '-O', cmd, '-l', self.user, self.host])

    def __enter__(self):
        self.start()
        return self

    def __exit__(self, type, value, traceback):
        self.stop()




