"""Solr client for DBnomics."""

from __future__ import annotations

import logging
import time as t
from dataclasses import dataclass
from datetime import date, datetime, time, timezone
from typing import Any, Iterable, Iterator, Optional, cast
from urllib.parse import quote

import daiquiri
import orjson
import requests
from dbnomics_data_model import ProviderCode
from dbnomics_data_model.model import DatasetCode, DatasetMetadata, ProviderMetadata, SeriesMetadata
from dbnomics_data_model.storage.adapters.filesystem import FileSystemStorage, SeriesJsonLinesOffset
from humanfriendly import format_timespan
from humanfriendly.text import pluralize
from pysolr import Results, Solr
from slugify import slugify
from solrq import Q
from tenacity import Retrying, after_log, retry_if_not_exception_type, stop_after_attempt, wait_random_exponential

from dbnomics_solr.errors import (
    DatasetAlreadyIndexed,
    DatasetNotFound,
    DuplicateDocuments,
    IndexationError,
    InvalidSolrDocument,
    ProviderNotFound,
)
from dbnomics_solr.text_utils import english_join
from dbnomics_solr.types import SolrDoc

from .hash_utils import compute_dir_hash

__all__ = ["DBnomicsSolrClient", "format_date_for_solr"]

logger = daiquiri.getLogger(__name__)


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

    def __init__(
        self,
        solr_url: str,
        *,
        debug_progress_nb_series: int = 10000,
        dirhash_jobs: int = 1,
        indexed_at: Optional[datetime] = None,
        retry_num_attempts: Optional[int] = None,
        timeout: int = 60,
    ):
        """Init DBnomicsSolrClient from Solr URL."""
        self._debug_progress_nb_series = debug_progress_nb_series
        self._dirhash_jobs = dirhash_jobs

        if solr_url.endswith("/"):
            solr_url = solr_url[:-1]
        self._solr = 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

        if retry_num_attempts is None:
            retry_num_attempts = 3
        self._retry_num_attempts = retry_num_attempts

        self._timeout = timeout

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

    def delete_obsolete_series(self, provider_code: ProviderCode, dataset_codes: Iterable[DatasetCode]):
        """Delete obsolete series related to that dataset.

        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 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_dataset_docs(self, provider_code: ProviderCode, dataset_code: DatasetCode):
        """Delete all the documents related to that dataset.

        This includes the document representing the dataset, but also all the ones representing
        the series of that dataset.
        """
        dataset_id = f"{provider_code}/{dataset_code}"

        logger.info("Deleting the dataset document...", dataset_id=dataset_id)
        self._solr.delete(q=Q(type="dataset", id=dataset_id))

        logger.info("Deleting series documents related to that dataset...", dataset_id=dataset_id)
        self._solr.delete(q=Q(type="series", dataset_id=dataset_id))

    def delete_provider_docs(self, provider_code: ProviderCode):
        """Delete all the documents related to that provider.

        This includes the document representing that provider, but also the one representing the datasets
        of that provider, and the ones representing the series of the datasets of that provider.
        """
        logger.info("Deleting the provider document...", provider_code=provider_code)
        self._solr.delete(q=Q(type="provider", code=provider_code))

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

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

    def search_dataset_docs(self, provider_code: ProviderCode, dataset_code: DatasetCode) -> DBnomicsSolrResults:
        """Search and return the documents related to that dataset.

        This includes the document representing the dataset, but also all the ones representing
        the series of that dataset.
        """
        dataset_id = f"{provider_code}/{dataset_code}"
        dataset_results = self._solr.search(q=Q(type="dataset", id=dataset_id))
        series_results = self._solr.search(q=Q(type="series", dataset_id=dataset_id))
        return DBnomicsSolrResults(dataset_results=dataset_results, series_results=series_results)

    def search_extra_datasets(
        self,
        provider_code: ProviderCode,
        *,
        storage: FileSystemStorage,
        storage_dataset_codes: Optional[Iterable[DatasetCode]],
    ) -> set[DatasetCode]:
        """Search the extra datasets of that provider.

        These are the ones that are in Solr but not in the storage.
        """
        if storage_dataset_codes is None:
            storage_dataset_codes = sorted(storage.iter_dataset_codes(on_error="log"))
        provider_datasets = self.search_provider_datasets(provider_code)
        solr_dataset_codes = {result["code"] for result in provider_datasets}
        return solr_dataset_codes - set(storage_dataset_codes)

    def search_provider_datasets(self, provider_code: ProviderCode) -> Results:
        """Search and return the documents representing the datasets of that provider."""
        return self._solr.search(q=Q(type="dataset", provider_code=provider_code))

    def search_provider_docs(self, provider_code: ProviderCode) -> DBnomicsSolrResults:
        """Search and return the documents related to that provider.

        This includes the document representing that provider, but also the one representing the datasets
        of that provider, and the ones representing the series of the datasets of that provider.
        """
        provider_results = self._solr.search(q=Q(type="provider", code=provider_code))
        dataset_results = self._solr.search(q=Q(type="dataset", provider_code=provider_code))
        series_results = self._solr.search(q=Q(type="series", provider_code=provider_code))
        return DBnomicsSolrResults(
            provider_results=provider_results, dataset_results=dataset_results, series_results=series_results
        )

    def find_provider_code_from_slug(self, provider_slug: str) -> ProviderCode:
        """Return the provider code corresponding to the given slug."""
        provider_solr = self.search_provider_by_slug(provider_slug, strict=False)
        provider_code = provider_solr.get("code")
        if provider_code is None:
            raise InvalidSolrDocument(
                f'Could not find the "code" property in the Solr document of the provider '
                f"which slug is {provider_slug!r}",
                solr_document=provider_solr,
            )
        return provider_code

    def index_dataset(
        self,
        dataset_code: DatasetCode,
        *,
        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.
        """
        provider_code = provider_metadata.code
        storage_dir_hash = self._compute_dir_hash(dataset_code, storage=storage)
        solr_dir_hash = self._fetch_dir_hash_from_solr(provider_code, dataset_code)
        if storage_dir_hash == solr_dir_hash:
            if force:
                logger.debug(
                    "The hash of the directory of the dataset %r (%r) is identical to the one stored in Solr, "
                    'but the dataset will be reindexed due to the "force" option.',
                    dataset_code,
                    storage_dir_hash,
                )
            else:
                raise DatasetAlreadyIndexed(
                    provider_code=provider_code, dataset_code=dataset_code, dir_hash=storage_dir_hash
                )
        else:
            logger.debug(
                "The hash of the directory of the dataset %r (%r) is different from the one stored in Solr (%r), "
                "so the dataset will be reindexed.",
                dataset_code,
                storage_dir_hash,
                solr_dir_hash,
            )

        # Sometimes Solr is too busy and responds an error, so add a retry strategy
        # instead of failing the dataset indexation.
        for attempt in Retrying(
            after=after_log(cast(logging.Logger, logger), logging.ERROR),
            reraise=True,
            retry=retry_if_not_exception_type(DatasetAlreadyIndexed),
            stop=stop_after_attempt(self._retry_num_attempts),
            wait=wait_random_exponential(multiplier=1, max=60 * 5),
        ):
            with attempt:
                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=storage_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 search_dataset(self, provider_code: ProviderCode, dataset_code: DatasetCode, *, strict: bool = True) -> SolrDoc:
        """Search a dataset document in Solr.

        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 search_provider_by_slug(self, provider_slug: str, *, strict: bool = True) -> SolrDoc:
        """Search a provider document in Solr.

        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 _compute_dir_hash(self, dataset_code: DatasetCode, *, storage: FileSystemStorage) -> str:
        logger.debug("Computing the hash of the directory of the dataset %r...", dataset_code)
        dataset_dir = storage.get_dataset_dir(dataset_code)
        return compute_dir_hash(dataset_dir, jobs=self._dirhash_jobs)

    def _fetch_dir_hash_from_solr(self, provider_code: ProviderCode, dataset_code: DatasetCode) -> Optional[str]:
        try:
            dataset_solr = self.search_dataset(provider_code, dataset_code)
        except DatasetNotFound:
            return None

        return dataset_solr.get("dir_hash")

    def _iter_series_metadata_and_offsets(
        self, dataset_code: DatasetCode, *, storage: FileSystemStorage
    ) -> Iterator[tuple[SeriesMetadata, Optional[SeriesJsonLinesOffset]]]:
        storage_variant = storage.get_storage_variant(dataset_code)
        if storage_variant == "jsonl":
            dataset_json = storage.load_dataset_json(dataset_code, storage_variant=storage_variant)
            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
        """
        response = requests.post(
            f"{self._solr_url}/update/json/docs?commit=true", data=docs_iter, timeout=self._timeout
        )
        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_id = f"{provider_metadata.code}/{dataset_code}"

        logger.debug("Loading dataset metadata...", dataset_id=dataset_id)
        dataset_metadata = storage.load_dataset_metadata(dataset_code)

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

        logger.debug("Loading series of dataset...", dataset_id=dataset_id)

        t0 = t.time()
        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 % self._debug_progress_nb_series == 0:
                logger.debug("Loaded %d series so far...", nb_series_produced, dataset_id=dataset_id)
            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(
            "Loaded a total of %d series in %s",
            dataset_solr["nb_series"],
            format_timespan(t.time() - t0),
            dataset_id=dataset_id,
        )

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

        yield dataset_solr


@dataclass
class DBnomicsSolrResults:
    dataset_results: Results
    series_results: Results
    provider_results: Optional[Results] = None

    def format_hits_by_type(self) -> str:
        items = []
        if self.provider_results is not None:
            items.append(pluralize(self.provider_results.hits, "provider document"))
        items.extend(
            [
                pluralize(self.dataset_results.hits, "dataset document"),
                pluralize(self.series_results.hits, "series document"),
            ]
        )
        return english_join(items)

    def hit_sum(self) -> int:
        provider_hits = self.provider_results.hits if self.provider_results is not None else 0
        return provider_hits + self.dataset_results.hits + self.series_results.hits


def build_dataset_solr(
    provider_metadata: ProviderMetadata, dataset_metadata: DatasetMetadata, indexed_at, dir_hash: str
) -> SolrDoc:
    """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],
) -> SolrDoc:
    """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: SolrDoc = {
        "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


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}
