from __future__ import annotations

import ast
import inspect
import sys
import time
import warnings
from contextlib import ExitStack
from textwrap import dedent, indent
from typing import TYPE_CHECKING, Any, Callable, Dict, Optional

import pandas as pd

import fused_batch
from fused_batch._udf.context import ExecutionContextProtocol, context
from fused_batch._udf.state import (
    decorator_src_override_context,
    noop_decorators_context,
)
from fused_batch.models.schema import Schema
from fused_batch.models.udf import EMPTY_UDF, AnyBaseUdf
from fused_batch.models.udf._eval_result import UdfEvaluationResult
from fused_batch.models.udf.output import Output, PandasOutput
from fused_batch.warnings import FusedWarning

if TYPE_CHECKING:
    from fused_batch.models.input import BaseInput


def _add_special_kwargs(
    fn: Callable,
    args: Dict[str, Any],
    output: Output,
    context: ExecutionContextProtocol,
):
    signature = inspect.signature(fn)
    param_names = signature.parameters.keys()
    # note: we don't check that this is a kwarg on the function, so if it's the first parameter it could cause
    # a call failure.
    if "output" in param_names:
        args["output"] = output
    if "context" in param_names:
        args["context"] = context

    return args


def execute_against_sample(
    udf: AnyBaseUdf,
    input: BaseInput,
    update_schema: bool = True,
    validate_output: bool = True,
    validate_imports: Optional[bool] = None,
    **kwargs,
) -> UdfEvaluationResult:
    if udf is EMPTY_UDF:
        raise ValueError("Empty UDF cannot be evaluated. Use `set_udf` to set a UDF.")

    output = Output()
    try:
        # Define custom function in environment

        # Wrap cell in a function
        # This is believed to fix scoping issues with closures and list comprehensions
        # See also https://stackoverflow.com/a/2749806,
        # https://stackoverflow.com/q/38318370
        # {cell_text} is not itself indented because we need to make sure that every
        # line of {cell_text} is indented, not just the first.
        FN_RESULT_NAME = "__fused_fn_result"
        wrapped_cell = """\
            def __fused_exec():
            {cell_text}

                return {entrypoint}

            {FN_RESULT_NAME} = __fused_exec()
            """
        original_src = udf.code
        src = dedent(wrapped_cell).format(
            cell_text=indent(dedent(original_src), " " * 4),
            entrypoint=udf.entrypoint,
            FN_RESULT_NAME=FN_RESULT_NAME,
        )

        exec_locals = {}
        exec_globals = {"fused_batch": fused_batch}

        # Store names of modules prior to rendering Headers
        preceding_sys_modules = list(sys.modules.keys())

        # Validate import stamements correspond to valid modules
        validate_imports_whitelist(udf, validate_imports=validate_imports)

        with ExitStack() as stack:
            stack.enter_context(noop_decorators_context(True))

            # Add headers to sys.meta_path
            if udf.headers is not None:
                for header in udf.headers:
                    stack.enter_context(
                        header._register_custom_finder(preceding_sys_modules)
                    )

            exec(src, exec_globals, exec_locals)

            _fn = exec_locals[FN_RESULT_NAME]

            kwargs_to_pass = {**udf.parameters, **kwargs}
            _add_special_kwargs(
                _fn, args=kwargs_to_pass, output=output, context=context
            )
            args_to_pass = (
                input.as_udf_args() if hasattr(input, "as_udf_args") else input
            )

            # Run UDF
            time_start = time.perf_counter()
            _output = _fn(*args_to_pass, **kwargs_to_pass)
            time_end = time.perf_counter()
            time_taken_seconds = time_end - time_start

            if not validate_output:
                return _output

            if _output is not None:
                output.data = _output

    except Exception:
        raise

    new_output = _transform_output(output=output)

    if new_output is None:
        # TODO: Assumes table_schema is present, which doesn't match type
        if udf.table_schema not in [None, {}, Schema()]:
            warnings.warn(
                "UDF is configured with a schema but returns `None`. An empty schema was set for this execution.",
                FusedWarning,
            )

        return UdfEvaluationResult(
            data=None,
            sidecar=None,
            udf=udf,
            table_schema={},
            time_taken_seconds=time_taken_seconds,
        )

    # Validate the dataframe after assigning it to the output variable in user_ns
    # so the user can inspect the output if anything is wrong.
    new_output.validate_data_with_schema()

    if update_schema and udf.table_schema is None:
        udf.table_schema = new_output.table_schema

    return UdfEvaluationResult(
        data=new_output.data,
        sidecar=new_output.sidecar_output,
        udf=udf,
        table_schema=new_output.table_schema,
        time_taken_seconds=time_taken_seconds,
    )


