from datetime import datetime

import pytest
from pytest_mock import MockerFixture

from tests.helpers import ColumnCreator
from vivarium.framework.population.manager import SimulantData
from vivarium.framework.randomness import RandomnessStream
from vivarium.framework.randomness.index_map import IndexMap
from vivarium.framework.resource.exceptions import ResourceError
from vivarium.framework.resource.group import ResourceGroup
from vivarium.framework.resource.resource import Column, NullResource, Resource
from vivarium.framework.values import AttributePipeline, Pipeline, ValueModifier, ValueSource


@pytest.mark.parametrize("resource_type", ["column", "resource"])
def test_resource_group(resource_type: str, mocker: MockerFixture) -> None:
    component = ColumnCreator()
    resources: list[Column] | Resource
    if resource_type == "column":
        resources = [Column(f"resource_{i}", component) for i in range(5)]
    else:
        resources = Resource("test", "some_resource", component)
    r_required_resources = [
        AttributePipeline("an_interesting_attribute", None),
        Pipeline("baz"),
        RandomnessStream("bar", lambda: datetime.now(), 1, IndexMap(), mocker.Mock()),
        ValueSource(Pipeline("foo"), lambda: 1, None),
    ]

    rg = ResourceGroup(
        initialized_resources=resources,
        required_resources=r_required_resources,
        initializer=component.initialize_test_columns,
    )

    assert rg.component == component
    assert rg.type == "column" if resource_type == "column" else "test"
    assert (
        rg.names == [f"column.{res.name}" for res in resources]
        if isinstance(resources, list)
        else ["test.some_resource"]
    )
    assert rg.initializer == component.initialize_test_columns
    assert rg.required_resources == [
        "attribute.an_interesting_attribute",
        "value.baz",
        "stream.bar",
        "value_source.foo",
    ]
    assert list(rg) == rg.names


@pytest.mark.parametrize(
    "resource",
    [
        Pipeline("foo", ColumnCreator()),
        AttributePipeline("foo", ColumnCreator()),
        ValueSource(Pipeline("bar", ColumnCreator()), lambda: 1, ColumnCreator()),
        ValueModifier(Pipeline("baz", ColumnCreator()), lambda: 1, ColumnCreator()),
        Column("foo", ColumnCreator()),
        RandomnessStream("bar", lambda: datetime.now(), 1, IndexMap(), ColumnCreator()),
        NullResource(0, ColumnCreator()),
    ],
)
@pytest.mark.parametrize("initialized", [True, False])
def test_resource_group_is_initializer(resource: Resource, initialized: bool) -> None:
    def some_initializer(pop_data: SimulantData) -> None:
        pass

    rg = ResourceGroup(
        resource, [Resource("test", "bar", None)], some_initializer if initialized else None
    )
    assert rg.is_initialized == initialized


def test_resource_group_with_no_resources() -> None:
    with pytest.raises(ResourceError, match="must have at least one resource"):
        _ = ResourceGroup([], [Resource("test", "foo", None)], None)


def test_resource_group_with_multiple_components() -> None:
    # This test is not terribly relevant since ResourceGroup now only accepts a list
    # of Columns or a single Resource. We keep it around in case that changes.
    resources = [
        Column("foo", ColumnCreator()),
        Column("bar", ColumnCreator()),
    ]

    with pytest.raises(ResourceError, match="resources must have the same component"):
        _ = ResourceGroup(resources, [], None)


def test_resource_group_with_multiple_resource_types() -> None:

    # This test is not terribly relevant since ResourceGroup now only accepts a list
    # of Columns or a single Resource. We keep it around in case that changes.
    component = ColumnCreator()
    resources = [
        ValueModifier(Pipeline("foo"), lambda: 1, component),
        ValueSource(Pipeline("bar"), lambda: 2, component),
    ]

    with pytest.raises(ResourceError, match="resources must be of the same type"):
        _ = ResourceGroup(resources, [], None)  # type: ignore[arg-type]


def test_set_required_resources(mocker: MockerFixture) -> None:
    some_attribute = AttributePipeline("some_attribute", mocker.Mock())
    some_other_attribute = AttributePipeline("some_other_attribute", mocker.Mock())
    resource = Resource("test", "some_resource", mocker.Mock())
    required_resources: list[AttributePipeline | str] = [
        some_attribute,
        "some_other_attribute",
    ]

    rg = ResourceGroup(
        initialized_resources=resource,
        required_resources=required_resources,
        initializer=None,
    )
    assert rg._required_resources == [some_attribute, "some_other_attribute"]
    # Mock the attribute pipelines dict
    attribute_pipelines = {
        "some_attribute": some_attribute,
        "some_other_attribute": some_other_attribute,
    }
    rg.set_required_resources(attribute_pipelines)
    # Check that the 'some_other_attribute' string has been replaced by the pipeline
    assert rg._required_resources == [some_attribute, some_other_attribute]
