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

import re
from typing import Any
from typing import Dict
from typing import List
from typing import Optional
from typing import Tuple

import jinja2

from synalinks.src import ops
from synalinks.src.api_export import synalinks_export
from synalinks.src.backend import ChatMessage
from synalinks.src.backend import ChatMessages
from synalinks.src.backend import ChatRole
from synalinks.src.backend import DataModel
from synalinks.src.backend import SymbolicDataModel
from synalinks.src.modules.module import Module
from synalinks.src.saving import serialization_lib

XML_TAGS_REGEX = re.compile(
    r"<("
    + "|".join([ChatRole.SYSTEM, ChatRole.USER, ChatRole.ASSISTANT])
    + r")\s*(?:[^>]*)>\s*([\s\S]*?)\s*</\1>",
    re.MULTILINE,
)


@synalinks_export("synalinks.default_prompt_template")
def default_prompt_template():
    """
    Returns the default prompt template

    Returns:
        (str): The default prompt template.
    """
    return """
<system>
{% if inputs_schema %}
You will be given an input JSON object with the following schema. 
Input JSON Schema:
{{ inputs_schema }}
{% else %}
You will be given an input JSON object.
{% endif %}
{% if outputs_schema %}
Your task is to answer with a JSON object and follow this output JSON schema.
Output JSON Schema:
{{ outputs_schema }}
{% else %}
Your task is to answer with a JSON object.
{% endif %}
{% if examples %}
Examples:
{% for example in examples %}
Input:
{{ example[0] }}
Output:
{{ example[1] }}
{% endfor %}
{% endif %}
{% if hints %}
Hints:
{% for hint in hints %}
 - {{ hint }}
{% endfor %}
{% endif %}
</system>
<user>
Input:
{{ inputs }}
Output:
</user>"""


@synalinks_export("synalinks.chat_prompt_template")
def chat_prompt_template():
    """
    Returns the default chat prompt template

    Returns:
        (str): The default chat prompt template.
    """
    return """
{% if hints %}
<system>
Hints:
{% for hint in hints %}
 - {{ hint }}
{% endfor %}
</system>
{% endif %}
{% for message in inputs.messages %}
{% if message.role == "assistant" %}
<assistant>
{{ message.content }}
</assistant>
{% elif message.role == "user" %}
<user>
{{ message.content }}
</user>
{% else %}{% endif %}
{% endfor %}
"""


class GeneratorVariables(DataModel):
    """The Generator variables."""

    prompt_template: str = None
    examples: List[
        Tuple[
            Dict[str, Any],  # input JSON value
            Dict[str, Any],  # output JSON value
            Optional[float],  # reward
        ]
    ] = []
    hints: List[str] = []
    predictions: List[
        Tuple[
            Dict[str, Any],  # input JSON value
            Dict[str, Any],  # output JSON value
            Optional[float],  # reward (None if not yet backpropagated)
        ],
    ] = []


