# dbnomics-solr: Index DBnomics data into Apache Solr for full-text and faceted search.
# By: Christophe Benz <christophe.benz@cepremap.org>
#
# Copyright (C) 2017-2020 Cepremap
# https://git.nomics.world/dbnomics/dbnomics-solr
#
# dbnomics-solr is free software; you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# dbnomics-solr is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

"""Solr client for DBnomics."""

from dataclasses import dataclass
from datetime import date, datetime, time, timezone
from typing import Any, Iterator, Optional, Tuple
from urllib.parse import quote

import daiquiri
import orjson
import pysolr
import requests
from dbnomics_data_model.model import DatasetCode, DatasetMetadata, ProviderMetadata, SeriesMetadata
from dbnomics_data_model.storage.adapters.filesystem import FileSystemStorage, SeriesJsonLinesOffset
from slugify import slugify
from solrq import Q

from .hash_utils import compute_dir_hash

logger = daiquiri.getLogger(__name__)


__all__ = ["DBnomicsSolrClient", "format_date_for_solr"]


class DBnomicsSolrClient:
    """Solr client for DBnomics."""

    def __init__(self, solr_url: str, *, indexed_at: Optional[datetime] = None, timeout: int = 60):
        """Init DBnomicsSolrClient from Solr URL."""
        if solr_url.endswith("/"):
            solr_url = solr_url[:-1]
        self._solr = pysolr.Solr(solr_url, timeout=timeout)
        self._solr_url = solr_url
        if indexed_at is None:
            indexed_at = datetime.now(timezone.utc)
        self._indexed_at = indexed_at

    def commit(self):
        """Commit changes to Solr."""
        self._solr.commit()

    def delete_obsolete_series(self, provider_code, desired_dataset_codes):
        """Delete obsolete series from Solr index.

        Obsolete series are those with a different "indexed_at" property.
        """

        def build_query():
            """Build Solr query matching obsolete documents.

            Hack: build Solr query half with solrq, half manually, because I could not
            achieve working with parentheses generated by solrq when doing
            Q(...) & ~Q(...) => (...) AND (!...)
            Solr did not work with the "!" inside parentheses.
            """
            q1 = Q(type="series", provider_code=provider_code, dataset_code=dataset_code)
            q2 = Q(indexed_at=indexed_at)
            return f"{q1} AND NOT {q2}"

        indexed_at = self._indexed_at

        for dataset_code in desired_dataset_codes:
            query = build_query()
            nb_obsolete_docs = self._solr.search(query).hits
            if nb_obsolete_docs > 0:
                logger.info(
                    "Deleting %d series of dataset %r with indexed_at != %r...",
                    nb_obsolete_docs,
                    dataset_code,
                    format_date_for_solr(indexed_at),
                )
                self._solr.delete(q=query, commit=False)

    def delete_by_provider_code(self, provider_code: str):
        """Delete all documents related to a provider from Solr index.

        Deletes provider document but also dataset and series documents.
        """
        logger.info("Deleting provider document...", provider_code=provider_code)
        self._solr.delete(q=Q(type="provider", code=provider_code))

        logger.info("Deleting dataset documents...", provider_code=provider_code)
        self._solr.delete(q=Q(type="dataset", provider_code=provider_code))

        logger.info("Deleting series documents...", provider_code=provider_code)
        self._solr.delete(q=Q(type="series", provider_code=provider_code))

    def delete_by_provider_slug(self, provider_slug: str):
        """Delete all documents related to a provider from Solr index.

        Deletes provider document but also dataset and series documents.
        """
        try:
            provider_solr = self.query_provider_by_slug(provider_slug, strict=False)
        except ProviderNotFound:
            logger.error(
                "Could not find the Solr document for the provider, skipping deletion", provider_slug=provider_slug
            )
            return
        provider_code = provider_solr.get("code")
        if provider_code is None:
            logger.error(
                "Could not find the provider code in the Solr document for the provider, skipping deletion",
                provider_slug=provider_slug,
            )
            return
        self.delete_by_provider_code(provider_code)

    def index_dataset(
        self, dataset_code: str, *, force: bool = False, provider_metadata: ProviderMetadata, storage: FileSystemStorage
    ) -> None:
        """Index a dataset.

        If the dataset found in storage is already indexed in this state, skip it.
        To implement this, a directory hash is computed and stored in Solr index.
        """
        provider_code = provider_metadata.code
        is_already_indexed, dir_hash = self._is_dataset_already_indexed(provider_code, dataset_code, storage)
        if is_already_indexed and not force:
            raise DatasetAlreadyIndexed(provider_code=provider_code, dataset_code=dataset_code, dir_hash=dir_hash)

        json_docs_iter = (
            serialize_json(without_none_values(solr_doc))
            for solr_doc in self._iter_docs_for_dataset(
                dataset_code, storage=storage, provider_metadata=provider_metadata, dir_hash=dir_hash
            )
        )
        self._index_json_lines(json_docs_iter)

    def index_provider(self, provider_metadata: ProviderMetadata) -> None:
        """Index a provider."""
        provider_code = provider_metadata.code
        provider_solr = without_none_values(
            {
                "id": provider_code,
                "indexed_at": format_date_for_solr(self._indexed_at),
                "type": "provider",
                "code": provider_code,
                "slug": slugify(provider_code),
                "name": provider_metadata.name,
                "region": provider_metadata.region,
                "terms_of_use": provider_metadata.terms_of_use,
                "website": provider_metadata.website,
            }
        )
        self._solr.add(docs=[provider_solr])

    def query_provider_by_slug(self, provider_slug: str, *, strict: bool = True) -> dict:
        """Query a provider document in Solr index.

        Raise ProviderNotFound if not found.

        If `strict`, raise a `DuplicateDocuments` error if more than one document is found.
        Otherwise, return the first matching document.
        This allows, for example, the deletion of the documents when there are duplicates.
        It seems that sometimes after an upgrade of Solr, the `uniqueKey` setting of `schema.xml` is ignored
        when adding documents, resulting in many documents having the same slug.
        """
        query = Q(type="provider", slug=provider_slug)
        result = self._solr.search(query)
        if strict and result.hits > 1:
            raise DuplicateDocuments(query=query)
        if not result:
            raise ProviderNotFound(provider_slug=provider_slug)
        dataset_solr = result.docs[0]
        return dataset_solr

    def query_dataset(self, provider_code: str, dataset_code: str, *, strict: bool = True) -> dict:
        """Query a dataset document in Solr index.

        Raise DatasetNotFound if not found.

        If `strict`, raise a `DuplicateDocuments` error if more than one document is found.
        Otherwise, return the first matching document.
        This allows, for example, the deletion of the documents when there are duplicates.
        It seems that sometimes after an upgrade of Solr, the `uniqueKey` setting of `schema.xml` is ignored
        when adding documents, resulting in many documents having the same slug.
        """
        query = Q(type="dataset", provider_code=provider_code, code=dataset_code)
        result = self._solr.search(query)
        if strict and result.hits > 1:
            raise DuplicateDocuments(query=query)
        if not result:
            raise DatasetNotFound(provider_code=provider_code, dataset_code=dataset_code)
        dataset_solr = result.docs[0]
        return dataset_solr

    def _iter_series_metadata_and_offsets(
        self, dataset_code: DatasetCode, *, storage: FileSystemStorage
    ) -> Iterator[tuple[SeriesMetadata, Optional[SeriesJsonLinesOffset]]]:
        dataset_storage_variant = storage.get_storage_variant(dataset_code)
        if dataset_storage_variant == "jsonl":
            dataset_json = storage.load_dataset_json(dataset_code)
            for series, offset in storage.iter_series_jsonl_variant(
                dataset_code,
                dataset_json=dataset_json,
                include_observations=False,
                on_error="log",
            ):
                yield (series.metadata, offset)
            return

        # Offsets are only used with the JSON Lines storage variant.
        # Set them to None otherwise.
        for series in storage.iter_dataset_series(dataset_code, include_observations=False, on_error="log"):
            yield (series.metadata, None)

    def _index_json_lines(self, docs_iter):
        """Index JSON lines from iterator in stream mode.

        pysolr doesn't support this so let's call directly a specific endpoint.
        Cf https://lucene.apache.org/solr/guide/transforming-and-indexing-custom-json.html#TransformingandIndexingCustomJSON-MultipledocumentsinaSinglePayload # noqa
        Cf https://requests.readthedocs.io/en/master/user/advanced/#chunk-encoded-requests # noqa
        """
        response = requests.post(f"{self._solr_url}/update/json/docs?commit=true", data=docs_iter)  # type: ignore
        if not response.ok:
            response_data = response.json()
            error = response_data.get("error")
            raise IndexationError(error=error)

    def _iter_docs_for_dataset(
        self,
        dataset_code: DatasetCode,
        *,
        dir_hash: str,
        provider_metadata: ProviderMetadata,
        storage: FileSystemStorage,
    ):
        """Yield Solr documents for dataset and series for a specific dataset."""
        dataset_metadata = storage.load_dataset_metadata(dataset_code)

        dataset_solr = build_dataset_solr(provider_metadata, dataset_metadata, self._indexed_at, dir_hash)

        for nb_series_produced, (series_metadata, series_jsonl_offset) in enumerate(
            self._iter_series_metadata_and_offsets(dataset_code, storage=storage), start=1
        ):
            if nb_series_produced % 1000 == 0:
                logger.debug("Produced %d series so far...", nb_series_produced)
            series_solr = build_series_solr(
                provider_metadata,
                dataset_metadata,
                series_metadata,
                indexed_at=self._indexed_at,
                series_jsonl_offset=series_jsonl_offset,
            )
            yield series_solr
            dataset_solr["nb_series"] += 1

        logger.debug(
            "Produced a total of %d Solr documents for the series of the dataset %r",
            dataset_solr["nb_series"],
            dataset_code,
        )

        # Obsolete series of the current dataset will be removed later.

        yield dataset_solr

    def _is_dataset_already_indexed(self, provider_code: str, dataset_code: str, storage) -> Tuple[bool, str]:
        """Return whether the dataset is already indexed, and the dataset dir hash."""
        dataset_dir = storage.get_dataset_dir(dataset_code)
        storage_dir_hash = compute_dir_hash(dataset_dir)

        try:
            dataset_solr = self.query_dataset(provider_code, dataset_code)
        except DatasetNotFound:
            return (False, storage_dir_hash)

        dir_hash = dataset_solr.get("dir_hash")
        return (storage_dir_hash == dir_hash, storage_dir_hash)


