# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import pydantic
from langchain_core import documents

from burr import core
from burr.core import State, action, expr, persistence, state
from burr.tracking import client as tracking_client


@action(reads=[], writes=["dict"])
def basic_action(state: State, user_input: str) -> tuple[dict, State]:
    v = {"foo": 1, "bar": "2", "bool": True, "None": None, "input": user_input}
    return {"dict": v}, state.update(dict=v)


class PydanticField(pydantic.BaseModel):
    f1: int = 0
    f2: bool = False


@action(reads=["dict"], writes=["pydantic_field"])
def pydantic_action(state: State) -> tuple[dict, State]:
    v = PydanticField(f1=state["dict"]["foo"], f2=state["dict"]["bool"])
    return {"pydantic_field": v}, state.update(pydantic_field=v)


@action(reads=["pydantic_field"], writes=["lc_doc"])
def langchain_action(state: State) -> tuple[dict, State]:
    v = documents.Document(
        page_content=f"foo: {state['pydantic_field'].f1}, bar: {state['pydantic_field'].f2}"
    )
    return {"lc_doc": v}, state.update(lc_doc=v)


@action(reads=["lc_doc"], writes=[])
def terminal_action(state: State) -> tuple[dict, State]:
    return {"output": state["lc_doc"].page_content}, state


def build_application(sqllite_persister, tracker, partition_key, app_id):
    persister = sqllite_persister or tracker
    app_builder = (
        core.ApplicationBuilder()
        .with_actions(basic_action, pydantic_action, langchain_action, terminal_action)
        .with_transitions(
            ("basic_action", "terminal_action", expr("dict['foo'] == 0")),
            ("basic_action", "pydantic_action"),
            ("pydantic_action", "langchain_action"),
            ("langchain_action", "terminal_action"),
        )
        .with_identifiers(partition_key=partition_key, app_id=app_id)
        .initialize_from(
            persister,
            resume_at_next_action=True,
            default_state={
                "custom_field": documents.Document(
                    page_content="this is a custom field to serialize"
                )
            },
            default_entrypoint="basic_action",
        )
    )
    if sqllite_persister:
        app_builder.with_state_persister(sqllite_persister)
    if tracker:
        app_builder.with_tracker(tracker)
    return app_builder.build()


def register_custom_serde_for_lc_document(field_name: str):
    """Register a custom serde for a field"""

    def my_field_serializer(value: documents.Document, **kwargs) -> dict:
        serde_value = f"serialized::{value.page_content}"
        return {"value": serde_value}

    def my_field_deserializer(value: dict, **kwargs) -> documents.Document:
        serde_value = value["value"]
        return documents.Document(page_content=serde_value.replace("serialized::", ""))

    state.register_field_serde(field_name, my_field_serializer, my_field_deserializer)


def test_whole_application_tracker(tmp_path):
    """This test creates an application and then steps through it rebuilding the
    application each time. This is a test of things being serialized and deserialized."""
    tracker = tracking_client.LocalTrackingClient("integration-test", tmp_path)
    app_id = "integration-test"
    partition_key = ""
    # step 1
    app = build_application(None, tracker, partition_key, app_id)
    register_custom_serde_for_lc_document("custom_field")
    action1, result1, state1 = app.step(inputs={"user_input": "hello"})
    # check custom serde
    assert state1.serialize() == {
        "__PRIOR_STEP": "basic_action",
        "__SEQUENCE_ID": 0,
        "custom_field": {"value": "serialized::this is a custom field to serialize"},
        "dict": {"None": None, "bar": "2", "bool": True, "foo": 1, "input": "hello"},
    }
    assert action1.name == "basic_action"
    # step 2
    app = build_application(None, tracker, partition_key, app_id)
    action2, result2, state2 = app.step()
    assert action2.name == "pydantic_action"
    # step 3
    app = build_application(None, tracker, partition_key, app_id)
    action3, result3, state3 = app.step()
    assert action3.name == "langchain_action"
    # step 4
    app = build_application(None, tracker, partition_key, app_id)
    action4, result4, state4 = app.step()
    assert action4.name == "terminal_action"

    # assert that state is basically the same across different steps
    assert state1["dict"] == {"foo": 1, "bar": "2", "bool": True, "None": None, "input": "hello"}
    assert state1["dict"] == state4["dict"]

    assert state2["pydantic_field"].f1 == 1
    assert state2["pydantic_field"].f2 is True
    assert state2["pydantic_field"] == state3["pydantic_field"]

    assert state3["lc_doc"].page_content == "foo: 1, bar: True"
    assert state3["lc_doc"] == state4["lc_doc"]

    # assert that tracker has things in it too
    final_tracker_state = tracker.load(partition_key, app_id=app_id)
    for k, v in final_tracker_state["state"].items():
        assert v == state4[k]


def test_whole_application_sqllite(tmp_path):
    """This test creates an application and then steps through it rebuilding the
    application each time. This is a test of things being serialized and deserialized."""
    sqllite_persister = persistence.SQLLitePersister(tmp_path / "test.db")
    sqllite_persister.initialize()
    app_id = "integration-test"
    partition_key = ""
    # step 1
    app = build_application(sqllite_persister, None, partition_key, app_id)
    register_custom_serde_for_lc_document("custom_field")
    action1, result1, state1 = app.step(inputs={"user_input": "hello"})
    # check custom serde
    assert state1.serialize() == {
        "__PRIOR_STEP": "basic_action",
        "__SEQUENCE_ID": 0,
        "custom_field": {"value": "serialized::this is a custom field to serialize"},
        "dict": {"None": None, "bar": "2", "bool": True, "foo": 1, "input": "hello"},
    }
    # check actions
    assert action1.name == "basic_action"
    # step 2
    app = build_application(sqllite_persister, None, partition_key, app_id)
    action2, result2, state2 = app.step()
    assert action2.name == "pydantic_action"
    # step 3
    app = build_application(sqllite_persister, None, partition_key, app_id)
    action3, result3, state3 = app.step()
    assert action3.name == "langchain_action"
    # step 4
    app = build_application(sqllite_persister, None, partition_key, app_id)
    action4, result4, state4 = app.step()
    assert action4.name == "terminal_action"

    # assert that state is basically the same across different steps
    assert state1["dict"] == {"foo": 1, "bar": "2", "bool": True, "None": None, "input": "hello"}
    assert state1["dict"] == state4["dict"]

    assert state2["pydantic_field"].f1 == 1
    assert state2["pydantic_field"].f2 is True
    assert state2["pydantic_field"] == state3["pydantic_field"]

    assert state3["lc_doc"].page_content == "foo: 1, bar: True"
    assert state3["lc_doc"] == state4["lc_doc"]

    final_sqllite_state = sqllite_persister.load("", app_id=app_id)
    assert final_sqllite_state["state"] == state4
    assert sqllite_persister.list_app_ids(partition_key="") == ["integration-test"]
