import inspect
from typing import List, Set

import pandas as pd
import pytest

import hamilton.function_modifiers
from hamilton import base, driver, function_modifiers, models, node
from hamilton.function_modifiers import does
from hamilton.function_modifiers.dependencies import source, value
from hamilton.function_modifiers.macros import Applicable, ensure_function_empty, pipe, step
from hamilton.node import DependencyType

import tests.resources.pipe


def test_no_code_validator():
    def no_code():
        pass

    def no_code_with_docstring():
        """This should still show up as having no code, even though it has a docstring"""
        pass

    def yes_code():
        """This should show up as having no code"""
        a = 0
        return a

    ensure_function_empty(no_code)
    ensure_function_empty(no_code_with_docstring)
    with pytest.raises(hamilton.function_modifiers.base.InvalidDecoratorException):
        ensure_function_empty(yes_code)


# Functions for  @does -- these are the functions we're "replacing"
def _no_params() -> int:
    pass


def _one_param(a: int) -> int:
    pass


def _two_params(a: int, b: int) -> int:
    pass


def _three_params(a: int, b: int, c: int) -> int:
    pass


def _three_params_with_defaults(a: int, b: int = 1, c: int = 2) -> int:
    pass


def _empty() -> int:
    return 1


def _kwargs(**kwargs: int) -> int:
    return sum(kwargs.values())


def _kwargs_with_a(a: int, **kwargs: int) -> int:
    return a + sum(kwargs.values())


def _just_a(a: int) -> int:
    return a


def _just_b(b: int) -> int:
    return b


def _a_b_c(a: int, b: int, c: int) -> int:
    return a + b + c


@pytest.mark.parametrize(
    "fn,replace_with,argument_mapping,matches",
    [
        (_no_params, _empty, {}, True),
        (_no_params, _kwargs, {}, True),
        (_no_params, _kwargs_with_a, {}, False),
        (_no_params, _just_a, {}, False),
        (_no_params, _a_b_c, {}, False),
        (_one_param, _empty, {}, False),
        (_one_param, _kwargs, {}, True),
        (_one_param, _kwargs_with_a, {}, True),
        (_one_param, _just_a, {}, True),
        (_one_param, _just_b, {}, False),
        (_one_param, _just_b, {"b": "a"}, True),  # Replacing a with b makes the signatures match
        (_one_param, _just_b, {"c": "a"}, False),  # Replacing a with b makes the signatures match
        (_two_params, _empty, {}, False),
        (_two_params, _kwargs, {}, True),
        (_two_params, _kwargs_with_a, {}, True),  # b gets fed to kwargs
        (_two_params, _kwargs_with_a, {"foo": "b"}, True),  # Any kwargs work
        (_two_params, _kwargs_with_a, {"bar": "a"}, False),  # No param bar
        (_two_params, _just_a, {}, False),
        (_two_params, _just_b, {}, False),
        (_three_params, _a_b_c, {}, True),
        (_three_params, _a_b_c, {"d": "a"}, False),
        (_three_params, _a_b_c, {}, True),
        (_three_params, _a_b_c, {"a": "b", "b": "a"}, True),  # Weird case but why not?
        (_three_params, _kwargs_with_a, {}, True),
        (_three_params_with_defaults, _a_b_c, {}, True),
        (_three_params_with_defaults, _a_b_c, {"d": "a"}, False),
        (_three_params_with_defaults, _a_b_c, {}, True),
    ],
)
def test_ensure_function_signatures_compatible(fn, replace_with, argument_mapping, matches):
    assert (
        does.test_function_signatures_compatible(
            inspect.signature(fn), inspect.signature(replace_with), argument_mapping
        )
        == matches
    )


def test_does_function_modifier():
    def sum_(**kwargs: int) -> int:
        return sum(kwargs.values())

    def to_modify(param1: int, param2: int) -> int:
        """This sums the inputs it gets..."""
        pass

    annotation = does(sum_)
    (node_,) = annotation.generate_nodes(to_modify, {})
    assert node_.name == "to_modify"
    assert node_.callable(param1=1, param2=1) == 2
    assert node_.documentation == to_modify.__doc__