@synalinks_export(["synalinks.modules.Generator", "synalinks.Generator"])
class Generator(Module):
    """
    Use a `LanguageModel` to generate a data model from an arbitrary input data model.

    Example:

    ```python
    import synalinks
    import asyncio

    async def main():

        class Query(DataModel):
            query: str = synalinks.Field(
                description="The user query",
            )

        class AnswerWithCritique(synalinks.DataModel):
            thinking: str = synalinks.Field(
                description="Your step by step thinking",
            )
            critique: str = synalinks.Field(
                description="The critique of the above thinking",
            )
            answer: str = synalinks.Field(
                description="The correct answer",
            )

        language_model = synalinks.LanguageModel("ollama_chat/deepseek-r1")

        x0 = synalinks.Input(data_model=Query)
        x1 = await synalinks.Generator(
            data_model=AnswerWithCritique,
            language_model=language_model,
        )(x0)

        program = synalinks.Program(
            inputs=x0,
            outputs=x1,
            name="chain_of_thought_with_critique",
            description="Useful to answer step by step and evaluate your answer",
        )

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

    Args:
        schema (dict): The target JSON schema.
            If not provided use the `data_model` to infer it.
        data_model (DataModel | SymbolicDataModel | JsonDataModel): The target data_model.
        language_model (LanguageModel): The language model to use.
        prompt_template (str): The jinja2 prompt template.
        examples (list): The default list of examples, the examples
            are a list of tuples containing input/output JSON pairs.
        hints (list): The default hints being a list of string containing
            addtional hints for the language model.
        return_inputs (bool): Optional. Whether or not to concatenate the inputs to
            the outputs (Default to False).
        use_inputs_schema (bool): Optional. Whether or not use the inputs schema in
            the prompt (Default to False).
        use_outputs_schema (bool): Optional. Whether or not use the outputs schema in
            the prompt (Default to False).
        streaming (str): Optional. If true stream the LM response, enabled only if
            `schema` is `None` and only during inference (not during training).
        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,
        schema=None,
        data_model=None,
        language_model=None,
        prompt_template=None,
        examples=None,
        hints=None,
        return_inputs=False,
        use_inputs_schema=False,
        use_outputs_schema=False,
        streaming=False,
        name=None,
        description=None,
        trainable=True,
    ):
        super().__init__(
            name=name,
            description=description,
            trainable=trainable,
        )
        if not schema and data_model:
            schema = data_model.schema()
        self.schema = schema
        self.language_model = language_model
        if not prompt_template:
            prompt_template = default_prompt_template()
        self.prompt_template = prompt_template
        if not examples:
            examples = []
        self.examples = examples
        if not hints:
            hints = []
        self.hints = hints
        self.return_inputs = return_inputs
        self.use_inputs_schema = use_inputs_schema
        self.use_outputs_schema = use_outputs_schema
        if schema and streaming:
            streaming = False
        self.streaming = streaming
        self.state = self.add_variable(
            initializer=GeneratorVariables(
                prompt_template=prompt_template,
                examples=examples,
                hints=hints,
            ).json(),
            data_model=GeneratorVariables,
            trainable=True,
            name=self.name + "_state",
        )

    async def call(self, inputs, training=False):
        if not inputs:
            return None
        msgs = ChatMessages()
        msgs.messages = self.format_messages(inputs)
        if self.streaming and not training:
            streaming = True
        else:
            streaming = False
        result = await ops.predict(
            msgs,
            schema=self.schema,
            language_model=self.language_model,
            streaming=streaming,
            name=self.name + "_prediction",
        )
        if streaming:
            return result
        if result:
            if training:
                self.state.get("predictions").append(
                    (
                        inputs.json(),
                        result.json(),
                        None,
                    )
                )
            if self.return_inputs:
                return await ops.concat(
                    inputs,
                    result,
                    name=self.name + "_with_inputs",
                )
            else:
                return result
        return None

    async def compute_output_spec(self, inputs, training=False):
        if self.schema:
            if self.return_inputs:
                return await ops.concat(
                    inputs,
                    SymbolicDataModel(
                        schema=self.schema,
                        name=self.name,
                    ),
                    name=self.name + "_with_inputs",
                )
            else:
                return SymbolicDataModel(
                    schema=self.schema,
                    name=self.name,
                )
        else:
            if self.return_inputs:
                return await ops.concat(
                    inputs,
                    SymbolicDataModel(
                        schema=ChatMessage.schema(),
                        name=self.name,
                    ),
                    name=self.name + "_with_inputs",
                )
            else:
                return SymbolicDataModel(
                    schema=ChatMessage.schema(),
                    name=self.name,
                )

    def format_messages(self, inputs):
        template = jinja2.Template(self.state.get("prompt_template"))
        rendered_prompt = template.render(
            inputs_schema=inputs.schema() if self.use_inputs_schema else None,
            outputs_schema=self.schema if self.use_outputs_schema else None,
            examples=self.state.get("examples"),
            hints=self.state.get("hints"),
            inputs=inputs.json(),
        )
        matches = XML_TAGS_REGEX.findall(rendered_prompt)
        extracted_tags = [(match[0], match[1].strip()) for match in matches]
        messages = []
        for message in extracted_tags:
            role, content = message
            if content:
                messages.append(ChatMessage(role=role, content=content))
        return messages

    def get_config(self):
        config = {
            "schema": self.schema,
            "prompt_template": self.prompt_template,
            "examples": self.examples,
            "hints": self.hints,
            "use_inputs_schema": self.use_inputs_schema,
            "use_outputs_schema": self.use_outputs_schema,
            "return_inputs": self.return_inputs,
            "name": self.name,
            "description": self.description,
            "trainable": self.trainable,
        }
        language_model_config = {
            "language_model": serialization_lib.serialize_synalinks_object(
                self.language_model
            )
        }
        return {**config, **language_model_config}

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