def execute_for_decorator(udf: AnyBaseUdf) -> AnyBaseUdf:
    """Evaluate a UDF for the purpose of getting the UDF object out of it."""
    # Define custom function in environment

    # This is a stripped-down version of execute_against_sample, above.

    FN_RESULT_NAME = "__fused_fn_result"
    wrapped_cell = """\
        def __fused_exec():
        {cell_text}

            return {entrypoint}

        {FN_RESULT_NAME} = __fused_exec()
        """
    original_src = udf.code
    src = dedent(wrapped_cell).format(
        cell_text=indent(dedent(original_src), " " * 4),
        entrypoint=udf.entrypoint,
        FN_RESULT_NAME=FN_RESULT_NAME,
    )

    exec_locals = {}
    exec_globals = {"fused_batch": fused_batch}

    with ExitStack() as stack:
        stack.enter_context(decorator_src_override_context(original_src))

        # Add headers to sys.meta_path
        if udf.headers is not None:
            for header in udf.headers:
                stack.enter_context(header._register_custom_finder())

        exec(src, exec_globals, exec_locals)

        _fn = exec_locals[FN_RESULT_NAME]

        return _fn


def _transform_output(output: Output) -> PandasOutput:
    # TODO: Support PyArrow
    if isinstance(output, PandasOutput):
        return output

    if isinstance(output.data, pd.DataFrame):
        # Force all column names to be strings
        output.data.columns = [str(x) for x in output.data.columns]

        # Note that we don't pass a _new_ schema argument here
        # - If the user defined a schema onto output.table_schema, that will be
        #   preserved.
        # - If the user did not define output.table_schema, then output.table_schema
        #   will be None, and PandasOutput will default to a schema inferred from
        #   the DataFrame.
        return PandasOutput(
            data=output.data,
            table_schema=output.table_schema,
            sidecar_output=output.sidecar_output,
            skip_fused_index_validation=output.skip_fused_index_validation,
        )
    elif output.data is None:
        return None

    raise TypeError(f"Unexpected result type {type(output.data)}")


def validate_imports_whitelist(
    udf: AnyBaseUdf, validate_imports: Optional[bool] = None
):
    # Skip import validation if the option is set
    if (
        not fused_batch.options.default_validate_imports
        and validate_imports is not True
    ):
        return

    # Skip import validation if not logged in
    if not fused_batch.api.AUTHORIZATION.is_configured():
        return

    from fused_batch._global_api import get_api

    # Get the dependency whitelist from the cached API endpoint
    api = get_api()
    package_dependencies = api.dependency_whitelist()

    # Initialize a list to store the import statements
    import_statements = []

    # Parse the source code into an AST
    tree = ast.parse(udf.code)

    # Traverse the AST to find import statements
    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            for alias in node.names:
                import_statements.append(alias.name)
        elif isinstance(node, ast.ImportFrom):
            module_name = node.module
            import_statements.append(module_name)

    # Check for unavailable modules
    header_modules = [header.module_name for header in udf.headers]
    fused_modules = ["fused"]  # assume fused is always available
    available_modules = (
        list(package_dependencies["dependency_whitelist"].keys())
        + header_modules
        + fused_modules
    )
    unavailable_modules = []
    for import_statement in import_statements:
        if import_statement.split(".", 1)[0] not in available_modules:
            unavailable_modules.append(import_statement)

    if unavailable_modules:
        raise ValueError(
            f"The following imports in the UDF might not be available: {repr(unavailable_modules)}. Please check the UDF headers and imports and try again."
        )

    # TODO: check major versions for some packages
