# License Apache 2.0: (c) 2025 Yoan Sallami (Synalinks Team)

import copy
import warnings

from synalinks.src import ops
from synalinks.src.api_export import synalinks_export
from synalinks.src.backend import Embedding as EmbeddingVector
from synalinks.src.backend import JsonDataModel
from synalinks.src.backend import is_entities
from synalinks.src.backend import is_entity
from synalinks.src.backend import is_knowledge_graph
from synalinks.src.backend import is_relation
from synalinks.src.backend import is_relations
from synalinks.src.modules.module import Module
from synalinks.src.saving import serialization_lib


@synalinks_export(
    [
        "synalinks.modules.Embedding",
        "synalinks.Embedding",
    ]
)
class Embedding(Module):
    """Extracts and updates the embedding vector of entities.

    This module is designed to work with `Entity`, `Relation`, `Entities`,
    `Relations` or `KnowledgeGraph` data models. It supports to mask the
    entity fields in order to keep **only one** field to embed per entity.

    **Note**: Each entity should have the *same field* to compute the embedding
        from like a `name` or `description` field using `in_mask`.
        **Or** every entity should have *only one field left* after masking using
        `out_mask` argument.

    ```python
    import synalinks
    import asyncio
    from typing import Literal

    class Document(synalinks.Entity):
        label: Literal["Document"]
        text: str = synalinks.Field(
            description="The document content",
        )

    async def main():
        inputs = synalinks.Input(data_model=Document)
        outputs = await synalinks.Embedding(
            embedding_model=embedding_model,
            in_mask=["text"],
        )(inputs)

        program = synalinks.Program(
            inputs=inputs,
            outputs=outputs,
            name="embbed_document",
            description="Embbed the given documents"
        )

        doc = Document(
            label="Document",
            text="my document",
        )

        result = await program(doc)

    if __name__ == "__main__":
        asyncio.run(main())
    ```

    If you want to process batch asynchronously
    use `program.predict()` instead, see the [FAQ](https://synalinks.github.io/synalinks/FAQ/#whats-the-difference-between-program-methods-predict-and-__call__)
    to understand the difference between `program()` and `program.predict()`

    Here is an example:

    ```python
    import synalinks
    import asyncio
    import numpy as np
    from typing import Literal

    class Document(synalinks.Entity):
        label: Literal["Document"]
        text: str = synalinks.Field(
            description="The document content",
        )

    async def main():
        inputs = synalinks.Input(data_model=Document)
        outputs = await synalinks.Embedding(
            embedding_model=embedding_model,
            in_mask=["text"],
        )(inputs)

        program = synalinks.Program(
            inputs=inputs,
            outputs=outputs,
            name="embbed_document",
            description="Embbed the given documents"
        )

        doc1 = Document(label="Document", text="my document 1")
        doc2 = Document(label="Document", text="my document 2")
        doc3 = Document(label="Document", text="my document 3")

        docs = np.array([doc1, doc2, doc3], dtype="object")

        embedded_docs = await program.predict(docs)

    if __name__ == "__main__":
        asyncio.run(main())
    ```

    Args:
        embedding_model (EmbeddingModel): The embedding model to use.
        in_mask (list): A mask applied to keep specific entity fields.
        out_mask (list): A mask applied to remove specific entity fields.
        name (str): Optional. The name of the module.
        description (str): Optional. The description of the module.
        trainable (bool): Whether the module's variables should be trainable.
    """

    def __init__(
        self,
        embedding_model=None,
        in_mask=None,
        out_mask=None,
        name=None,
        description=None,
        trainable=False,
    ):
        super().__init__(
            name=name,
            description=description,
            trainable=trainable,
        )
        self.embedding_model = embedding_model
        self.in_mask = in_mask
        self.out_mask = out_mask

    async def _embed_entity(self, entity):
        # # Check if entity is already embedded and has valid embeddings
        embeddings = entity.get("embeddings")
        if embeddings:
            warnings.warn(
                "Embeddings already generated for entity.Returning original entity."
            )
            return JsonDataModel(
                json=entity.get_json(),
                schema=entity.get_schema(),
                name=entity.name + "_embedded",
            )
        # Apply masking to the entity
        filtered_entity = entity  # Default to original entity

        if self.out_mask:
            filtered_entity = await ops.out_mask(
                entity,
                mask=self.out_mask,
                recursive=False,
                name=entity.name + "_out_mask",
            )
        elif self.in_mask:
            filtered_entity = await ops.in_mask(
                entity,
                mask=self.in_mask,
                recursive=False,
                name=entity.name + "_in_mask",
            )

        # Generate embeddings
        embeddings = await ops.embedding(
            filtered_entity,
            embedding_model=self.embedding_model,
            name=entity.name + "_embedding",
        )

        # Validate embeddings
        if not embeddings or not embeddings.get("embeddings"):
            warnings.warn(
                f"No embeddings generated for entity {entity.name}. "
                "Please check that your schema is correct."
            )
            return None

        embedding_list = embeddings.get("embeddings")
        if len(embedding_list) != 1:
            warnings.warn(
                "Entities can only have one embedding vector per entity, "
                "adjust `Embedding` module's `in_mask` or `out_mask` "
                "to keep only one field. Skipping embedding."
            )
            return None

        # Add embedding to entity
        vector = embedding_list[0]
        return await ops.concat(
            entity,
            EmbeddingVector(embedding=vector),
            name=entity.name + "_embedded",
        )

    async def _embed_relation(self, relation):
        subj = relation.get_nested_entity("subj")
        obj = relation.get_nested_entity("obj")
        if not subj or not obj:
            return None
        embedded_subj = await self._embed_entity(subj)
        embedded_obj = await self._embed_entity(obj)

        relation_json = copy.deepcopy(relation.get_json())
        relation_json.update(
            {
                "subj": embedded_subj.get_json(),
                "obj": embedded_obj.get_json(),
            }
        )
        outputs_schema = copy.deepcopy(relation.get_schema())

        # Update schema definitions for embedded entities
        if outputs_schema.get("$defs"):
            subj_label = subj.get("label")
            obj_label = obj.get("label")

            if subj_label and subj_label in outputs_schema["$defs"]:
                embedded_subj_schema = embedded_subj.get_schema()
                if embedded_subj_schema.get("properties"):
                    outputs_schema["$defs"][subj_label]["properties"].update(
                        embedded_subj_schema["properties"]
                    )

            if obj_label and obj_label in outputs_schema["$defs"]:
                embedded_obj_schema = embedded_obj.get_schema()
                if embedded_obj_schema.get("properties"):
                    outputs_schema["$defs"][obj_label]["properties"].update(
                        embedded_obj_schema["properties"]
                    )

        return JsonDataModel(
            json=relation_json,
            schema=outputs_schema,
            name=relation.name + "_embedded",
        )

    async def call(self, inputs):
        if not inputs:
            return None
        if is_knowledge_graph(inputs):
            entities_json = []
            relations_json = []
            outputs_schema = copy.deepcopy(inputs.get_schema())

            # Process entities
            for entity in inputs.get_nested_entity_list("entities"):
                embedded_entity = await self._embed_entity(entity)
                if embedded_entity:
                    entities_json.append(embedded_entity.get_json())

                    # Update schema definitions
                    if outputs_schema.get("$defs"):
                        entity_label = entity.get("label")
                        if entity_label and entity_label in outputs_schema["$defs"]:
                            embedded_schema = embedded_entity.get_schema()
                            if embedded_schema.get("properties"):
                                outputs_schema["$defs"][entity_label][
                                    "properties"
                                ].update(embedded_schema["properties"])
            # Process relations
            for relation in inputs.get_nested_entity_list("relations"):
                embedded_relation = await self._embed_relation(relation)
                if embedded_relation:
                    relations_json.append(embedded_relation.get_json())

                    embedded_schema = embedded_relation.get_schema()
                    if embedded_schema.get("$defs") and outputs_schema.get("$defs"):
                        for def_key, def_value in embedded_schema["$defs"].items():
                            if def_key in outputs_schema["$defs"]:
                                # Merge properties if they exist
                                if def_value.get("properties") and outputs_schema[
                                    "$defs"
                                ][def_key].get("properties"):
                                    outputs_schema["$defs"][def_key]["properties"].update(
                                        def_value["properties"]
                                    )
                                else:
                                    outputs_schema["$defs"][def_key] = def_value

            # Update output JSON
            outputs_json = inputs.get_json()
            outputs_json.update({"entities": entities_json, "relations": relations_json})
            return JsonDataModel(
                json=outputs_json,
                schema=outputs_schema,
                name=inputs.name + "_embedded",
            )

        elif is_entities(inputs):
            entities_json = []
            outputs_schema = copy.deepcopy(inputs.get_schema())

            # Process all entities and collect schema updates
            for entity in inputs.get_nested_entity_list("entities"):
                embedded_entity = await self._embed_entity(entity)
                if embedded_entity:
                    entities_json.append(embedded_entity.get_json())

                    # Update schema definitions
                    if outputs_schema.get("$defs"):
                        entity_label = entity.get("label")
                        if entity_label and entity_label in outputs_schema["$defs"]:
                            embedded_schema = embedded_entity.get_schema()
                            if embedded_schema.get("properties"):
                                outputs_schema["$defs"][entity_label][
                                    "properties"
                                ].update(embedded_schema["properties"])

            # Update output JSON with embedded entities
            outputs_json = inputs.get_json()
            outputs_json.update({"entities": entities_json})

            return JsonDataModel(
                json=outputs_json,
                schema=outputs_schema,
                name=inputs.name + "_embedded",
            )

        elif is_relations(inputs):
            relations_json = []
            outputs_schema = copy.deepcopy(inputs.get_schema())

            # Process all relations
            for relation in inputs.get_nested_entity_list("relations"):
                embedded_relation = await self._embed_relation(relation)
                if embedded_relation:
                    relations_json.append(embedded_relation.get_json())

                    # Merge schema definitions from embedded relation
                    embedded_schema = embedded_relation.get_schema()
                    if embedded_schema.get("$defs") and outputs_schema.get("$defs"):
                        for def_key, def_value in embedded_schema["$defs"].items():
                            if def_key in outputs_schema["$defs"]:
                                # Merge properties if they exist
                                if def_value.get("properties") and outputs_schema[
                                    "$defs"
                                ][def_key].get("properties"):
                                    outputs_schema["$defs"][def_key]["properties"].update(
                                        def_value["properties"]
                                    )
                                else:
                                    outputs_schema["$defs"][def_key] = def_value

            # Update output JSON
            outputs_json = inputs.get_json()
            outputs_json.update({"relations": relations_json})

            return JsonDataModel(
                json=outputs_json,
                schema=outputs_schema,
                name=inputs.name + "_embedded",
            )
        elif is_relation(inputs):
            return await self._embed_relation(inputs)
        elif is_entity(inputs):
            return await self._embed_entity(inputs)
        else:
            return None

    async def compute_output_spec(self, inputs):
        return inputs.clone(name=inputs.name + "_embedded")

    def get_config(self):
        config = {
            "in_mask": self.in_mask,
            "out_mask": self.out_mask,
            "name": self.name,
            "description": self.description,
            "trainable": self.trainable,
        }
        embedding_model_config = {
            "embedding_model": serialization_lib.serialize_synalinks_object(
                self.embedding_model
            )
        }
        return {**embedding_model_config, **config}

    @classmethod
    def from_config(cls, config):
        embedding_model = serialization_lib.deserialize_synalinks_object(
            config.pop("embedding_model")
        )
        return cls(embedding_model=embedding_model, **config)
