from dataclasses import dataclass, field
from typing import Any

from aiodataloader import DataLoader

from infrahub.core.branch.models import Branch
from infrahub.core.manager import NodeManager
from infrahub.core.metadata.model import MetadataQueryOptions
from infrahub.core.node import Node
from infrahub.core.timestamp import Timestamp
from infrahub.database import InfrahubDatabase

from .shared import to_frozen_set


@dataclass
class GetManyParams:
    branch: Branch | str
    fields: dict | None = None
    at: Timestamp | str | None = None
    include_metadata: MetadataQueryOptions = field(default_factory=MetadataQueryOptions)
    prefetch_relationships: bool = False
    branch_agnostic: bool = False

    def __hash__(self) -> int:
        frozen_fields: frozenset | None = None
        if self.fields:
            frozen_fields = to_frozen_set(self.fields)
        timestamp = Timestamp(self.at)
        branch = self.branch.name if isinstance(self.branch, Branch) else self.branch
        hash_str = "|".join(
            [
                str(hash(frozen_fields)),
                timestamp.to_string(),
                branch,
                str(hash(self.include_metadata)),
                str(self.prefetch_relationships),
                str(self.branch_agnostic),
            ]
        )
        return hash(hash_str)


class NodeDataLoader(DataLoader[str, Node | None]):
    def __init__(self, db: InfrahubDatabase, query_params: GetManyParams, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self.query_params = query_params
        self.db = db

    async def batch_load_fn(self, keys: list[Any]) -> list[Node | None]:
        async with self.db.start_session(read_only=True) as db:
            nodes_by_id = await NodeManager.get_many(
                db=db,
                ids=keys,
                fields=self.query_params.fields,
                at=self.query_params.at,
                branch=self.query_params.branch,
                include_metadata=self.query_params.include_metadata,
                prefetch_relationships=self.query_params.prefetch_relationships,
                branch_agnostic=self.query_params.branch_agnostic,
            )
        results = []
        for node_id in keys:
            results.append(nodes_by_id.get(node_id, None))
        return results
