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

# Copyright (c) Saga Inc.
# Distributed under the terms of the GPL License.
from typing import Any, Dict, List, Optional, Set, Tuple, Union

import numpy as np

from mitosheet.code_chunks.code_chunk import CodeChunk
from mitosheet.code_chunks.step_performers.set_cell_value_code_chunk import \
    SetCellValueCodeChunk
from mitosheet.errors import make_cast_value_to_type_error
from mitosheet.is_type_utils import (is_none_type, is_number_dtype,
                                     is_string_dtype)
from mitosheet.public.v1.sheet_functions.types import \
    get_function_to_convert_to_series
from mitosheet.state import State
from mitosheet.step_performers.step_performer import StepPerformer
from mitosheet.step_performers.utils.utils import get_param
from mitosheet.types import ColumnID, StepType


class SetCellValueStepPerformer(StepPerformer):
    """
    A set_cell_value step, allows you to set the value
    of a given cell in the sheet.

    NOTE: this step performer is no longer generated by the frontend, so we don't worry
    about making it work with new public interface versions.
    """

    @classmethod
    def step_version(cls) -> int:
        return 1

    @classmethod
    def step_type(cls) -> str:
        return 'set_cell_value'

    @classmethod
    def saturate(cls, prev_state: State, params: Dict[str, Any], previous_steps: List[StepType]) -> Dict[str, Any]:
        # Mito doesn't allow empty cells, so if the new value is empty, change it to None.
        if params['new_value'] == '':
            params['new_value'] = None

        # Get the old value so we can check if the new value is different
        sheet_index = params['sheet_index']
        column_id = params['column_id']
        row_index = params['row_index']
        column_header = prev_state.column_ids.get_column_header_by_id(sheet_index, column_id)

        # Cast the old value to a string to avoid errors while writing the saved analysis
        params['old_value'] = str(prev_state.dfs[sheet_index].at[row_index, column_header])
        
        return params

    @classmethod
    def execute(cls, prev_state: State, params: Dict[str, Any]) -> Tuple[State, Optional[Dict[str, Any]]]:
        sheet_index: int = get_param(params, 'sheet_index')
        column_id: ColumnID = get_param(params, 'column_id')
        row_index: int = get_param(params, 'row_index')
        old_value: str = get_param(params, 'old_value')
        new_value: Union[str, None] = get_param(params, 'new_value')

        # If nothings changed, there's no work to do
        if old_value == new_value:
            return prev_state, None

        column_header = prev_state.column_ids.get_column_header_by_id(sheet_index, column_id)

        # Update the value of the cell, we handle it differently depending on the type of the column
        column_dtype = str(prev_state.dfs[sheet_index][column_header].dtype)
        type_corrected_new_value = cast_value_to_type(new_value, column_dtype)

        execution_data = {
            'type_corrected_new_value': type_corrected_new_value,
        }

        return cls.execute_through_transpile(
            prev_state, 
            params, 
            execution_data,
        )

    @classmethod
    def transpile(
        cls,
        prev_state: State,
        params: Dict[str, Any],
        execution_data: Optional[Dict[str, Any]],
    ) -> List[CodeChunk]:

        return [
            SetCellValueCodeChunk(
                prev_state, 
                get_param(params, 'sheet_index'), 
                get_param(params, 'column_id'), 
                get_param(params, 'row_index'), 
                get_param(params, 'old_value'), 
                get_param(params, 'new_value'), 
                execution_data['type_corrected_new_value'] if (execution_data is not None and 'type_corrected_new_value' in execution_data) else None)
        ]


    @classmethod
    def get_modified_dataframe_indexes(cls, params: Dict[str, Any]) -> Set[int]:
        return {get_param(params, 'sheet_index')}


def cast_value_to_type(value: Union[str, None], column_dtype: str) -> Optional[Any]:
    """
    Helper function for converting a value into the correct type for the 
    series that it is going to be added to. 
    """
    # If the user is trying to make the value None, let them.
    if is_none_type(value) or value is None:
        return None

    try:
        conversion_function = get_function_to_convert_to_series(column_dtype)
        casted_value_series = conversion_function(value, on_uncastable_arg_element=np.nan)

        type_corrected_new_value = casted_value_series.iat[0]

        # If the value is a string and it has a " in it, replace it with a ' so the transpiled code does not error
        if is_string_dtype(column_dtype) and '"' in type_corrected_new_value:
            type_corrected_new_value = type_corrected_new_value.replace('"', "'")

        # If the typed value is not a float, then we do not make it one
        if is_number_dtype(column_dtype) and '.' not in value:
            return round(type_corrected_new_value)

        return type_corrected_new_value
    except:
        raise make_cast_value_to_type_error(value, column_dtype, error_modal=False)
