# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0

import logging
from typing import Any
from unittest.mock import patch

import arrow
import pytest
from jinja2 import TemplateSyntaxError

from haystack import component
from haystack.components.builders.prompt_builder import PromptBuilder
from haystack.core.pipeline.pipeline import Pipeline
from haystack.dataclasses.document import Document


class TestPromptBuilder:
    def test_init(self):
        builder = PromptBuilder(template="This is a {{ variable }}")
        assert builder.template is not None
        assert builder.required_variables == "*"
        assert builder._template_string == "This is a {{ variable }}"
        assert builder._variables is None
        assert builder._required_variables == "*"

        # we have inputs that contain: template, template_variables + inferred variables
        inputs = builder.__haystack_input__._sockets_dict
        assert set(inputs.keys()) == {"template", "template_variables", "variable"}
        assert inputs["template"].type == str | None
        assert inputs["template_variables"].type == dict[str, Any] | None
        assert inputs["variable"].type == Any

        # response is always prompt
        outputs = builder.__haystack_output__._sockets_dict
        assert set(outputs.keys()) == {"prompt"}
        assert outputs["prompt"].type == str

    def test_init_with_required_variables(self):
        builder = PromptBuilder(template="This is a {{ variable }}", required_variables=["variable"])
        assert builder.template is not None
        assert builder.required_variables == ["variable"]
        assert builder._template_string == "This is a {{ variable }}"
        assert builder._variables is None
        assert builder._required_variables == ["variable"]

        # we have inputs that contain: template, template_variables + inferred variables
        inputs = builder.__haystack_input__._sockets_dict
        assert set(inputs.keys()) == {"template", "template_variables", "variable"}
        assert inputs["template"].type == str | None
        assert inputs["template_variables"].type == dict[str, Any] | None
        assert inputs["variable"].type == Any

        # response is always prompt
        outputs = builder.__haystack_output__._sockets_dict
        assert set(outputs.keys()) == {"prompt"}
        assert outputs["prompt"].type == str

    def test_init_with_custom_variables(self):
        variables = ["var1", "var2", "var3"]
        template = "Hello, {{ var1 }}, {{ var2 }}!"
        builder = PromptBuilder(template=template, variables=variables)
        assert builder.template is not None
        assert builder.required_variables == "*"
        assert builder._variables == variables
        assert builder._template_string == "Hello, {{ var1 }}, {{ var2 }}!"
        assert builder._required_variables == "*"

        # we have inputs that contain: template, template_variables + variables
        inputs = builder.__haystack_input__._sockets_dict
        assert set(inputs.keys()) == {"template", "template_variables", "var1", "var2", "var3"}
        assert inputs["template"].type == str | None
        assert inputs["template_variables"].type == dict[str, Any] | None
        assert inputs["var1"].type == Any
        assert inputs["var2"].type == Any
        assert inputs["var3"].type == Any

        # response is always prompt
        outputs = builder.__haystack_output__._sockets_dict
        assert set(outputs.keys()) == {"prompt"}
        assert outputs["prompt"].type == str

    def test_to_dict(self):
        builder = PromptBuilder(
            template="This is a {{ variable }}", variables=["var1", "var2"], required_variables=["var1", "var3"]
        )
        res = builder.to_dict()
        assert res == {
            "type": "haystack.components.builders.prompt_builder.PromptBuilder",
            "init_parameters": {
                "template": "This is a {{ variable }}",
                "variables": ["var1", "var2"],
                "required_variables": ["var1", "var3"],
            },
        }

    def test_to_dict_without_optional_params(self):
        builder = PromptBuilder(template="This is a {{ variable }}")
        res = builder.to_dict()
        assert res == {
            "type": "haystack.components.builders.prompt_builder.PromptBuilder",
            "init_parameters": {"template": "This is a {{ variable }}", "variables": None, "required_variables": "*"},
        }

    def test_run(self):
        builder = PromptBuilder(template="This is a {{ variable }}")
        res = builder.run(variable="test")
        assert res == {"prompt": "This is a test"}

    def test_run_template_variable(self):
        builder = PromptBuilder(template="This is a {{ variable }}")
        res = builder.run(template_variables={"variable": "test"})
        assert res == {"prompt": "This is a test"}

    def test_run_template_variable_overrides_variable(self):
        builder = PromptBuilder(template="This is a {{ variable }}")
        res = builder.run(template_variables={"variable": "test_from_template_var"}, variable="test")
        assert res == {"prompt": "This is a test_from_template_var"}

    def test_run_without_input(self):
        builder = PromptBuilder(template="This is a template without input")
        res = builder.run()
        assert res == {"prompt": "This is a template without input"}

    def test_run_with_missing_optional_input(self):
        builder = PromptBuilder(template="This is a {{ variable }}", required_variables=None)
        res = builder.run()
        assert res == {"prompt": "This is a "}

    def test_run_with_missing_required_input_default(self):
        builder = PromptBuilder(template="This is a {{ variable }}")
        with pytest.raises(ValueError, match="variable"):
            builder.run()

    def test_run_with_missing_required_input(self):
        builder = PromptBuilder(template="This is a {{ foo }}, not a {{ bar }}", required_variables=["foo", "bar"])
        with pytest.raises(ValueError, match="foo"):
            builder.run(bar="bar")
        with pytest.raises(ValueError, match="bar"):
            builder.run(foo="foo")
        with pytest.raises(ValueError, match="foo, bar"):
            builder.run()

    def test_run_with_missing_required_input_using_star(self):
        builder = PromptBuilder(template="This is a {{ foo }}, not a {{ bar }}", required_variables="*")
        with pytest.raises(ValueError, match="foo"):
            builder.run(bar="bar")
        with pytest.raises(ValueError, match="bar"):
            builder.run(foo="foo")
        with pytest.raises(ValueError, match="bar, foo"):
            builder.run()

    def test_run_with_variables(self):
        variables = ["var1", "var2", "var3"]
        template = "Hello, {{ name }}! {{ var1 }}"
        builder = PromptBuilder(template=template, variables=variables, required_variables=["name", "var1"])
        assert builder.run(name="John", var1="How are you?") == {"prompt": "Hello, John! How are you?"}

    def test_run_overwriting_default_template(self):
        builder = PromptBuilder(template="Hello, {{ name }}!")
        assert builder.run("Hello, {{ var1 }}{{ name }}!!", name="John") == {"prompt": "Hello, John!!"}

    def test_run_overwriting_default_template_with_template_variables(self):
        builder = PromptBuilder(template="Hello, {{ name }}!")
        template = "Hello, {{ var1 }} {{ name }}!"
        assert builder.run(template, var1="Big", name="John") == {"prompt": "Hello, Big John!"}

    def test_run_overwriting_default_template_with_variables(self):
        variables = ["var1", "var2", "name"]
        builder = PromptBuilder(template="Hello, {{ name }}!", variables=variables, required_variables=["name"])
        template = "Hello, {{ var1 }} {{ name }}!"
        assert builder.run(template, name="John", var1="Big") == {"prompt": "Hello, Big John!"}

    def test_run_with_invalid_template(self):
        builder = PromptBuilder(template="Hello, {{ name }}!")

        template = "Hello, {{ name }!"
        template_variables = {"name": "John"}
        with pytest.raises(TemplateSyntaxError):
            builder.run(template, template_variables)

    def test_init_with_invalid_template(self):
        template = "Hello, {{ name }!"
        with pytest.raises(TemplateSyntaxError):
            PromptBuilder(template)

    def test_provided_template_variables(self):
        prompt_builder = PromptBuilder(template="", variables=["documents"], required_variables=["city"])

        # both variables are provided
        prompt_builder._validate_variables({"name", "city"})

        # provided variables are a superset of the required variables
        prompt_builder._validate_variables({"name", "city", "age"})

        with pytest.raises(ValueError):
            prompt_builder._validate_variables({"name"})

    def test_example_in_pipeline(self):
        default_template = "Here is the document: {{documents[0].content}} \\n Answer: {{query}}"
        prompt_builder = PromptBuilder(template=default_template, variables=["documents"])

        @component
        class DocumentProducer:
            @component.output_types(documents=list[Document])
            def run(self, doc_input: str):
                return {"documents": [Document(content=doc_input)]}

        pipe = Pipeline()
        pipe.add_component("doc_producer", DocumentProducer())
        pipe.add_component("prompt_builder", prompt_builder)
        pipe.connect("doc_producer.documents", "prompt_builder.documents")

        template = "Here is the document: {{documents[0].content}} \n Query: {{query}}"
        result = pipe.run(
            data={
                "doc_producer": {"doc_input": "Hello world, I live in Berlin"},
                "prompt_builder": {
                    "template": template,
                    "template_variables": {"query": "Where does the speaker live?"},
                },
            }
        )

        assert result == {
            "prompt_builder": {
                "prompt": "Here is the document: Hello world, I live in Berlin \n Query: Where does the speaker live?"
            }
        }

    def test_example_in_pipeline_simple(self):
        default_template = "This is the default prompt:\n Query: {{query}}"
        prompt_builder = PromptBuilder(template=default_template)

        pipe = Pipeline()
        pipe.add_component("prompt_builder", prompt_builder)

        # using the default prompt
        result = pipe.run(data={"query": "Where does the speaker live?"})
        expected_default = {
            "prompt_builder": {"prompt": "This is the default prompt:\n Query: Where does the speaker live?"}
        }
        assert result == expected_default

        # using the dynamic prompt
        result = pipe.run(
            data={"query": "Where does the speaker live?", "template": "This is the dynamic prompt:\n Query: {{query}}"}
        )
        expected_dynamic = {
            "prompt_builder": {"prompt": "This is the dynamic prompt:\n Query: Where does the speaker live?"}
        }
        assert result == expected_dynamic

    def test_warning_when_required_variables_explicitly_none(self, caplog):
        with caplog.at_level(logging.WARNING):
            _ = PromptBuilder(template="This is a {{ variable }}", required_variables=None)
            assert "explicitly set to `None`" in caplog.text

    def test_no_warning_with_default_required_variables(self, caplog):
        with caplog.at_level(logging.WARNING):
            _ = PromptBuilder(template="This is a {{ variable }}")
            assert "required_variables" not in caplog.text

    def test_variables_correct_with_assignment(self) -> None:
        template = """{% if existing_documents is not none %}
{% set existing_doc_len = existing_documents|length %}
{% else %}
{% set existing_doc_len = 0 %}
{% endif %}
{% for doc in docs %}
<document reference="{{loop.index + existing_doc_len}}">
{{ doc.content }}
</document>
{% endfor %}
"""
        builder = PromptBuilder(template=template, required_variables="*")
        assert set(builder.variables) == {"docs", "existing_documents"}
        assert builder.required_variables == "*"

    def test_variables_correct_with_tuple_assignment(self):
        template = """{% if existing_documents is not none -%}
{% set x, y = (existing_documents|length, 1) -%}
{% else -%}
{% set x, y = (0, 1) -%}
{% endif -%}
x={{ x }}, y={{ y }}
"""
        builder = PromptBuilder(template=template, required_variables="*")
        assert builder.variables == ["existing_documents"]
        assert builder.required_variables == "*"
        res = builder.run(existing_documents=None)
        assert res["prompt"] == "x=0, y=1"

    def test_variables_correct_with_list_assignment(self):
        template = """{% if existing_documents is not none -%}
{% set x, y = [existing_documents|length, 1] -%}
{% else -%}
{% set x, y = [0, 1] -%}
{% endif -%}
x={{ x }}, y={{ y }}
"""
        builder = PromptBuilder(template=template, required_variables="*")
        assert builder.variables == ["existing_documents"]
        assert builder.required_variables == "*"
        res = builder.run(existing_documents=None)
        assert res["prompt"] == "x=0, y=1"


