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

# Copyright (c) Mito.
# Distributed under the terms of the Modified BSD License.

"""
TODO: Add module docstring
"""
from mitosheet.api import handle_api_event
from mitosheet.mito_analytics import log
import sys
import pandas as pd
from typing import List
import traceback

from ipywidgets import DOMWidget
import traitlets as t

from ._frontend import module_name, module_version
from .errors import EditError, get_recent_traceback_as_list, log_recent_error
from .save_utils import (
    read_and_upgrade_analysis, 
    write_analysis,
    saved_analysis_names_json
)
from mitosheet.widget_state_container import WidgetStateContainer
from mitosheet.profiling import timeit

from mitosheet.mito_analytics import analytics, static_user_id

class MitoWidget(DOMWidget):
    """TODO: Add docstring here
    """
    _model_name = t.Unicode('ExampleModel').tag(sync=True)
    _model_module = t.Unicode(module_name).tag(sync=True)
    _model_module_version = t.Unicode(module_version).tag(sync=True)
    _view_name = t.Unicode('ExampleView').tag(sync=True)
    _view_module = t.Unicode(module_name).tag(sync=True)
    _view_module_version = t.Unicode(module_version).tag(sync=True)

    value = t.Unicode('Hello World').tag(sync=True)
    analysis_name = t.Unicode('').tag(sync=True)
    sheet_json = t.Unicode('').tag(sync=True)
    code_json = t.Unicode('').tag(sync=True)
    df_names_json = t.Unicode('').tag(sync=True)
    df_shape_json = t.Unicode('').tag(sync=True)
    saved_analysis_names_json = t.Unicode('').tag(sync=True)
    user_id = t.Unicode(static_user_id).tag(sync=True)
    column_spreadsheet_code_json = t.Unicode('').tag(sync=True)
    column_filters_json = t.Unicode('').tag(sync=True)
    column_type_json = t.Unicode('').tag(sync=True)


    def __init__(self, *args, **kwargs):
        """
        Takes a list of dataframes, passed through *args. 

        NOTE: assumes that the passed arguments are all dataframes, 
        and also have all valid headers. These conditions are checked in the
        mitosheet.sheet function. 
        """
        # Call the DOMWidget constructor to set up the widget properly
        super(MitoWidget, self).__init__()

        # Set up the state container to hold private widget state
        self.widget_state_container = WidgetStateContainer(args)

        # Set up starting shared state variables
        self.update_shared_state_variables()

        # Set up message handler
        self.on_msg(self.receive_message)

    def update_shared_state_variables(self):
        """
        Helper function for updating all the variables that are shared
        between the backend and the frontend through trailets.
        """
        self.sheet_json = self.widget_state_container.sheet_json
        self.column_spreadsheet_code_json = self.widget_state_container.column_spreadsheet_code_json
        self.code_json = self.widget_state_container.code_json
        self.df_names_json = self.widget_state_container.df_names_json
        self.df_shape_json = self.widget_state_container.df_shape_json
        self.analysis_name = self.widget_state_container.analysis_name
        self.saved_analysis_names_json = saved_analysis_names_json()
        self.column_filters_json = self.widget_state_container.column_filters_json
        self.column_type_json = self.widget_state_container.column_type_json

    @timeit 
    def send(self, msg):
        """
        We overload the DOMWidget's send function, so that 
        we log all outgoing messages
        """
        # Send the message though the DOMWidget's send function
        super().send(msg)

    @timeit
    def saturate(self, event):
        """
        Saturation is when the server fills in information that
        is missing from the event client-side. This is for consistency
        and because the client may not have easy access to the info
        all the time.
        """
        curr_step = self.widget_state_container.curr_step

        if event['event'] == 'edit_event':
            if event['type'] == 'set_column_formula_edit':
                sheet_index = event['sheet_index']
                column_header = event['column_header']
                event['old_formula'] = curr_step['column_spreadsheet_code'][sheet_index][column_header]
        if event['event'] == 'update_event':
            if event['type'] == 'use_existing_analysis_update':
                # If we're getting an event telling us to update, we read in the steps from the file
                analysis_name = event['analysis_name']
                analysis = read_and_upgrade_analysis(analysis_name)
                for key, value in analysis.items():
                    event[key] = value

                # When replaying an analysis with import events, you can also send over
                # new params to the import events to replace them. We replace them in the steps here
                if 'import_summaries' in event:
                    for step_id, params in event['import_summaries'].items():
                        for key, value in params.items():
                            event['steps'][step_id][key] = value
                
        if event['event'] == 'api_call':
            handle_api_event(self.send, event)

    @timeit
    def handle_edit_event(self, event):
        """
        Handles an edit_event. Per the spec, an edit_event
        updates both the sheet and the codeblock, and as such
        the sheet is re-evaluated and the code for the codeblock
        is re-transpiled.

        Useful for any event that changes the state of both the sheet
        and the codeblock!
        """
        # First, we log the edit that we just got
        log(event['type'], event)

        # First, we send this new edit to the evaluator
        self.widget_state_container.handle_edit_event(event)

        # We update the state variables 
        self.update_shared_state_variables()

        # Also, write the analysis to a file!
        write_analysis(self.widget_state_container)

        # Tell the front-end to render the new sheet and new code
        self.send({"event": "update_sheet"})
        self.send({"event": "update_code"})


    @timeit
    def handle_update_event(self, event):
        """
        This event is not the user editing the sheet, but rather information
        that has been collected from the frontend (after render) that is being
        passed back.

        For example:
        - Names of the dataframes
        - Name of an existing analysis
        """
        # First, we log the update event that we just got
        log(event['type'], event)

        self.widget_state_container.handle_update_event(event)

        # Update all state variables and send out an update
        self.update_shared_state_variables()

        # Also, write the analysis to a file!
        write_analysis(self.widget_state_container)

        # Tell the front-end to render the new sheet and new code
        self.send({"event": "update_sheet"})
        self.send({"event": "update_code"})

    def receive_message(self, widget, content, buffers=None):
        """
        Handles all incoming messages from the JS widget. There are two main
        types of events:

        1. edit_event: any event that updates the state of the sheet and the
        code block at once. Leads to reevaluation, and a re-transpile.

        2. update_event: any event that isn't caused by an edit, but instead
        other types of new data coming from the frontend (e.g. the df names 
        or some existing steps).

        3. A log_event is just an event that should get logged on the backend.
        """
        # First, we saturate the event
        event = content
        self.saturate(event)

        try:
            if event['event'] == 'edit_event':
                self.handle_edit_event(event)
            if event['event'] == 'update_event':
                self.handle_update_event(event)
            if event['event'] == 'log_event':
                log( 
                    event['type'], 
                    event
                )
        except EditError as e:
            print(e)
            # If we hit an error during editing, log that it has occured
            log(
                f'{e.type_}_log_event', 
                {
                    'header': e.header, 
                    'to_fix': e.to_fix,
                    'traceback': get_recent_traceback_as_list()
                }
            )
            # Report it to the user, and then return
            self.send({
                'event': 'edit_error',
                'type': e.type_,
                'header': e.header,
                'to_fix': e.to_fix
            })
        except:
            # We log the error
            log_recent_error()
            # Report it to the user, and then return
            self.send({
                'event': 'edit_error',
                'type': 'execution_error',
                'header': 'Execution Error',
                'to_fix': 'Sorry, there was an error during executing this code.'
            })

