from awspipe.aws import AwsSessionContext
from awspipe.backend import BackendAdapter
from awspipe.engine import QueryEngine
from awspipe.providers.base import Column, TableProvider


class FakeSessionFactory:
    def __init__(self) -> None:
        self.resolve_context_calls: list[tuple[str | None, str | None]] = []

    def resolve_context(self, profile=None, region=None):
        self.resolve_context_calls.append((profile, region))
        return AwsSessionContext(
            session=None,
            profile=profile,
            region=region or "eu-west-1",
            account_id="123456789012",
            caller_identity={"Account": "123456789012", "Arn": "arn:aws:iam::123456789012:user/test", "UserId": "AIDA"},
        )

    def resolve_regions(self, session, region=None, all_regions=False):
        if all_regions:
            return ["eu-west-1", "us-east-1"]
        return [region or "eu-west-1"]

    def doctor(self, profile=None, region=None, cache_path=""):
        raise AssertionError("doctor not expected")


class FakeRegistry:
    def __init__(self, provider):
        self.provider = provider

    def get(self, table_name):
        if table_name == self.provider.name:
            return self.provider
        return None

    def all(self):
        return [self.provider]


def test_query_engine_populates_cache_and_reuses_it(tmp_path) -> None:
    calls = {"count": 0}

    def fetcher(context):
        calls["count"] += 1
        return [
            {
                "profile_name": context.profile or "default",
                "account_id": context.account_id,
                "region": context.region,
                "arn": "arn:aws:test:::resource/1",
                "title": "resource-1",
                "resource_id": "resource-1",
            }
        ]

    provider = TableProvider(
        name="aws_dummy",
        description="dummy",
        columns=(
            Column("profile_name", "VARCHAR", "Profile."),
            Column("account_id", "VARCHAR", "Account."),
            Column("region", "VARCHAR", "Region."),
            Column("arn", "VARCHAR", "Arn."),
            Column("title", "VARCHAR", "Title."),
            Column("resource_id", "VARCHAR", "Resource id."),
        ),
        scope="regional",
        fetcher=fetcher,
    )
    engine = QueryEngine(
        backend=BackendAdapter(tmp_path / "cache.duckdb"),
        session_factory=FakeSessionFactory(),
        registry=FakeRegistry(provider),
        ttl_seconds=300,
    )

    first = engine.query("select resource_id from aws_dummy", profile="dev", region="eu-west-1")
    second = engine.query("select resource_id from aws_dummy", profile="dev", region="eu-west-1")
    refreshed = engine.query("select resource_id from aws_dummy", profile="dev", region="eu-west-1", refresh=True)

    assert first.rows == [("resource-1",)]
    assert second.rows == [("resource-1",)]
    assert refreshed.rows == [("resource-1",)]
    assert calls["count"] == 2


def test_query_engine_fetches_regional_tables_for_all_regions(tmp_path) -> None:
    calls: list[str] = []

    def fetcher(context):
        calls.append(context.region)
        return [
            {
                "profile_name": context.profile or "default",
                "account_id": context.account_id,
                "region": context.region,
                "arn": f"arn:aws:test:::{context.region}",
                "title": context.region,
                "resource_id": context.region,
            }
        ]

    provider = TableProvider(
        name="aws_dummy",
        description="dummy",
        columns=(
            Column("profile_name", "VARCHAR", "Profile."),
            Column("account_id", "VARCHAR", "Account."),
            Column("region", "VARCHAR", "Region."),
            Column("arn", "VARCHAR", "Arn."),
            Column("title", "VARCHAR", "Title."),
            Column("resource_id", "VARCHAR", "Resource id."),
        ),
        scope="regional",
        fetcher=fetcher,
    )
    session_factory = FakeSessionFactory()
    engine = QueryEngine(
        backend=BackendAdapter(tmp_path / "cache.duckdb"),
        session_factory=session_factory,
        registry=FakeRegistry(provider),
        ttl_seconds=300,
    )

    result = engine.query(
        "select resource_id from aws_dummy order by resource_id",
        profile="dev",
        region="eu-west-1",
        all_regions=True,
    )

    assert result.rows == [("eu-west-1",), ("us-east-1",)]
    assert calls == ["eu-west-1", "us-east-1"]
    assert session_factory.resolve_context_calls == [("dev", "eu-west-1"), ("dev", "us-east-1")]
