from __future__ import annotations

from datetime import datetime
from itertools import chain
from math import ceil

from maggma.builders import Builder
from maggma.core import Store
from maggma.utils import grouper

from emmet.builders.settings import EmmetBuildSettings
from emmet.core.molecules.summary import MoleculeSummaryDoc
from emmet.core.utils import jsanitize


from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from collections.abc import Iterable, Iterator
    from typing import Any

# from monty.serialization import loadfn, dumpfn


__author__ = "Evan Spotte-Smith"

SETTINGS = EmmetBuildSettings()


class SummaryBuilder(Builder):
    """
    The SummaryBuilder collects all property documents and gathers their properties
    into a single MoleculeSummaryDoc

    The process is as follows:
        1. Gather MoleculeDocs by formula
        2. For each doc, grab the relevant property docs
        3. Convert property docs to MoleculeSummaryDoc
    """

    def __init__(
        self,
        molecules: Store,
        charges: Store,
        spins: Store,
        bonds: Store,
        multipoles: Store,
        metal_binding: Store,
        orbitals: Store,
        redox: Store,
        thermo: Store,
        vibes: Store,
        summary: Store,
        query: dict | None = None,
        settings: EmmetBuildSettings | None = None,
        **kwargs,
    ):
        self.molecules = molecules
        self.charges = charges
        self.spins = spins
        self.bonds = bonds
        self.multipoles = multipoles
        self.metal_binding = metal_binding
        self.orbitals = orbitals
        self.redox = redox
        self.thermo = thermo
        self.vibes = vibes
        self.summary = summary
        self.query = query if query else dict()
        self.settings = EmmetBuildSettings.autoload(settings)
        self.kwargs = kwargs

        super().__init__(
            sources=[
                molecules,
                charges,
                spins,
                bonds,
                multipoles,
                metal_binding,
                orbitals,
                redox,
                thermo,
                vibes,
            ],
            targets=[summary],
            **kwargs,
        )
        # Uncomment in case of issue with mrun not connecting automatically to collections
        # for i in [
        #     self.molecules,
        #     self.charges,
        #     self.spins,
        #     self.bonds,
        #     self.multipoles,
        #     self.metal_binding,
        #     self.orbitals,
        #     self.redox,
        #     self.thermo,
        #     self.vibes,
        #     self.summary
        # ]:
        #     try:
        #         i.connect()
        #     except Exception as e:
        #         print("Could not connect,", e)

    def ensure_indexes(self):
        """
        Ensures indices on the collections needed for building
        """

        # Search index for molecules
        self.molecules.ensure_index("molecule_id")
        self.molecules.ensure_index("last_updated")
        self.molecules.ensure_index("task_ids")
        self.molecules.ensure_index("formula_alphabetical")
        self.molecules.ensure_index("species_hash")

        # Search index for charges
        self.charges.ensure_index("molecule_id")
        self.charges.ensure_index("method")
        self.charges.ensure_index("task_id")
        self.charges.ensure_index("solvent")
        self.charges.ensure_index("lot_solvent")
        self.charges.ensure_index("property_id")
        self.charges.ensure_index("last_updated")
        self.charges.ensure_index("formula_alphabetical")

        # Search index for charges
        self.spins.ensure_index("molecule_id")
        self.spins.ensure_index("method")
        self.spins.ensure_index("task_id")
        self.spins.ensure_index("solvent")
        self.spins.ensure_index("lot_solvent")
        self.spins.ensure_index("property_id")
        self.spins.ensure_index("last_updated")
        self.spins.ensure_index("formula_alphabetical")

        # Search index for charges
        self.bonds.ensure_index("molecule_id")
        self.bonds.ensure_index("method")
        self.bonds.ensure_index("task_id")
        self.bonds.ensure_index("solvent")
        self.bonds.ensure_index("lot_solvent")
        self.bonds.ensure_index("property_id")
        self.bonds.ensure_index("last_updated")
        self.bonds.ensure_index("formula_alphabetical")

        # Search index for multipoles
        self.multipoles.ensure_index("molecule_id")
        self.multipoles.ensure_index("task_id")
        self.multipoles.ensure_index("solvent")
        self.multipoles.ensure_index("lot_solvent")
        self.multipoles.ensure_index("property_id")
        self.multipoles.ensure_index("last_updated")
        self.multipoles.ensure_index("formula_alphabetical")

        # Search index for metal_binding
        self.metal_binding.ensure_index("molecule_id")
        self.metal_binding.ensure_index("solvent")
        self.metal_binding.ensure_index("lot_solvent")
        self.metal_binding.ensure_index("property_id")
        self.metal_binding.ensure_index("last_updated")
        self.metal_binding.ensure_index("formula_alphabetical")
        self.metal_binding.ensure_index("method")

        # Search index for orbitals
        self.orbitals.ensure_index("molecule_id")
        self.orbitals.ensure_index("task_id")
        self.orbitals.ensure_index("solvent")
        self.orbitals.ensure_index("lot_solvent")
        self.orbitals.ensure_index("property_id")
        self.orbitals.ensure_index("last_updated")
        self.orbitals.ensure_index("formula_alphabetical")

        # Search index for orbitals
        self.redox.ensure_index("molecule_id")
        self.redox.ensure_index("task_id")
        self.redox.ensure_index("solvent")
        self.redox.ensure_index("lot_solvent")
        self.redox.ensure_index("property_id")
        self.redox.ensure_index("last_updated")
        self.redox.ensure_index("formula_alphabetical")

        # Search index for thermo
        self.thermo.ensure_index("molecule_id")
        self.thermo.ensure_index("task_id")
        self.thermo.ensure_index("solvent")
        self.thermo.ensure_index("lot_solvent")
        self.thermo.ensure_index("property_id")
        self.thermo.ensure_index("last_updated")
        self.thermo.ensure_index("formula_alphabetical")

        # Search index for vibrational properties
        self.vibes.ensure_index("molecule_id")
        self.vibes.ensure_index("task_id")
        self.vibes.ensure_index("solvent")
        self.vibes.ensure_index("lot_solvent")
        self.vibes.ensure_index("property_id")
        self.vibes.ensure_index("last_updated")
        self.vibes.ensure_index("formula_alphabetical")

        # Search index for molecules
        self.summary.ensure_index("molecule_id")
        self.summary.ensure_index("last_updated")
        self.summary.ensure_index("formula_alphabetical")

    def prechunk(self, number_splits: int) -> Iterable[dict]:  # pragma: no cover
        """Prechunk the builder for distributed computation"""

        temp_query = dict(self.query)
        temp_query["deprecated"] = False

        self.logger.info("Finding documents to process")
        all_mols = list(
            self.molecules.query(temp_query, [self.molecules.key, "species_hash"])
        )

        processed_docs = set([e for e in self.summary.distinct("molecule_id")])
        to_process_docs = {d[self.molecules.key] for d in all_mols} - processed_docs
        to_process_hashes = {
            d["species_hash"]
            for d in all_mols
            if d[self.molecules.key] in to_process_docs
        }

        N = ceil(len(to_process_hashes) / number_splits)

        for hash_chunk in grouper(to_process_hashes, N):
            query = dict(temp_query)
            query["species_hash"] = {"$in": list(hash_chunk)}
            yield {"query": query}

    def get_items(self) -> Iterator[list[dict]]:
        """
        Gets all items to process into summary documents.
        This does no datetime checking; relying on on whether
        task_ids are included in the summary Store

        Returns:
            generator or list relevant tasks and molecules to process into documents
        """

        self.logger.info("Summary builder started")
        self.logger.info("Setting indexes")
        self.ensure_indexes()

        # Save timestamp to mark buildtime
        self.timestamp = datetime.utcnow()

        # Get all processed molecules
        temp_query = dict(self.query)
        temp_query["deprecated"] = False

        self.logger.info("Finding documents to process")
        all_mols = list(
            self.molecules.query(temp_query, [self.molecules.key, "species_hash"])
        )

        processed_docs = set([e for e in self.summary.distinct("molecule_id")])
        to_process_docs = {d[self.molecules.key] for d in all_mols} - processed_docs
        to_process_hashes = {
            d["species_hash"]
            for d in all_mols
            if d[self.molecules.key] in to_process_docs
        }

        self.logger.info(f"Found {len(to_process_docs)} unprocessed documents")
        self.logger.info(f"Found {len(to_process_hashes)} unprocessed hashes")

        # Set total for builder bars to have a total
        self.total = len(to_process_hashes)

        for shash in to_process_hashes:
            mol_query = dict(temp_query)
            mol_query["species_hash"] = shash
            molecules = list(self.molecules.query(criteria=mol_query))

            yield molecules

    def process_item(self, items: list[dict]) -> list[dict]:
        """
        Process the tasks into a MoleculeSummaryDoc

        Args:
            tasks list[dict] : a list of MoleculeDocs in dict form

        Returns:
            [dict] : a list of new orbital docs
        """

        def _group_docs(docs: list[dict[str, Any]], by_method: bool = False):
            """Helper function to group docs by solvent"""
            grouped: dict[str, Any] = dict()

            for doc in docs:
                solvent = doc.get("solvent")
                method = doc.get("method")
                if not solvent:
                    # Need to group by solvent
                    continue
                if by_method and method is None:
                    # Trying to group by method, but no method present
                    continue

                if not by_method:
                    grouped[solvent] = doc
                else:
                    if solvent not in grouped:
                        grouped[solvent] = {method: doc}
                    else:
                        grouped[solvent][method] = doc

            return grouped

        mols = items
        shash = mols[0]["species_hash"]
        mol_ids = [m["molecule_id"] for m in mols]
        self.logger.debug(f"Processing {shash} : {mol_ids}")

        summary_docs = list()

        for mol in mols:
            mol_id = mol["molecule_id"]

            d = {
                "molecules": mol,
                "partial_charges": _group_docs(
                    list(self.charges.query({"molecule_id": mol_id})), True
                ),
                "partial_spins": _group_docs(
                    list(self.spins.query({"molecule_id": mol_id})), True
                ),
                "bonding": _group_docs(
                    list(self.bonds.query({"molecule_id": mol_id})), True
                ),
                "metal_binding": _group_docs(
                    list(self.metal_binding.query({"molecule_id": mol_id})), True
                ),
                "multipole_moments": _group_docs(
                    list(self.multipoles.query({"molecule_id": mol_id})), False
                ),
                "orbitals": _group_docs(
                    list(self.orbitals.query({"molecule_id": mol_id})), False
                ),
                "redox": _group_docs(
                    list(self.redox.query({"molecule_id": mol_id})), False
                ),
                "thermo": _group_docs(
                    list(self.thermo.query({"molecule_id": mol_id})), False
                ),
                "vibration": _group_docs(
                    list(self.vibes.query({"molecule_id": mol_id})), False
                ),
            }

            to_delete = list()

            for k, v in d.items():
                if isinstance(v, dict) and len(v) == 0:
                    to_delete.append(k)

            for td in to_delete:
                del d[td]

            # # For debugging; keep because it might be needed again
            # dumpfn(d, f"{mol_id}.json.gz")
            # break

            summary_doc = MoleculeSummaryDoc.from_docs(molecule_id=mol_id, docs=d)
            summary_docs.append(summary_doc)

        self.logger.debug(f"Produced {len(summary_docs)} summary docs for {shash}")

        return jsanitize([doc.model_dump() for doc in summary_docs], allow_bson=True)

    def update_targets(self, items: list[list[dict]]):
        """
        Inserts the new documents into the summary collection

        Args:
            items [[dict]]: A list of documents to update
        """

        docs = list(chain.from_iterable(items))  # type: ignore

        # Add timestamp
        for item in docs:
            item.update(
                {
                    "_bt": self.timestamp,
                }
            )

        molecule_ids = list({item["molecule_id"] for item in docs})

        if len(items) > 0:
            self.logger.info(f"Updating {len(docs)} summary documents")
            self.summary.remove_docs({self.summary.key: {"$in": molecule_ids}})
            self.summary.update(
                docs=docs,
                key=["molecule_id"],
            )
        else:
            self.logger.info("No items to update")