def sheet(*args: List[pd.DataFrame]) -> MitoWidget:
    """
    Renders a Mito sheet. If no arguments are passed, renders an empty sheet. Otherwise, renders
    any dataframes that are passed. Errors if any given arguments are not dataframes.

    If running this function just prints text that looks like `MitoWidget(...`, then you need to 
    install the JupyterLab extension manager by running:

    jupyter labextension install @jupyter-widgets/jupyterlab-manager;

    Run this command in the terminal where you installed Mito. It should take 5-10 minutes to complete.

    Then, restart your JupyterLab instance, and refresh your browser. Mito should now render.

    NOTE: if you have any issues with installation, please book a demo at https://calendly.com/trymito/mito-early-access-demo?utm_source=sheetdocstring
    """
    args = list(args)

    for df in args:
        if not isinstance(df, pd.DataFrame):
            # Log the error
            log(
                'sheet_error_log_event', 
                {'error': f'Invalid argument passed to sheet: {df}. Please pass all dataframes.'}
            )

            raise ValueError(f'Invalid argument passed to sheet: {df}. Please pass all dataframes.')

    try:
        widget = MitoWidget(*args) 
    except:
        # We log the error
        log_recent_error('sheet_error_log_event')
        # And then bubble it to the user
        raise

    return widget