def test_does_function_modifier_complex_types():
    def setify(**kwargs: List[int]) -> Set[int]:
        return set(sum(kwargs.values(), []))

    def to_modify(param1: List[int], param2: List[int]) -> int:
        """This sums the inputs it gets..."""
        pass

    annotation = does(setify)
    (node_,) = annotation.generate_nodes(to_modify, {})
    assert node_.name == "to_modify"
    assert node_.callable(param1=[1, 2, 3], param2=[4, 5, 6]) == {1, 2, 3, 4, 5, 6}
    assert node_.documentation == to_modify.__doc__


def test_does_function_modifier_optionals():
    def sum_(param0: int, **kwargs: int) -> int:
        return sum(kwargs.values())

    def to_modify(param0: int, param1: int = 1, param2: int = 2) -> int:
        """This sums the inputs it gets..."""
        pass

    annotation = does(sum_)
    (node_,) = annotation.generate_nodes(to_modify, {})
    assert node_.name == "to_modify"
    assert node_.input_types["param0"][1] == DependencyType.REQUIRED
    assert node_.input_types["param1"][1] == DependencyType.OPTIONAL
    assert node_.input_types["param2"][1] == DependencyType.OPTIONAL
    assert node_.callable(param0=0) == 3
    assert node_.callable(param0=0, param1=0, param2=0) == 0
    assert node_.documentation == to_modify.__doc__


def test_does_with_argument_mapping():
    def _sum_multiply(param0: int, param1: int, param2: int) -> int:
        return param0 + param1 * param2

    def to_modify(parama: int, paramb: int = 1, paramc: int = 2) -> int:
        """This sums the inputs it gets..."""
        pass

    annotation = does(_sum_multiply, param0="parama", param1="paramb", param2="paramc")
    (node_,) = annotation.generate_nodes(to_modify, {})
    assert node_.name == "to_modify"
    assert node_.input_types["parama"][1] == DependencyType.REQUIRED
    assert node_.input_types["paramb"][1] == DependencyType.OPTIONAL
    assert node_.input_types["paramc"][1] == DependencyType.OPTIONAL
    assert node_.callable(parama=0) == 2
    assert node_.callable(parama=0, paramb=1, paramc=2) == 2
    assert node_.callable(parama=1, paramb=4) == 9
    assert node_.documentation == to_modify.__doc__


def test_model_modifier():
    config = {
        "my_column_model_params": {
            "col_1": 0.5,
            "col_2": 0.5,
        }
    }

    class LinearCombination(models.BaseModel):
        def get_dependents(self) -> List[str]:
            return list(self.config_parameters.keys())

        def predict(self, **columns: pd.Series) -> pd.Series:
            return sum(
                self.config_parameters[column_name] * column
                for column_name, column in columns.items()
            )

    def my_column() -> pd.Series:
        """Column that will be annotated by a model"""
        pass

    annotation = function_modifiers.model(LinearCombination, "my_column_model_params")
    annotation.validate(my_column)
    (model_node,) = annotation.generate_nodes(my_column, config)
    assert model_node.input_types["col_1"][0] == model_node.input_types["col_2"][0] == pd.Series
    assert model_node.type == pd.Series
    pd.testing.assert_series_equal(
        model_node.callable(col_1=pd.Series([1]), col_2=pd.Series([2])), pd.Series([1.5])
    )

    def bad_model(col_1: pd.Series, col_2: pd.Series) -> pd.Series:
        return col_1 * 0.5 + col_2 * 0.5

    with pytest.raises(hamilton.function_modifiers.base.InvalidDecoratorException):
        annotation.validate(bad_model)


def _test_apply_function(foo: int, bar: int, baz: int = 100) -> int:
    return foo + bar + baz


def _test_apply_function_2(foo: int) -> int:
    return foo + 1


@pytest.mark.parametrize(
    "args,kwargs,chain_first_param",
    [
        ([source("foo_upstream"), value(1)], {}, False),
        ([value(1)], {}, True),
        ([source("foo_upstream"), value(1), value(2)], {}, False),
        ([value(1), value(2)], {}, True),
        ([source("foo_upstream")], {"bar": value(1)}, False),
        ([], {"bar": value(1)}, True),
        ([], {"foo": source("foo_upstream"), "bar": value(1)}, False),
        ([], {"bar": value(1)}, True),
        ([], {"foo": source("foo_upstream"), "bar": value(1), "baz": value(1)}, False),
        ([], {"bar": value(1), "baz": value(1)}, True),
    ],
)
def test_applicable_validates_correctly(args, kwargs, chain_first_param: bool):
    applicable = Applicable(_test_apply_function, args=args, kwargs=kwargs)
    applicable.validate(chain_first_param=chain_first_param, allow_custom_namespace=True)


