# stdlib
from typing import Callable
from typing import List
from typing import Union

# third party
import numpy as np
import pandas as pd

# relative
from .action_object import ActionObject
from .response import SyftError
from .response import SyftResponseMessage
from .response import SyftSuccess


def verify_result(
    func: Callable,
    private_inputs: Union[ActionObject, List[ActionObject]],
    private_outputs: Union[ActionObject, List[ActionObject]],
) -> SyftResponseMessage:
    """Verify a single result of Code Verification"""
    trace_assets = []
    if not isinstance(private_inputs, List):
        private_inputs = [private_inputs]

    for asset in private_inputs:
        if not isinstance(asset, ActionObject):
            raise Exception(f"ActionObject expected, instead received: {type(asset)}")
        # Manual type casting for now, to automate later
        if isinstance(asset.syft_action_data, np.ndarray):
            empty_obj: ActionObject = ActionObject(
                id=asset.id, syft_result_obj=np.ndarray([])
            )
        elif isinstance(asset.syft_action_data, pd.DataFrame):
            empty_obj: ActionObject = ActionObject(
                id=asset.id, syft_result_obj=pd.DataFrame()
            )
        else:
            raise NotImplementedError(
                f"Trace mode not yet automated for type: {type(asset.syft_action_data)}"
            )

        trace_assets.append(empty_obj)

    print("Code Verification in progress.")
    traced_results = func(*trace_assets)

    if isinstance(private_outputs, List):
        target_hashes = [output.syft_history_hash for output in private_outputs]
        traced_hashes = [result.syft_history_hash for result in traced_results]
    else:
        target_hashes = private_outputs.syft_history_hash
        traced_hashes = traced_results.syft_history_hash

    if target_hashes == traced_hashes:
        msg = "Code Verification passed with matching hashes! Congratulations, and thank you for supporting PySyft!"
        return SyftSuccess(message=msg)
    else:
        msg = (
            f"Hashes do not match! Target hashes were: {target_hashes} but Traced hashes were: {traced_results}. "
            f"Please try checking the logs."
        )
        return SyftError(message=msg)


def code_verification(func: Callable):
    """Compares history hashes of an Empty Action Object to that of the real action object

    Inputs:
    - func:: a Callable whose sole argument should be the Private Dataset(s) being used. Constraints:
        - Input arguments are the private datasets being used
        - Output arguments are the results requested.

    Outputs:
    - boolean:: if history hashes match
    """

    def wrapper(*args, **kwargs):
        trace_assets = []
        for asset in args:
            if not isinstance(asset, ActionObject):
                raise Exception(
                    f"ActionObject expected, instead received: {type(asset)}"
                )
            # Manual type casting for now, to automate later
            if isinstance(asset.syft_action_data, np.ndarray):
                empty_obj: ActionObject = ActionObject(
                    id=asset.id, syft_result_obj=np.ndarray([])
                )
            elif isinstance(asset.syft_action_data, pd.DataFrame):
                empty_obj: ActionObject = ActionObject(
                    id=asset.id, syft_result_obj=pd.DataFrame()
                )
            else:
                raise NotImplementedError(
                    f"Trace mode not yet automated for type: {type(asset.syft_action_data)}"
                )

            trace_assets.append(empty_obj)

        print("Evaluating function normally to obtain history hash")
        results = func(*args, **kwargs).syft_history_hash
        print(8 * "(-(-_(-_-)_-)-)")

        print("Tracing function to obtain history hash")
        traced_results = func(*trace_assets, **kwargs).syft_history_hash
        print(8 * "(-(-_(-_-)_-)-)")

        # assert len(results) == len(traced_results)
        hashes_match = results == traced_results
        if hashes_match:
            msg = (
                f"Code Verification passed with matching hashes of {results}! Congratulations, and thank you for "
                f"supporting PySyft!"
            )
            return SyftSuccess(message=msg)
        else:
            msg = (
                f"Hashes do not match! Target hashes were: {results} but Traced hashes were: {traced_results}. "
                f"Please try checking the logs."
            )
            return SyftError(message=msg)

    return wrapper
