#  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.

import dataclasses
from typing import List, Optional, Type

from fastapi import Depends

from argilla_server import telemetry
from argilla_server.commons.config import TasksFactory
from argilla_server.commons.models import TaskStatus
from argilla_server.daos.backend.base import WrongLogDataError
from argilla_server.daos.records import DatasetRecordsDAO
from argilla_server.errors import BulkDataError, ForbiddenOperationError
from argilla_server.models import User
from argilla_server.policies import DatasetPolicy, is_authorized
from argilla_server.services.datasets import ServiceDataset
from argilla_server.services.search.model import ServiceBaseRecordsQuery
from argilla_server.services.tasks.commons import ServiceRecord


@dataclasses.dataclass
class DeleteRecordsOut:
    processed: int = 0
    discarded: int = 0
    deleted: int = 0


class RecordsStorageService:
    _INSTANCE: "RecordsStorageService" = None

    @classmethod
    def get_instance(
        cls,
        dao: DatasetRecordsDAO = Depends(DatasetRecordsDAO.get_instance),
    ) -> "RecordsStorageService":
        if not cls._INSTANCE:
            cls._INSTANCE = cls(dao)
        return cls._INSTANCE

    def __init__(self, dao: DatasetRecordsDAO):
        self.__dao__ = dao

    async def store_records(
        self,
        dataset: ServiceDataset,
        records: List[ServiceRecord],
        record_type: Type[ServiceRecord],
    ) -> int:
        """Store a set of records"""
        await telemetry.track_bulk(task=dataset.task, records=len(records))

        metrics = TasksFactory.get_task_metrics(dataset.task)
        if metrics:
            for record in records:
                record.metrics = metrics.record_metrics(record)

        try:
            return self.__dao__.add_records(
                dataset=dataset,
                records=records,
                record_class=record_type,
            )
        except WrongLogDataError as ex:
            raise BulkDataError(
                detail=f"Cannot log data in dataset {dataset.id}",
                errors=ex.errors,
            )

    async def delete_records(
        self,
        user: User,
        dataset: ServiceDataset,
        query: Optional[ServiceBaseRecordsQuery] = None,
        mark_as_discarded: bool = False,
    ) -> DeleteRecordsOut:
        processed, discarded, deleted = None, None, None
        if mark_as_discarded:
            processed, discarded = await self.__dao__.update_records_by_query(
                dataset,
                query=query,
                status=TaskStatus.discarded,
            )
        else:
            if not await is_authorized(user, DatasetPolicy.delete_records(dataset)):
                raise ForbiddenOperationError(
                    "You don't have the necessary permissions to delete records on this dataset. "
                    "Only administrators can delete datasets"
                )

            processed, deleted = await self.__dao__.delete_records_by_query(dataset, query=query)

        return DeleteRecordsOut(
            processed=processed or 0,
            discarded=discarded or 0,
            deleted=deleted or 0,
        )

    async def update_record(
        self,
        dataset: ServiceDataset,
        record: ServiceRecord,
        **data,
    ) -> ServiceRecord:
        if data.get("metadata"):
            record.metadata = {
                **(record.metadata or {}),
                **data.pop("metadata"),
            }

        new_record = record.copy(update=data)  # Data validation
        await self.__dao__.update_record(
            dataset=dataset,
            record=record,
        )
        return new_record