@pytest.mark.parametrize(
    "args,kwargs,chain_first_param",
    [
        (
            [source("foo_upstream"), value(1)],
            {"foo": source("foo_upstream")},
            True,
        ),  # We chain the first parameter, not pass it in
        ([value(1)], {}, False),  # Not enough first parameters
        ([source("foo_upstream"), value(1), value(2)], {}, True),
        ([value(2)], {"foo": source("foo_upstream")}, False),
        ([source("foo_upstream")], {"bar": value(1)}, True),
        ([], {"bar": value(1)}, False),
        ([], {"foo": source("foo_upstream"), "bar": value(1)}, True),
        ([], {"bar": value(1)}, False),
        ([], {"foo": source("foo_upstream"), "bar": value(1), "baz": value(1)}, True),
        ([], {"bar": value(1), "baz": value(1)}, False),
    ],
)
def test_applicable_does_not_validate(args, kwargs, chain_first_param: bool):
    with pytest.raises(hamilton.function_modifiers.base.InvalidDecoratorException):
        applicable = Applicable(_test_apply_function, args=args, kwargs=kwargs)
        applicable.validate(chain_first_param=chain_first_param, allow_custom_namespace=True)


def test_applicable_does_not_validate_invalid_function_pos_only():
    def foo(a: int, /, b: int) -> int:
        return a + b

    with pytest.raises(hamilton.function_modifiers.base.InvalidDecoratorException):
        applicable = Applicable(foo, args=(source("a"), source("b")), kwargs={})
        applicable.validate(chain_first_param=True, allow_custom_namespace=True)


# We will likely start supporting this in the future, but for now we don't
def test_applicable_does_not_validate_no_param_type_hints():
    def foo(a, b) -> int:
        return a + b

    with pytest.raises(hamilton.function_modifiers.base.InvalidDecoratorException):
        applicable = Applicable(foo, args=(source("a"), source("b")), kwargs={})
        applicable.validate(chain_first_param=True, allow_custom_namespace=True)


def test_applicable_does_not_validate_no_return_type_hints():
    def foo(a: int, b: int):
        return a + b

    with pytest.raises(hamilton.function_modifiers.base.InvalidDecoratorException):
        applicable = Applicable(foo, args=(source("a"), source("b")), kwargs={})
        applicable.validate(chain_first_param=True, allow_custom_namespace=True)


def test_applicable_does_not_validate_invalid_function_no_params():
    def foo() -> int:
        return 1

    with pytest.raises(hamilton.function_modifiers.base.InvalidDecoratorException):
        applicable = Applicable(foo, args=(), kwargs={})
        applicable.validate(chain_first_param=True, allow_custom_namespace=True)


def general_downstream_function(result: int) -> int:
    return result


def test_pipe_decorator_positional_variable_args():
    n = node.Node.from_fn(general_downstream_function)

    decorator = pipe(
        step(_test_apply_function, source("bar_upstream"), baz=value(1000)).named("node_1"),
        namespace=None,
    )
    nodes = decorator.transform_dag([n], {}, general_downstream_function)
    nodes_by_name = {item.name: item for item in nodes}
    chain_node = nodes_by_name["node_1"]
    assert chain_node(result=1, bar_upstream=10) == 1011  # This chains it through
    assert sorted(chain_node.input_types) == ["bar_upstream", "result"]
    final_node = nodes_by_name["general_downstream_function"]
    assert final_node(node_1=1) == 1


def test_pipe_decorator_no_collapse_single_node():
    n = node.Node.from_fn(general_downstream_function)

    decorator = pipe(
        step(_test_apply_function, bar=source("bar_upstream"), baz=value(1000)).named("node_1"),
        namespace=None,
    )
    nodes = decorator.transform_dag([n], {}, general_downstream_function)
    nodes_by_name = {item.name: item for item in nodes}
    chain_node = nodes_by_name["node_1"]
    assert chain_node(result=1, bar_upstream=10) == 1011  # This chains it through
    assert sorted(chain_node.input_types) == ["bar_upstream", "result"]
    final_node = nodes_by_name["general_downstream_function"]
    assert final_node(node_1=1) == 1

    # It should take in the value result, and cascade it through to this
    # result -> "1" -> general_downstream_function