def build_dataset_solr(
    provider_metadata: ProviderMetadata, dataset_metadata: DatasetMetadata, indexed_at, dir_hash: str
) -> Any:
    """Build a JSON object for a dataset, following the Solr schema."""
    provider_code = provider_metadata.code
    dataset_code = dataset_metadata.code

    updated_at = dataset_metadata.updated_at
    updated_at_str = None if updated_at is None else format_date_for_solr(updated_at)

    dataset_solr = {
        "id": f"{provider_code}/{dataset_code}",
        "indexed_at": format_date_for_solr(indexed_at),
        "type": "dataset",
        "code": dataset_code,
        "provider_code": provider_code,
        "provider_name": provider_metadata.name,
        "name": dataset_metadata.name,
        "description": dataset_metadata.description,
        "dir_hash": dir_hash,
        "updated_at": updated_at_str,
        "nb_series": 0,  # Will be incremented if dataset has series.
    }

    return dataset_solr


def build_series_solr(
    provider_metadata: ProviderMetadata,
    dataset_metadata: DatasetMetadata,
    series_metadata: SeriesMetadata,
    *,
    indexed_at: datetime,
    series_jsonl_offset: Optional[SeriesJsonLinesOffset],
) -> Any:
    """Build a JSON object for a series, following the Solr schema."""
    provider_code = provider_metadata.code
    dataset_code = dataset_metadata.code
    series_code = series_metadata.code
    series_solr = {
        "id": f"{provider_code}/{dataset_code}/{series_code}",
        "indexed_at": format_date_for_solr(indexed_at),
        "type": "series",
        "code": series_code,
        "name": series_metadata.name,
        "provider_code": provider_code,
        "provider_name": provider_metadata.name,
        "dataset_id": f"{provider_code}/{dataset_code}",
        "dataset_code": dataset_code,
        "dataset_name": dataset_metadata.name,
    }
    if series_jsonl_offset is not None:
        series_solr["series_jsonl_offset"] = series_jsonl_offset

    for dimension_code, dimension_value_code in series_metadata.dimensions.items():
        # Index dimension codes to compute facets.
        series_solr[f"dimension_value_code.{quote(dimension_code)}"] = dimension_value_code

    dimensions_values_labels = list(iter_dimensions_values_labels(series_metadata, dataset_metadata=dataset_metadata))
    if dimensions_values_labels:
        series_solr["dimensions_values_labels"] = dimensions_values_labels

    return series_solr


