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.relationship.model import Relationship
from infrahub.core.schema.relationship_schema import RelationshipSchema
from infrahub.core.timestamp import Timestamp
from infrahub.database import InfrahubDatabase

from .shared import to_frozen_set


@dataclass
class QueryPeerParams:
    branch: Branch | str
    source_kind: str
    schema: RelationshipSchema
    filters: dict[str, Any]
    fields: dict | None = None
    at: Timestamp | str | None = None
    branch_agnostic: bool = False
    include_metadata: MetadataQueryOptions = field(default_factory=MetadataQueryOptions)

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


class PeerRelationshipsDataLoader(DataLoader[str, list[Relationship]]):
    def __init__(self, db: InfrahubDatabase, query_params: QueryPeerParams, *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[list[Relationship]]:
        async with self.db.start_session(read_only=True) as db:
            peer_rels = await NodeManager.query_peers(
                db=db,
                ids=keys,
                source_kind=self.query_params.source_kind,
                schema=self.query_params.schema,
                filters=self.query_params.filters,
                fields=self.query_params.fields,
                at=self.query_params.at,
                branch=self.query_params.branch,
                branch_agnostic=self.query_params.branch_agnostic,
                include_metadata=self.query_params.include_metadata,
                fetch_peers=True,
            )
        peer_rels_by_node_id: dict[str, list[Relationship]] = {}
        for rel in peer_rels:
            node_id = rel.node_id
            if node_id not in peer_rels_by_node_id:
                peer_rels_by_node_id[node_id] = []
            peer_rels_by_node_id[node_id].append(rel)

        results = []
        for node_id in keys:
            results.append(peer_rels_by_node_id.get(node_id, []))
        return results