class TestPromptBuilderWithJinja2TimeExtension:
    @pytest.fixture(autouse=True)
    def mock_now(self, monkeypatch):
        """Mock the arrow.now function to return a fixed datetime"""
        monkeypatch.setattr("arrow.now", lambda timezone="UTC": arrow.get("1970-01-01 00:00:00").to(timezone))

    @patch("haystack.components.builders.prompt_builder.Jinja2TimeExtension")
    def test_init_with_missing_extension_dependency(self, extension_mock):
        extension_mock.side_effect = ImportError
        builder = PromptBuilder(template="This is a {{ variable }}")
        assert builder._env.extensions == {}
        res = builder.run(variable="test")
        assert res == {"prompt": "This is a test"}

    def test_with_custom_dateformat(self) -> None:
        template = "Formatted date: {% now 'UTC', '%Y-%m-%d' %}"
        builder = PromptBuilder(template=template)

        result = builder.run()["prompt"]

        now_formatted = f"Formatted date: {arrow.now('UTC').strftime('%Y-%m-%d')}"

        assert now_formatted == result

    def test_with_different_timezone(self) -> None:
        template = "Current time in New York is: {% now 'America/New_York' %}"
        builder = PromptBuilder(template=template)

        result = builder.run()["prompt"]

        now_ny = f"Current time in New York is: {arrow.now('America/New_York').strftime('%Y-%m-%d %H:%M:%S')}"

        assert now_ny == result

    def test_date_with_addition_offset(self) -> None:
        template = "Time after 2 hours is: {% now 'UTC' + 'hours=2' %}"
        builder = PromptBuilder(template=template)

        result = builder.run()["prompt"]

        now_plus_2 = f"Time after 2 hours is: {(arrow.now('UTC').shift(hours=+2)).strftime('%Y-%m-%d %H:%M:%S')}"

        assert now_plus_2 == result

    def test_date_with_subtraction_offset(self) -> None:
        template = "Time after 12 days is: {% now 'UTC' - 'days=12' %}"
        builder = PromptBuilder(template=template)

        result = builder.run()["prompt"]

        now_plus_2 = f"Time after 12 days is: {(arrow.now('UTC').shift(days=-12)).strftime('%Y-%m-%d %H:%M:%S')}"

        assert now_plus_2 == result

    def test_invalid_timezone(self) -> None:
        template = "Current time is: {% now 'Invalid/Timezone' %}"
        builder = PromptBuilder(template=template)

        # Expect ValueError for invalid timezone
        with pytest.raises(ValueError, match="Invalid timezone"):
            builder.run()

    def test_invalid_offset(self) -> None:
        template = "Time after invalid offset is: {% now 'UTC' + 'invalid_offset' %}"
        builder = PromptBuilder(template=template)

        # Expect ValueError for invalid offset
        with pytest.raises(ValueError, match="Invalid offset or operator"):
            builder.run()