def iter_dimensions_values_labels(
    series_metadata: SeriesMetadata, *, dataset_metadata: DatasetMetadata
) -> Iterator[str]:
    for dimension_code, dimension_value_code in series_metadata.dimensions.items():
        dimension_def = dataset_metadata.find_dimension_by_code(dimension_code)
        if dimension_def is None:
            continue
        dimension_value_def = dimension_def.find_value_by_code(dimension_value_code)
        if dimension_value_def is None:
            continue
        dimension_value_label = dimension_value_def.label
        if dimension_value_label is not None:
            yield dimension_value_label


class DbnomicsSolrException(Exception):
    pass


@dataclass
class DatasetAlreadyIndexed(DbnomicsSolrException):
    provider_code: str
    dataset_code: str
    dir_hash: str


@dataclass
class DatasetNotFound(DbnomicsSolrException):
    provider_code: str
    dataset_code: str


@dataclass
class DuplicateDocuments(DbnomicsSolrException):
    query: Q


@dataclass
class ProviderNotFound(DbnomicsSolrException):
    provider_slug: Optional[str] = None


@dataclass
class IndexationError(DbnomicsSolrException):
    error: Optional[dict]

    def __str__(self):
        return str(self.error)


def format_date_for_solr(d: date) -> str:
    """Format a `date` or a `datetime` for Solr.

    If `d` is a `date` (and not a `datetime`), assume the time is `00:00:00` in the UTC timezone.

    Solr supports a subset of ISO-8601 and requires the datetime to be expressed in the UTC timezone
    with the "Z" suffix notation instead of the "+00:00" generated by Python.

    If the timezone of a datetime is not specified, UTC is assumed and a warning is logged.

    >>> from datetime import timedelta
    >>> format_date_for_solr(date(2000, 1, 1))
    '2000-01-01T00:00:00Z'
    >>> format_date_for_solr(datetime(2000, 1, 1, 5, 5, 5))
    '2000-01-01T05:05:05Z'
    >>> format_date_for_solr(datetime(2000, 1, 1, 5, 5, 5, tzinfo=timezone.utc))
    '2000-01-01T05:05:05Z'
    >>> format_date_for_solr(datetime(2000, 1, 1, 5, 5, 5, tzinfo=timezone(timedelta(hours=2))))
    '2000-01-01T03:05:05Z'
    """
    if not isinstance(d, datetime):
        return datetime.combine(d, time.min).isoformat() + "Z"

    if d.tzinfo is None:
        logger.warning("No timezone defined in %r, assuming it is UTC", d)
        d = d.replace(tzinfo=timezone.utc)
    else:
        d = d.astimezone(timezone.utc)

    return d.isoformat().removesuffix("+00:00") + "Z"


def serialize_json(data: Any) -> bytes:
    return orjson.dumps(data)


def without_none_values(d):
    """Return a copy of d without None values."""
    return {k: v for k, v in d.items() if v is not None}
