import io
import logging
import networkx as nx
import os
import threading
import tornado.ioloop
import tornado.web
import tornado.websocket
import tornado.escape
import tornado.gen
import yaml
import webbrowser
from contextlib import contextmanager
from time import sleep

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

class VisualizationElement:
    """
    Defines an element of the visualization.
    Attributes:
        package_includes: A list of external JavaScript files to include that
                          are part of the packages.
        local_includes: A list of JavaScript files that are local to the
                        directory that the server is being run in.
        js_code: A JavaScript code string to instantiate the element.
    Methods:
        render: Takes a model object, and produces JSON data which can be sent
                to the client.
    """

    package_includes = []
    local_includes = []
    js_code = ''
    render_args = {}

    def __init__(self):
        pass

    def render(self, model):
        return '<b>VisualizationElement goes here</b>.'


class PageHandler(tornado.web.RequestHandler):
    """ Handler for the HTML template which holds the visualization. """

    def get(self):
        self.render('index.html', port=self.application.port,
                    model_name=self.application.model_name,
                    package_includes=self.application.package_includes,
                    local_includes=self.application.local_includes,
                    scripts=self.application.js_code)


class SocketHandler(tornado.websocket.WebSocketHandler):
    """ Handler for websocket. """

    def open(self):
        if self.application.verbose:
            logger.info('Socket opened!')

    def check_origin(self, origin):
        return True

    def on_message(self, message):
        """ Receiving a message from the websocket, parse, and act accordingly. """

        msg = tornado.escape.json_decode(message)

        if msg['type'] == 'config_file':

            if self.application.verbose:
                print(msg['data'])

            self.config = list(yaml.load_all(msg['data']))

            if len(self.config) > 1:
                error = 'Please, provide only one configuration.'
                if self.application.verbose:
                    logger.error(error)
                self.write_message({'type': 'error',
                    'error': error})
                return
            else:
                self.config = self.config[0]
                self.send_log('INFO.soil', 'Using config: {name}'.format(name=self.config['name']))

            if 'visualization_params' in self.config:
                self.write_message({'type': 'visualization_params',
                    'data': self.config['visualization_params']})
            self.name = self.config['name']
            self.run_simulation()

            settings = []
            for key in self.config['environment_params']: 
                if type(self.config['environment_params'][key]) == float or type(self.config['environment_params'][key]) == int:
                    if self.config['environment_params'][key] <= 1:
                        setting_type = 'number'
                    else:
                        setting_type = 'great_number'
                elif type(self.config['environment_params'][key]) == bool:
                    setting_type = 'boolean'
                else:
                    setting_type = 'undefined'

                settings.append({
                    'label': key,
                    'type': setting_type,
                    'value': self.config['environment_params'][key]
                })

            self.write_message({'type': 'settings',
                'data': settings})

        elif msg['type'] == 'get_trial':
            if self.application.verbose:
                logger.info('Trial {} requested!'.format(msg['data']))
            self.send_log('INFO.user', 'Trial {} requested!'.format(msg['data']))
            self.write_message({'type': 'get_trial',
                'data': self.get_trial( int(msg['data'] )) })

        elif msg['type'] == 'run_simulation':
            if self.application.verbose:
                logger.info('Running new simulation for {name}'.format(name=self.config['name']))
            self.send_log('INFO.soil', 'Running new simulation for {name}'.format(name=self.config['name']))
            self.config['environment_params'] = msg['data']
            self.run_simulation()

        else:
            if self.application.verbose:
                logger.info('Unexpected message!')

    def update_logging(self):
        try:
            if (not self.log_capture_string.closed and self.log_capture_string.getvalue()):
                for i in range(len(self.log_capture_string.getvalue().split('\n')) - 1):
                    self.send_log('INFO.soil', self.log_capture_string.getvalue().split('\n')[i])
                self.log_capture_string.truncate(0)
                self.log_capture_string.seek(0)
        finally:
            if self.capture_logging:
                thread = threading.Timer(0.01, self.update_logging)
                thread.start()

    def on_close(self):
        if self.application.verbose:
            logger.info('Socket closed!')

    def send_log(self, logger, logging):
        self.write_message({'type': 'log',
            'logger': logger,
            'logging': logging })

    def run_simulation(self):
        # Run simulation and capture logs
        if 'visualization_params' in self.config:
            del self.config['visualization_params']
        with self.logging(self.application.model.name):
            try:
                self.simulation = self.application.model.run(self.config)
            except:
                error = 'Something went wrong. Please, try again.'
                self.write_message({'type': 'error',
                    'error': error})

        trials = []
        for i in range(self.config['num_trials']):
            trials.append('{}_trial_{}'.format(self.name, i))
        self.write_message({'type': 'trials',
            'data': trials })

    def get_trial(self, trial):
        G = self.simulation[trial].history_to_graph()
        return nx.node_link_data(G)

    @contextmanager
    def logging(self, logger):
        self.capture_logging = True
        self.logger_application = logging.getLogger(logger)
        self.log_capture_string = io.StringIO()
        ch = logging.StreamHandler(self.log_capture_string)
        self.logger_application.addHandler(ch)
        self.update_logging()
        yield self.capture_logging

        sleep(0.2)
        self.log_capture_string.close()
        self.logger_application.removeHandler(ch)
        self.capture_logging = False
        return self.capture_logging
    

class ModularServer(tornado.web.Application):
    """ Main visualization application. """

    portrayal_method = None
    port = 8001
    model_args = ()
    model_kwargs = {}
    page_handler = (r'/', PageHandler)
    socket_handler = (r'/ws', SocketHandler)
    static_handler = (r'/(.*)', tornado.web.StaticFileHandler,
                      {'path': 'templates'})
    local_handler = (r'/local/(.*)', tornado.web.StaticFileHandler,
                     {'path': ''})

    handlers = [page_handler, socket_handler, static_handler, local_handler]
    settings = {'debug': True,
                'template_path': os.path.dirname(__file__) + '/templates'}

    def __init__(self, model, visualization_element, name='SOIL', verbose=True,
                 *args, **kwargs):
        
        self.verbose = verbose
        self.package_includes = set()
        self.local_includes = set()
        self.js_code = []
        
        self.visualization_element = visualization_element

        self.model_name = name
        self.model = model
        self.model_args = args
        self.model_kwargs = kwargs
        #self.reset_model()

        # Initializing the application itself:
        super().__init__(self.handlers, **self.settings)

    '''
    def reset_model(self):
        self.model = self.model_cls(*self.model_args, **self.model_kwargs)
    '''

    def render_model(self):
        return self.visualization_element.render(self.model)

    def launch(self, port=None):
        """ Run the app. """
        
        if port is not None:
            self.port = port
        url = 'http://127.0.0.1:{PORT}'.format(PORT=self.port)
        print('Interface starting at {url}'.format(url=url))
        self.listen(self.port)
        # webbrowser.open(url)
        tornado.ioloop.IOLoop.instance().start()