# SPDX-License-Identifier: MIT
# Copyright (c) 2025 LlamaIndex Inc.

from typing import Any, cast

import pytest
from pydantic import PrivateAttr

from workflows.context import JsonSerializer
from workflows.events import Event, StopEvent


class _TestEvent(Event):
    param: str
    _private_param_1: str = PrivateAttr()
    _private_param_2: str = PrivateAttr(default_factory=str)


class _TestEvent2(Event):
    """
    Custom Test Event.

    Private Attrs:
        _private_param: doesn't get modified during construction
        _modified_private_param: gets processed before being set
    """

    _private_param: int = PrivateAttr()
    _modified_private_param: int = PrivateAttr()

    def __init__(self, _modified_private_param: int, **params: Any):
        super().__init__(**params)
        self._modified_private_param = _modified_private_param * 2


def test_event_init_basic() -> None:
    evt = Event(a=1, b=2, c="c")

    assert evt.a == 1
    assert evt.b == 2
    assert evt.c == "c"
    assert evt["a"] == evt.a
    assert evt["b"] == evt.b
    assert evt["c"] == evt.c
    assert evt.keys() == {"a": 1, "b": 2, "c": "c"}.keys()


def test_custom_event_with_fields_and_private_params() -> None:
    evt = _TestEvent(a=1, param="test_param", _private_param_1="test_private_param_1")  # type: ignore

    assert evt.a == 1
    assert evt["a"] == evt.a
    assert evt.param == "test_param"
    assert evt._data == {"a": 1}
    assert evt._private_param_1 == "test_private_param_1"
    assert evt._private_param_2 == ""


def test_custom_event_override_init() -> None:
    evt = _TestEvent2(a=1, b=2, _private_param=2, _modified_private_param=2)

    assert evt.a == 1
    assert evt.b == 2
    assert evt._data == {"a": 1, "b": 2}
    assert evt._private_param == 2
    assert evt._modified_private_param == 4


def test_event_missing_key() -> None:
    ev = _TestEvent(param="bar")
    with pytest.raises(AttributeError):
        ev.wrong_key


def test_event_not_a_field() -> None:
    ev = _TestEvent(param="foo", not_a_field="bar")  # type: ignore
    assert ev._data["not_a_field"] == "bar"
    ev.not_a_field = "baz"
    assert ev._data["not_a_field"] == "baz"
    ev["not_a_field"] = "barbaz"
    assert ev._data["not_a_field"] == "barbaz"
    assert ev.get("not_a_field") == "barbaz"


def test_event_dict_api() -> None:
    ev = _TestEvent(param="foo")
    assert len(ev) == 0
    ev["a_new_key"] = "bar"
    assert len(ev) == 1
    assert list(ev.values()) == ["bar"]
    k, v = next(iter(ev.items()))
    assert k == "a_new_key"
    assert v == "bar"
    assert next(iter(ev)) == "a_new_key"
    assert ev.to_dict() == {"a_new_key": "bar"}


def test_event_serialization() -> None:
    ev = _TestEvent(param="foo", not_a_field="bar")  # type: ignore
    serializer = JsonSerializer()
    serialized_ev = serializer.serialize(ev)
    deseriazlied_ev = serializer.deserialize(serialized_ev)

    assert type(deseriazlied_ev).__name__ == type(ev).__name__
    deseriazlied_ev = cast(
        _TestEvent,
        deseriazlied_ev,
    )
    assert ev.param == deseriazlied_ev.param
    assert ev._data == deseriazlied_ev._data


def test_bool() -> None:
    assert bool(_TestEvent(param="foo")) is True


def test_stop_event_serialization() -> None:
    ev = StopEvent(result="foo")
    data_dict = ev.model_dump()
    assert data_dict == {"result": "foo"}

    serializer = JsonSerializer()
    serialized_ev = serializer.serialize(ev)
    deseriazlied_ev = serializer.deserialize(serialized_ev)

    assert type(deseriazlied_ev).__name__ == type(ev).__name__
    deseriazlied_ev = cast(
        StopEvent,
        deseriazlied_ev,
    )
    assert ev.result == deseriazlied_ev.result


class CustomStopEvent(StopEvent):
    foo: str
    bar: int


def test_custom_stop_event_serialization() -> None:
    ev = CustomStopEvent(foo="foo", bar=42)
    data_dict = ev.model_dump()
    assert data_dict == {"foo": "foo", "bar": 42}

    serializer = JsonSerializer()
    serialized_ev = serializer.serialize(ev)
    deserialized_ev = serializer.deserialize(serialized_ev)

    assert type(deserialized_ev).__name__ == type(ev).__name__
    deserialized_ev = cast(
        CustomStopEvent,
        deserialized_ev,
    )
    assert ev.foo == deserialized_ev.foo
    assert ev.bar == deserialized_ev.bar


def test_stop_event_repr() -> None:
    ev = StopEvent(foo="foo", result=42)
    assert repr(ev) == "StopEvent(foo='foo', result=42)"


def test_custom_stop_event_repr_no_result() -> None:
    ev = CustomStopEvent(foo="foo", bar=42)
    rep = repr(ev)
    assert rep == "CustomStopEvent(foo='foo', bar=42)"
