#  Copyright 2021-present, the Recognai S.L. team.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

from typing import Any, Dict, Optional, Type

from fastapi import Depends

from argilla_server.daos.backend.search.model import SortConfig
from argilla_server.daos.models.records import DaoRecordsSearch
from argilla_server.daos.records import DatasetRecordsDAO
from argilla_server.services.datasets import ServiceDataset
from argilla_server.services.metrics.models import ServiceMetric, ServicePythonMetric
from argilla_server.services.search.model import ServiceBaseRecordsQuery, ServiceRecordsQuery
from argilla_server.services.tasks.commons import ServiceRecord


class MetricsService:
    """The dataset metrics service singleton"""

    _INSTANCE = None

    @classmethod
    def get_instance(
        cls,
        dao: DatasetRecordsDAO = Depends(DatasetRecordsDAO.get_instance),
    ) -> "MetricsService":
        """
        Creates the service instance.

        Parameters
        ----------
        dao:
            The dataset records dao

        Returns
        -------
            The metrics service instance

        """
        if not cls._INSTANCE:
            cls._INSTANCE = cls(dao)
        return cls._INSTANCE

    def __init__(self, dao: DatasetRecordsDAO):
        """
        Creates a service instance

        Parameters
        ----------
        dao:
            The dataset records dao
        """
        self.__dao__ = dao

    def summarize_metric(
        self,
        dataset: ServiceDataset,
        metric: ServiceMetric,
        record_class: Optional[Type[ServiceRecord]] = None,
        query: Optional[ServiceRecordsQuery] = None,
        **metric_params,
    ) -> Dict[str, Any]:
        """
        Applies a metric summarization.

        Parameters
        ----------
        dataset:
            The records dataset
        metric:
            The selected metric
        record_class:
            The record class type for python metrics computation
        query:
            An optional query passed for records filtering
        metric_params:
            Related metrics parameters

        Returns
        -------
            The metric summarization info
        """

        if isinstance(metric, ServicePythonMetric):
            query = query or ServiceBaseRecordsQuery()
            query = metric.prepare_query(query)
            records = self.__dao__.scan_dataset(
                dataset,
                search=DaoRecordsSearch(query=query, sort=SortConfig(shuffle=metric.shuffle_records)),
                limit=metric.records_to_fetch,
                exclude_fields={
                    "vectors",
                    "metrics",
                    "metadata",
                },  # Load records more efficiently
            )
            return metric.apply(map(record_class.parse_obj, records))

        return self.__dao__.compute_metric(
            metric_id=metric.id,
            metric_params=metric_params,
            dataset=dataset,
            query=query,
        )

    def annotated_records(self, dataset: ServiceDataset) -> int:
        """Return the number of annotated records for a dataset"""
        results = self.__dao__.search_records(
            dataset,
            size=0,
            search=DaoRecordsSearch(query=ServiceBaseRecordsQuery(has_annotation=True)),
        )
        return results.total

    def total_records(self, dataset: ServiceDataset) -> int:
        """Return the total number of records for a given dataset"""
        return self.__dao__.search_records(dataset, size=0).total