def test_pipe_decorator_no_collapse_multi_node():
    n = node.Node.from_fn(general_downstream_function)

    decorator = pipe(
        step(_test_apply_function, bar=source("bar_upstream"), baz=100).named("node_1"),
        step(_test_apply_function, bar=value(10), baz=value(100)).named("node_2"),
        namespace=None,
    )
    nodes = decorator.transform_dag([n], {}, general_downstream_function)
    nodes_by_name = {item.name: item for item in nodes}
    final_node = nodes_by_name["general_downstream_function"]
    assert len(nodes_by_name) == 3
    assert nodes_by_name["node_1"](result=1, bar_upstream=10) == 111
    assert nodes_by_name["node_2"](node_1=1) == 111
    assert final_node(node_2=100) == 100


def test_resolve_namespace_inherit():
    applicable = Applicable(
        _test_apply_function, args=(), kwargs=dict(bar=source("bar_upstream"), baz=100)
    ).named("node_1")
    assert applicable.resolve_namespace("inherited") == ("inherited",)


def test_resolve_namespace_discard():
    applicable = Applicable(
        _test_apply_function, args=(), kwargs=dict(bar=source("bar_upstream"), baz=100)
    ).named("node_1", namespace=None)
    assert applicable.resolve_namespace("unused") == ()


def test_resolve_namespace_replaced():
    applicable = Applicable(
        _test_apply_function, args=(), kwargs=dict(bar=source("bar_upstream"), baz=100)
    ).named("node_1", namespace="replaced")
    assert applicable.resolve_namespace("unused") == ("replaced",)


def test_validate_pipe_fails_with_conflicting_namespace():
    decorator = pipe(
        step(_test_apply_function, bar=source("bar_upstream"), baz=100).named(
            "node_1", namespace="custom"
        ),
        namespace=None,  # Not allowed to have custom namespacess if the namespace is None
    )
    with pytest.raises(hamilton.function_modifiers.base.InvalidDecoratorException):
        decorator.validate(general_downstream_function)


def test_inherits_null_namespace():
    n = node.Node.from_fn(general_downstream_function)
    decorator = pipe(
        step(_test_apply_function, bar=source("bar_upstream"), baz=100).named(
            "node_1", namespace=...
        ),
        namespace=None,  # Not allowed to have custom namespacess if the namespace is None
    )
    decorator.validate(general_downstream_function)
    nodes = decorator.transform_dag([n], {}, general_downstream_function)
    assert "node_1" in {item.name for item in nodes}
    assert "general_downstream_function" in {item.name for item in nodes}


def test_pipe_end_to_end_1():
    dr = (
        driver.Builder()
        .with_modules(tests.resources.pipe)
        .with_adapter(base.DefaultAdapter())
        .with_config({"calc_c": True})
        .build()
    )

    inputs = {
        "input_1": 10,
        "input_2": 20,
        "input_3": 30,
    }
    result = dr.execute(
        [
            "chain_1_using_pipe",
            "chain_2_using_pipe",
            "chain_1_not_using_pipe",
            "chain_2_not_using_pipe",
        ],
        inputs=inputs,
    )
    assert result["chain_1_using_pipe"] == result["chain_1_not_using_pipe"]
    assert result["chain_2_using_pipe"] == result["chain_2_not_using_pipe"]


def test_pipe_end_to_end_2():
    dr = (
        driver.Builder()
        .with_modules(tests.resources.pipe)
        .with_adapter(base.DefaultAdapter())
        .build()
    )

    inputs = {
        "input_1": 10,
        "input_2": 20,
        "input_3": 30,
    }
    result = dr.execute(
        [
            "chain_1_using_pipe",
            "chain_2_using_pipe",
            "chain_1_not_using_pipe",
            "chain_2_not_using_pipe",
        ],
        inputs=inputs,
    )
    assert result["chain_1_using_pipe"] == result["chain_1_not_using_pipe"]
    assert result["chain_2_using_pipe"] == result["chain_2_not_using_pipe"]