from __future__ import annotations

import copy
import logging
from typing import Any, Collection, Coroutine, Sequence

import aiohttp
from pydantic import BaseModel, ConfigDict

from ..types import Doc, DocDetails
from ..utils import gather_with_concurrency
from .client_models import MetadataPostProcessor, MetadataProvider
from .crossref import CrossrefProvider
from .journal_quality import JournalQualityPostProcessor
from .semantic_scholar import SemanticScholarProvider

logger = logging.getLogger(__name__)

ALL_CLIENTS: (
    Collection[type[MetadataPostProcessor | MetadataProvider]]
    | Sequence[Collection[type[MetadataPostProcessor | MetadataProvider]]]
) = {
    CrossrefProvider,
    SemanticScholarProvider,
    JournalQualityPostProcessor,
}


class DocMetadataTask(BaseModel):
    """Holder for provider and processor tasks."""

    providers: Collection[MetadataProvider]
    processors: Collection[MetadataPostProcessor]

    model_config = ConfigDict(arbitrary_types_allowed=True)

    def provider_queries(
        self, query: dict
    ) -> list[Coroutine[Any, Any, DocDetails | None]]:
        return [p.query(query) for p in self.providers]

    def processor_queries(
        self, doc_details: DocDetails, session: aiohttp.ClientSession
    ) -> list[Coroutine[Any, Any, DocDetails]]:
        return [
            p.process(copy.copy(doc_details), session=session) for p in self.processors
        ]

    def __repr__(self) -> str:
        return (
            f"DocMetadataTask(providers={self.providers}, processors={self.processors})"
        )


class DocMetadataClient:
    def __init__(
        self,
        session: aiohttp.ClientSession | None = None,
        clients: (
            Collection[type[MetadataPostProcessor | MetadataProvider]]
            | Sequence[Collection[type[MetadataPostProcessor | MetadataProvider]]]
        ) = ALL_CLIENTS,
    ) -> None:
        """Metadata client for querying multiple metadata providers and processors.

        Args:
            session: outer scope aiohttp session to allow for connection pooling
            clients: list of MetadataProvider and MetadataPostProcessor classes to query;
                if nested, will query in order looking for termination criteria after each.
                Will terminate early if either DocDetails.is_hydration_needed is False OR if
                all requested fields are present in the DocDetails object.

        """
        self._session = session
        self.tasks: list[DocMetadataTask] = []

        # first see if we are nested; i.e. we want order
        if isinstance(clients, Sequence) and all(
            isinstance(sub_clients, Collection) for sub_clients in clients
        ):
            for sub_clients in clients:
                self.tasks.append(
                    DocMetadataTask(
                        providers=[
                            c() for c in sub_clients if issubclass(c, MetadataProvider)
                        ],
                        processors=[
                            c()
                            for c in sub_clients
                            if issubclass(c, MetadataPostProcessor)
                        ],
                    )
                )
        # otherwise, we are a flat collection
        if not self.tasks and all(not isinstance(c, Collection) for c in clients):
            self.tasks.append(
                DocMetadataTask(
                    providers=[c() for c in clients if issubclass(c, MetadataProvider)],  # type: ignore[operator, arg-type]
                    processors=[
                        c() for c in clients if issubclass(c, MetadataPostProcessor)  # type: ignore[operator, arg-type]
                    ],
                )
            )

        if not self.tasks or (self.tasks and not self.tasks[0].providers):
            raise ValueError("At least one MetadataProvider must be provided.")

    async def query(self, **kwargs) -> DocDetails | None:

        session = aiohttp.ClientSession() if self._session is None else self._session

        query_args = kwargs if "session" in kwargs else kwargs | {"session": session}

        doc_details: DocDetails | None = None

        for ti, task in enumerate(self.tasks):
            logger.debug(
                f"Attempting to populate metadata query: {query_args} via {task}"
            )

            # first query all client_models and aggregate the results
            doc_details = (
                sum(
                    p
                    for p in (
                        await gather_with_concurrency(
                            len(task.providers), task.provider_queries(query_args)
                        )
                    )
                    if p
                )
                or None
            )

            # then process and re-aggregate the results
            if doc_details and task.processors:
                doc_details = sum(
                    await gather_with_concurrency(
                        len(task.processors),
                        task.processor_queries(doc_details, session),
                    )
                )

            if doc_details and not doc_details.is_hydration_needed(
                inclusion=kwargs.get("fields", [])
            ):
                logger.debug(
                    "All requested fields are present in the DocDetails "
                    f"object{', stopping early.' if ti != len(self.tasks) - 1 else '.'}"
                )
                break

        if self._session is None:
            await session.close()

        return doc_details

    async def bulk_query(
        self, queries: Collection[dict[str, Any]], concurrency: int = 10
    ) -> list[DocDetails]:
        return await gather_with_concurrency(
            concurrency, [self.query(**kwargs) for kwargs in queries]
        )

    async def upgrade_doc_to_doc_details(self, doc: Doc, **kwargs) -> DocDetails:

        # note we have some extra fields which may have come from reading the doc text,
        # but aren't in the doc object, we add them here too.
        extra_fields = {
            k: v for k, v in kwargs.items() if k in {"title", "authors", "doi"}
        }
        # abuse our doc_details object to be an int if it's empty
        # our __add__ operation supports int by doing nothing
        extra_doc: int | DocDetails = (
            0 if not extra_fields else DocDetails(**extra_fields)
        )

        if doc_details := await self.query(**kwargs):
            if doc.overwrite_fields_from_metadata:
                return extra_doc + doc_details

            # hard overwrite the details from the prior object
            doc_details.dockey = doc.dockey
            doc_details.doc_id = doc.dockey
            doc_details.docname = doc.docname
            doc_details.key = doc.docname
            doc_details.citation = doc.citation
            return extra_doc + doc_details

        # if we can't get metadata, just return the doc, but don't overwrite any fields
        prior_doc = doc.model_dump()
        prior_doc["overwrite_fields_from_metadata"] = False
        return DocDetails(**(prior_doc | extra_fields))
