#!/usr/bin/env python3

import subprocess
import logging
import shutil
import osrm
import time
import sys
import sh

from contextlib import AbstractContextManager
from urllib.error import HTTPError
from sh import CommandNotFound
from functools import partial
from pathlib import Path

from .RoutingProfile import RoutingProfile
from .OSMDataset import OSMDataset
from .OSRM import OSRM
from . import defaults
from . import utils

class Scenario():
    """
    Scenario is an abstraction of OSRM executables in order to compile, serve, and query
    a routable network.

    A Scenario is initialized with (1) an OSM Dataset and (2) a Routing Profile. When called,
    Scenario compiles the OSMDataset and RoutingProfile into an OSRM routable network,
    provides a context manager to serve the scenario's HTTP API, and returns various methods
    to query that API.

    The HTTP API methods (`match`, `nearest`, `simple_route`, `table`, `trip`) are provided by
    the `python-osrm` module (https://github.com/ustroetz/python-osrm)

    Example
    -------
    >>> from tebetebe.profiles import foot
    >>> import tebetebe as tb
    >>>
    >>> ## initialize scenario with GeoFabrik extract & default foot profile
    >>> scenario = tb.Scenario("./swaziland-latest.osm.pbf", foot)
    >>>
    >>> with scenario() as api: ## compile scenario and serve HTTP API
    >>>     api.simple_route()... ## query HTTP API for a simple route, match, nearest, trip...

    Parameters
    ----------
    osm_dataset : str / OSMDataset
        OSM dataset from which the route network will be extracted
    routing_profile : str / RoutingProfile
        Routing profile to be used in scenario
    name : str, optional
        Scenario name used in output file and log entries. If not supplied, it will be built
        based upon the OSMDataset and RoutingProfile names
    algorithm : str, optional
        Algorithm to be used (either "CH" or "MLD")
    tmp_dir : str, optional
        Temporary directory to store files generated by osrm binaries
    overwrite: bool, optional
        Overwrite scenario if .osrm* files already exist. Otherwise, the existing .osrm* files
        will be used.
    verbose: bool, optional
        Print output of OSRM compilation
    **kwargs
        To pass any custom arguments to the OSRM executables ("osrm-routed", "osrm-contract", ...)
        pass the kwarg "{executable}_args" (ex. "routed_args", "contract_args", ...) with a dictionary
        of key:values to be passed
    """

    def __init__(self, osm_dataset, routing_profile,
                 name=None, algorithm="MLD", 
                 tmp_dir=defaults.TMP_DIR,
                 overwrite=defaults.OVERWRITE,
                 verbose=defaults.VERBOSE,
                 **kwargs):

        self.log = logging.getLogger(defaults.LOGGER)
        self.algorithm = algorithm.upper()
        self.tmp_dir = Path(tmp_dir)
        self.overwrite = overwrite

        ## Handle if params are passed as string
        if isinstance(osm_dataset, str):
            osm_dataset = OSMDataset(osm_dataset, tmp_dir=self.tmp_dir)

        if isinstance(routing_profile, str):
            routing_profile = RoutingProfile(routing_profile)

        self.osm_dataset = osm_dataset
        self.routing_profile = routing_profile

        ## Initialize / check OSRM Binaries
        self.OSRM = OSRM(verbose=verbose)

        ## Scenario name and output osrm path
        self.name = name if name else "{}_{}".format(self.osm_dataset.get_name(),
                                                     self.routing_profile.get_name())
        self.path = self.tmp_dir / self.name / "{}.osrm".format(self.name)
        self.path.parent.mkdir(parents=True, exist_ok=True) # Make sure output_dir exists

        ## Add osrm executable args if they are passed
        self.args = {}
        for k_name, kwarg in kwargs.items():
            if k_name[-5:] == "_args":
                self.args[k_name[:-5]] = kwarg

    ##
    ## Scenario compilation
    ##

    def __call__(self):
        '''
        Compile OSMDataset and RoutingProfile into an OSRM routable network (.osrm file)
        '''

        ## Honor overwrite settings
        if self.path.is_file():
            if self.overwrite:
                self.log.info("{}: Overwriting scenario".format(self.name))

                ## Remove all .osrm* files
                for osrm_file in self.path.parent.glob("{}.osrm*".format(self.name)):
                    osrm_file.unlink()
            else:
                self.log.info("{}: Using existing scenario".format(self.name))
                return self

        try:
            self.log.info("{}: Compiling scenario ({})".format(self.name, self.algorithm))

            if self.routing_profile.is_default():
                self.log.warning("{}: Default {} profile may not be accurate for your use case".format(self.name,
                                                                                                       self.routing_profile.get_name()))

            ## Run osrm-extract
            extr = self.OSRM.extract(self.osm_dataset.get_path(),
                                     self.routing_profile.get_path(),
                                     **self._get_command_args("extract"))
            extr.wait() ## synchronous execution

            ## osrm-extract generates osrm file in same folder as OSMDataset
            ## so rename and move all files to self.path
            network_path = self.osm_dataset.get_path()

            for src_path in network_path.parent.glob("{}.osrm*".format(network_path.stem.split(".")[0])):
                dst_file = src_path.name.replace(src_path.stem.split(".")[0], self.name)
                dst_path = self.path.parent / dst_file

                shutil.move(src_path, dst_path)

            ## Run different binaries depending on algorithm used
            if self.algorithm == "MLD":
                part = self.OSRM.partition(self.path, **self._get_command_args("partition"))
                part.wait()

                cust = self.OSRM.customize(self.path, **self._get_command_args("customize"))
                cust.wait()

            elif self.algorithm == "CH":
              cont = self.OSRM.contract(self.path, **self._get_command_args("contract"))
              cont.wait()

            return self

        except Exception as exc:
            ## TODO custom handle more common exception
            self.log.error(exc)

    ##
    ## Context Manager for running Scenario HTTP server (osrm-routed)
    ##

    def __enter__(self):
        ## Compile the Scenario if it hasn't yet
        if not self.path.is_file():
            self()

        ## Point the osrm http api to correct port
        port = utils.find_open_port()
        self._set_config(port)

        self._set_command_args("routed", {
            "port": port,
            "algorithm": self.algorithm,
            "dataset_name": self.name
        })

        self.ready = False

        ## Start up OSRM HTTP server with OSRM-routed
        try:
            ## Callback to execute when OSRM process is ready to receive http requests
            def ready_callback(process):
                self.ready = True

            ## Callback when command exits
            def done_callback(cmd, success, exit_code):
                ## -9 = SIGKILL = it was (probably) killed by this script. aka, this is normal.
                if exit_code != -9:
                    self.log.error("{}: Scenario process exited with exit code {}".format(self.name, exit_code))
                    self.log.error("{}: Command: {}".format(self.name, cmd))


            self.log.info("{}: Initializing scenario".format(self.name))
            self.process = self.OSRM.routed(self.path,
                                            ready_callback,
                                            done_callback,
                                            **self._get_command_args("routed"))

        except Exception as exc:
            self.log.error(exc)

        ## Wait until OSRM-routed is ready before returning the api
        ## TODO better way to do this?
        while not self.ready:
            time.sleep(1)

        self.log.info("{}: Ready for requests".format(self.name))

        return self

    def __exit__(self, exc_type, exc_value, traceback):
        ## TODO handle common exceptions
        if exc_type:
            if exc_type is HTTPError: ## Log but not exit on HTTP Errors
                self.log.error(exc_value)
                return True
            else:
                self.log.error("{} {} {}".format(exc_type, exc_value, traceback))

        ## Kill the OSRM process
        try:
            self.process.kill()
        except:
            pass

    ##
    ## osrm-routed HTTP API

    simple_route = osrm.simple_route
    nearest = osrm.nearest
    match = osrm.match
    table = osrm.table
    trip = osrm.trip

    def _set_config(self, port):
        """Set up HTTP API configuration given the HTTP server port"""
        self.config = osrm.RequestConfig("127.0.0.1:{}/v1/skobuffs".format(port))

        self.simple_route = partial(osrm.simple_route, url_config=self.config)
        self.nearest = partial(osrm.nearest, url_config=self.config)
        self.match = partial(osrm.match, url_config=self.config)
        self.table = partial(osrm.table, url_config=self.config)
        self.trip = partial(osrm.trip, url_config=self.config)

    ##
    ## Other

    def _get_command_args(self, command):
        return self.args[command] if command in self.args else {}

    def _set_command_args(self, command, args):
        self.args[command] = self.args[command] if command in self.args else {}
        self.args[command] = {**self.args[command], **args}

    def get_name(self):
        """Return the Scenario name"""
        return self.name

    def get_path(self):
        """Return the Scenario path"""
        return self.path

    def is_alive(self):
        """True if Scenario process is running, False otherwise"""
        if hasattr(self, "process") and hasattr(self.process, "process"):
            return self.process.process.is_alive()
        else:
            return False



