#  coding=utf-8
#  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 datetime
from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Type

from fastapi import Depends

from argilla_server.daos.backend import GenericElasticEngineBackend
from argilla_server.daos.backend.base import ClosedIndexError, IndexNotFoundError
from argilla_server.daos.backend.generic_elastic import PaginatedSortInfo
from argilla_server.daos.backend.search.model import BaseRecordsQuery, SortableField
from argilla_server.daos.models.datasets import DatasetDB
from argilla_server.daos.models.records import DaoRecordsSearch, DaoRecordsSearchResults, RecordDB
from argilla_server.errors import ClosedDatasetError, MissingDatasetRecordsError


class DatasetRecordsDAO:
    """Datasets records DAO"""

    _INSTANCE = None

    @classmethod
    def get_instance(
        cls,
        es: GenericElasticEngineBackend = Depends(GenericElasticEngineBackend.get_instance),
    ) -> "DatasetRecordsDAO":
        """
        Creates a dataset records dao instance

        Parameters
        ----------
        es:
            The elasticserach wrapper dependency

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

    def __init__(self, es: GenericElasticEngineBackend):
        self._es = es

    def init(self):
        """Initializes dataset records dao. Used on app startup"""
        pass

    def add_records(
        self,
        dataset: DatasetDB,
        records: List[RecordDB],
        record_class: Type[RecordDB],
    ) -> int:
        """
        Add records to dataset

        Parameters
        ----------
        dataset:
            The dataset
        records:
            The list of records
        record_class:
            Record class used to convert records to
        Returns
        -------
            The number of failed records

        """

        now = datetime.datetime.utcnow()
        documents = []
        metadata_values = {}

        mapping = self._es.get_schema(id=dataset.id)
        exclude_fields = [
            name for name in record_class.schema()["properties"] if name not in mapping["mappings"]["properties"]
        ]

        vectors_configuration = {}
        for record in records:
            metadata_values.update(record.metadata or {})

            db_record = record_class.parse_obj(record)
            db_record.last_updated = now

            record_dict = db_record.dict(exclude_none=True, exclude=set(exclude_fields))

            if record.vectors is not None:
                # TODO: Create embeddings config by settings
                for (
                    vector_name,
                    vector_data_mapping,
                ) in record.vectors.items():
                    vector_dimension = vectors_configuration.get(vector_name, None)

                    if vector_dimension is None:
                        dimension = len(vector_data_mapping.value)
                        vectors_configuration[vector_name] = dimension

            documents.append(record_dict)

        self._es.create_dataset(
            id=dataset.id,
            task=dataset.task,
            metadata_values=metadata_values,
            vectors_cfg=vectors_configuration,
        )

        return self._es.add_dataset_records(
            id=dataset.id,
            documents=documents,
        )

    def compute_metric(
        self,
        dataset: DatasetDB,
        metric_id: str,
        metric_params: Dict[str, Any] = None,
        query: Optional[BaseRecordsQuery] = None,
    ):
        return self._es.compute_metric(
            id=dataset.id,
            metric_id=metric_id,
            query=query,
            params=metric_params,
        )

    def search_records(
        self,
        dataset: DatasetDB,
        search: Optional[DaoRecordsSearch] = None,
        size: int = 100,
        record_from: int = 0,
        exclude_fields: List[str] = None,
        highligth_results: bool = True,
    ) -> DaoRecordsSearchResults:
        try:
            search = search or DaoRecordsSearch()

            total, records = self._es.search_records(
                id=dataset.id,
                query=search.query,
                sort=search.sort,
                record_from=record_from,
                size=size,
                exclude_fields=exclude_fields,
                enable_highlight=highligth_results,
            )
            if isinstance(total, dict):
                total = total["value"]
            return DaoRecordsSearchResults(
                total=total,
                records=records,
            )
        except ClosedIndexError:
            raise ClosedDatasetError(dataset.name)
        except IndexNotFoundError:
            raise MissingDatasetRecordsError(f"No records index found for dataset {dataset.name}")

    def scan_dataset(
        self,
        dataset: DatasetDB,
        search: Optional[DaoRecordsSearch] = None,
        limit: Optional[int] = 1000,
        id_from: Optional[str] = None,
        include_fields: Optional[Set[str]] = None,
        exclude_fields: Optional[Set[str]] = None,
    ) -> Iterable[Dict[str, Any]]:
        """
        Iterates over a dataset records

        Parameters
        ----------
        dataset:
            The dataset
        search:
            The search parameters. Optional
        limit:
            Batch size to extract, only works if an `id_from` is provided
        id_from:
            From which ID should we start iterating
        include_fields:
            A set of record fields to retrieve. Wildcard are allowed
        exclude_fields:
            A set of record fields to exclude. Wildcard are allowed

        Returns
        -------
            An iterable over found dataset records
        """

        search = search or DaoRecordsSearch()
        next_search_params = [id_from] if id_from else None
        paginated_sort = PaginatedSortInfo(
            sort_by=[SortableField(id="id")], shuffle=search.sort.shuffle, next_search_params=next_search_params
        )

        return self._es.scan_records(
            id=dataset.id,
            query=search.query,
            sort=paginated_sort,
            limit=limit,
            include_fields=list(include_fields) if include_fields else None,
            exclude_fields=list(exclude_fields) if exclude_fields else None,
        )

    async def delete_records_by_query(
        self,
        dataset: DatasetDB,
        query: Optional[BaseRecordsQuery] = None,
    ) -> Tuple[int, int]:
        total, deleted = await self._es.delete_records_by_query(
            id=dataset.id,
            query=query,
        )
        return total, deleted

    async def update_records_by_query(
        self,
        dataset: DatasetDB,
        query: Optional[BaseRecordsQuery] = None,
        **content,
    ) -> Tuple[int, int]:
        total, updated = await self._es.update_records_content(
            id=dataset.id,
            content=content,
            query=query,
        )
        return total, updated

    async def update_record(
        self,
        dataset: DatasetDB,
        record: RecordDB,
    ):
        self._es.update_record(
            dataset_id=dataset.id,
            record_id=record.id,
            content=record.dict(exclude_none=True),
        )

    async def get_record_by_id(
        self,
        dataset: DatasetDB,
        id: str,
    ) -> Optional[Dict[str, Any]]:
        return self._es.find_record_by_id(
            dataset_id=dataset.id,
            record_id=id,
        )
