"""Test flow state persistence functionality."""

import os
from typing import Dict, List

import pytest
from crewai.flow.flow import Flow, FlowState, listen, start
from crewai.flow.persistence import persist
from crewai.flow.persistence.sqlite import SQLiteFlowPersistence
from pydantic import BaseModel


class TestState(FlowState):
    """Test state model with required id field."""

    counter: int = 0
    message: str = ""


def test_persist_decorator_saves_state(tmp_path, caplog):
    """Test that @persist decorator saves state in SQLite."""
    db_path = os.path.join(tmp_path, "test_flows.db")
    persistence = SQLiteFlowPersistence(db_path)

    class TestFlow(Flow[Dict[str, str]]):
        initial_state = dict()

        @start()
        @persist(persistence)
        def init_step(self):
            self.state["message"] = "Hello, World!"
            self.state["id"] = "test-uuid"  # Ensure we have an ID for persistence

    flow = TestFlow(persistence=persistence)
    flow.kickoff()

    saved_state = persistence.load_state(flow.state["id"])
    assert saved_state is not None
    assert saved_state["message"] == "Hello, World!"


def test_structured_state_persistence(tmp_path):
    """Test persistence with Pydantic model state."""
    db_path = os.path.join(tmp_path, "test_flows.db")
    persistence = SQLiteFlowPersistence(db_path)

    class StructuredFlow(Flow[TestState]):
        initial_state = TestState

        @start()
        @persist(persistence)
        def count_up(self):
            self.state.counter += 1
            self.state.message = f"Count is {self.state.counter}"

    flow = StructuredFlow(persistence=persistence)
    flow.kickoff()

    saved_state = persistence.load_state(flow.state.id)
    assert saved_state is not None
    assert saved_state["counter"] == 1
    assert saved_state["message"] == "Count is 1"


def test_flow_state_restoration(tmp_path):
    """Test restoring flow state from persistence with various restoration methods."""
    db_path = os.path.join(tmp_path, "test_flows.db")
    persistence = SQLiteFlowPersistence(db_path)

    class RestorableFlow(Flow[TestState]):
        @start()
        @persist(persistence)
        def set_message(self):
            if self.state.message == "":
                self.state.message = "Original message"
            if self.state.counter == 0:
                self.state.counter = 42

    flow1 = RestorableFlow(persistence=persistence)
    flow1.kickoff()
    original_uuid = flow1.state.id

    flow2 = RestorableFlow(persistence=persistence)
    flow2.kickoff(inputs={"id": original_uuid, "counter": 43})

    assert flow2.state.id == original_uuid
    assert flow2.state.message == "Original message"
    assert flow2.state.counter == 43

    flow3 = RestorableFlow(persistence=persistence)
    flow3.kickoff(inputs={"id": original_uuid, "message": "Updated message"})

    assert flow3.state.id == original_uuid
    assert flow3.state.counter == 43
    assert flow3.state.message == "Updated message"


def test_multiple_method_persistence(tmp_path):
    """Test state persistence across multiple method executions."""
    db_path = os.path.join(tmp_path, "test_flows.db")
    persistence = SQLiteFlowPersistence(db_path)

    class MultiStepFlow(Flow[TestState]):
        @start()
        @persist(persistence)
        def step_1(self):
            if self.state.counter == 1:
                self.state.counter = 99999
                self.state.message = "Step 99999"
            else:
                self.state.counter = 1
                self.state.message = "Step 1"

        @listen(step_1)
        @persist(persistence)
        def step_2(self):
            if self.state.counter == 1:
                self.state.counter = 2
                self.state.message = "Step 2"

    flow = MultiStepFlow(persistence=persistence)
    flow.kickoff()

    flow2 = MultiStepFlow(persistence=persistence)
    flow2.kickoff(inputs={"id": flow.state.id})

    final_state = flow2.state
    assert final_state is not None
    assert final_state.counter == 2
    assert final_state.message == "Step 2"

    class NoPersistenceMultiStepFlow(Flow[TestState]):
        @start()
        @persist(persistence)
        def step_1(self):
            if self.state.counter == 1:
                self.state.counter = 99999
                self.state.message = "Step 99999"
            else:
                self.state.counter = 1
                self.state.message = "Step 1"

        @listen(step_1)
        def step_2(self):
            if self.state.counter == 1:
                self.state.counter = 2
                self.state.message = "Step 2"

    flow = NoPersistenceMultiStepFlow(persistence=persistence)
    flow.kickoff()

    flow2 = NoPersistenceMultiStepFlow(persistence=persistence)
    flow2.kickoff(inputs={"id": flow.state.id})

    final_state = flow2.state
    assert final_state.counter == 99999
    assert final_state.message == "Step 99999"


def test_persist_decorator_verbose_logging(tmp_path, caplog):
    """Test that @persist decorator's verbose parameter controls logging."""
    # Set logging level to ensure we capture all logs
    caplog.set_level("INFO")

    db_path = os.path.join(tmp_path, "test_flows.db")
    persistence = SQLiteFlowPersistence(db_path)

    class QuietFlow(Flow[Dict[str, str]]):
        initial_state = dict()

        @start()
        @persist(persistence)
        def init_step(self):
            self.state["message"] = "Hello, World!"
            self.state["id"] = "test-uuid-1"

    flow = QuietFlow(persistence=persistence)
    flow.kickoff()
    assert "Saving flow state" not in caplog.text

    caplog.clear()

    class VerboseFlow(Flow[Dict[str, str]]):
        initial_state = dict()

        @start()
        @persist(persistence, verbose=True)
        def init_step(self):
            self.state["message"] = "Hello, World!"
            self.state["id"] = "test-uuid-2"

    flow = VerboseFlow(persistence=persistence)
    flow.kickoff()
    assert "Saving flow state" in caplog.text


def test_persistence_with_base_model(tmp_path):
    db_path = os.path.join(tmp_path, "test_flows.db")
    persistence = SQLiteFlowPersistence(db_path)

    class Message(BaseModel):
        role: str
        type: str
        content: str

    class State(FlowState):
        latest_message: Message | None = None
        history: List[Message] = []

    @persist(persistence)
    class BaseModelFlow(Flow[State]):
        initial_state = State(latest_message=None, history=[])

        @start()
        def init_step(self):
            self.state.latest_message = Message(
                role="user", type="text", content="Hello, World!"
            )
            self.state.history.append(self.state.latest_message)

    flow = BaseModelFlow(persistence=persistence)
    flow.kickoff()

    latest_message = flow.state.latest_message
    (message,) = flow.state.history

    assert latest_message is not None
    assert latest_message.role == "user"
    assert latest_message.type == "text"
    assert latest_message.content == "Hello, World!"

    assert len(flow.state.history) == 1
    assert message.role == "user"
    assert message.type == "text"
    assert message.content == "Hello, World!"
    assert isinstance(flow.state, State)


def test_fork_with_restore_from_state_id(tmp_path):
    """Fork: restore_from_state_id hydrates state from source flow_uuid; new run gets a
    fresh state.id; source's history is preserved (the fork's @persist writes go under
    the new state.id, not the source's)."""
    db_path = os.path.join(tmp_path, "test_flows.db")
    persistence = SQLiteFlowPersistence(db_path)

    class ForkableFlow(Flow[TestState]):
        @start()
        @persist(persistence)
        def step(self):
            self.state.counter += 1

    flow1 = ForkableFlow(persistence=persistence)
    flow1.kickoff()
    source_uuid = flow1.state.id
    assert flow1.state.counter == 1

    # Resume on the same uuid bumps counter to 2 in the SAME flow_uuid history.
    flow1b = ForkableFlow(persistence=persistence)
    flow1b.kickoff(inputs={"id": source_uuid})
    assert flow1b.state.counter == 2
    assert persistence.load_state(source_uuid)["counter"] == 2

    # Fork: hydrate from source, but persist under a fresh state.id.
    flow2 = ForkableFlow(persistence=persistence)
    flow2.kickoff(restore_from_state_id=source_uuid)

    # Fork has a different state.id from the source.
    assert flow2.state.id != source_uuid
    # Hydrated from source's latest snapshot (counter=2), then incremented to 3.
    assert flow2.state.counter == 3

    # Source's history is unchanged after the fork.
    assert persistence.load_state(source_uuid)["counter"] == 2

    # Fork's writes landed under its own state.id.
    assert persistence.load_state(flow2.state.id)["counter"] == 3


def test_fork_with_pinned_state_id(tmp_path):
    """Fork into a pinned state.id (inputs.id supplied alongside restore_from_state_id):
    the new run uses inputs.id as state.id and hydrates from restore_from_state_id."""
    db_path = os.path.join(tmp_path, "test_flows.db")
    persistence = SQLiteFlowPersistence(db_path)

    class PinnableFlow(Flow[TestState]):
        @start()
        @persist(persistence)
        def step(self):
            self.state.counter += 1

    flow1 = PinnableFlow(persistence=persistence)
    flow1.kickoff()
    source_uuid = flow1.state.id
    assert flow1.state.counter == 1

    pinned_uuid = "pinned-fork-uuid-1234"
    flow2 = PinnableFlow(persistence=persistence)
    flow2.kickoff(
        inputs={"id": pinned_uuid},
        restore_from_state_id=source_uuid,
    )

    # state.id pinned to inputs.id, NOT the source uuid.
    assert flow2.state.id == pinned_uuid
    # Hydrated from source: counter started at 1, step incremented to 2.
    assert flow2.state.counter == 2
    # Source's history is unchanged.
    assert persistence.load_state(source_uuid)["counter"] == 1
    # Fork's writes are under the pinned uuid.
    assert persistence.load_state(pinned_uuid)["counter"] == 2


def test_restore_from_state_id_not_found_silent_fallback(tmp_path):
    """Lookup miss on restore_from_state_id silently falls through to default behavior."""
    db_path = os.path.join(tmp_path, "test_flows.db")
    persistence = SQLiteFlowPersistence(db_path)

    class FallbackFlow(Flow[TestState]):
        @start()
        @persist(persistence)
        def step(self):
            self.state.counter += 1

    flow = FallbackFlow(persistence=persistence)
    # No source UUID exists — should not raise.
    flow.kickoff(restore_from_state_id="no-such-uuid")

    # Default state path: counter starts at 0 and step increments to 1.
    assert flow.state.counter == 1
    # state.id is the auto-generated one, NOT the missing source.
    assert flow.state.id != "no-such-uuid"


def test_restore_from_state_id_not_found_raises_when_strict(tmp_path):
    """raise_on_missing_state=True turns a restore_from_state_id lookup miss into a
    ValueError instead of a silent fresh run."""
    db_path = os.path.join(tmp_path, "test_flows.db")
    persistence = SQLiteFlowPersistence(db_path)

    class StrictFlow(Flow[TestState]):
        @start()
        @persist(persistence)
        def step(self):
            self.state.counter += 1

    flow = StrictFlow(persistence=persistence)
    with pytest.raises(ValueError) as excinfo:
        flow.kickoff(
            restore_from_state_id="no-such-uuid",
            raise_on_missing_state=True,
        )

    msg = str(excinfo.value)
    assert "restore_from_state_id" in msg
    assert "no-such-uuid" in msg
    # The flow never ran: state is untouched and the outcome signal reports the miss.
    assert flow.state.counter == 0
    assert flow.last_restore_succeeded is False


def test_restore_from_state_id_strict_with_valid_id_restores(tmp_path):
    """raise_on_missing_state=True does not change behavior when the source state
    exists: hydrate from the snapshot, mint a fresh state.id, report success."""
    db_path = os.path.join(tmp_path, "test_flows.db")
    persistence = SQLiteFlowPersistence(db_path)

    class StrictForkFlow(Flow[TestState]):
        @start()
        @persist(persistence)
        def step(self):
            self.state.counter += 1

    flow1 = StrictForkFlow(persistence=persistence)
    flow1.kickoff()
    source_uuid = flow1.state.id
    assert flow1.state.counter == 1

    flow2 = StrictForkFlow(persistence=persistence)
    flow2.kickoff(
        restore_from_state_id=source_uuid,
        raise_on_missing_state=True,
    )

    assert flow2.state.id != source_uuid
    assert flow2.state.counter == 2
    assert flow2.last_restore_succeeded is True


def test_strict_restore_without_persistence_raises():
    """raise_on_missing_state=True with restore_from_state_id but no persistence
    backend raises instead of silently skipping hydration."""

    class NoPersistenceFlow(Flow[TestState]):
        @start()
        def step(self):
            self.state.counter += 1

    flow = NoPersistenceFlow()
    with pytest.raises(ValueError) as excinfo:
        flow.kickoff(
            restore_from_state_id="some-uuid",
            raise_on_missing_state=True,
        )
    msg = str(excinfo.value)
    assert "restore_from_state_id" in msg
    assert "persistence" in msg
    assert flow.last_restore_succeeded is False


def test_stream_events_strict_restore_miss_raises_eagerly(tmp_path):
    """stream_events with raise_on_missing_state=True raises at session creation
    time on a restore_from_state_id miss, not while consuming frames."""
    db_path = os.path.join(tmp_path, "test_flows.db")
    persistence = SQLiteFlowPersistence(db_path)

    class StreamStrictFlow(Flow[TestState]):
        @start()
        @persist(persistence)
        def step(self):
            self.state.counter += 1

    flow = StreamStrictFlow(persistence=persistence)
    with pytest.raises(ValueError) as excinfo:
        flow.stream_events(
            restore_from_state_id="no-such-uuid",
            raise_on_missing_state=True,
        )
    assert "no-such-uuid" in str(excinfo.value)
    # Nothing ran.
    assert flow.state.counter == 0


def test_kickoff_stream_strict_restore_miss_raises_eagerly(tmp_path):
    """kickoff with stream=True and raise_on_missing_state=True raises immediately
    on a restore miss instead of returning a StreamSession."""
    db_path = os.path.join(tmp_path, "test_flows.db")
    persistence = SQLiteFlowPersistence(db_path)

    class StreamingKickoffFlow(Flow[TestState]):
        @start()
        @persist(persistence)
        def step(self):
            self.state.counter += 1

    flow = StreamingKickoffFlow(persistence=persistence)
    flow.stream = True
    with pytest.raises(ValueError) as excinfo:
        flow.kickoff(
            restore_from_state_id="no-such-uuid",
            raise_on_missing_state=True,
        )
    assert "no-such-uuid" in str(excinfo.value)
    assert flow.state.counter == 0


def test_stream_events_strict_restore_valid_id_streams(tmp_path):
    """stream_events with raise_on_missing_state=True streams normally when the
    source state exists: hydrates, mints a fresh state.id, reports success."""
    db_path = os.path.join(tmp_path, "test_flows.db")
    persistence = SQLiteFlowPersistence(db_path)

    class StreamForkFlow(Flow[TestState]):
        @start()
        @persist(persistence)
        def step(self):
            self.state.counter += 1

    flow1 = StreamForkFlow(persistence=persistence)
    flow1.kickoff()
    source_uuid = flow1.state.id

    flow2 = StreamForkFlow(persistence=persistence)
    with flow2.stream_events(
        restore_from_state_id=source_uuid,
        raise_on_missing_state=True,
    ) as stream:
        list(stream.events)

    assert flow2.state.id != source_uuid
    assert flow2.state.counter == 2
    assert flow2.last_restore_succeeded is True


@pytest.mark.asyncio
async def test_astream_strict_restore_miss_raises_eagerly(tmp_path):
    """astream and kickoff_async(stream=True) raise at call time on a restore miss
    with raise_on_missing_state=True."""
    db_path = os.path.join(tmp_path, "test_flows.db")
    persistence = SQLiteFlowPersistence(db_path)

    class AsyncStreamStrictFlow(Flow[TestState]):
        @start()
        @persist(persistence)
        def step(self):
            self.state.counter += 1

    flow = AsyncStreamStrictFlow(persistence=persistence)
    with pytest.raises(ValueError) as excinfo:
        flow.astream(
            restore_from_state_id="no-such-uuid",
            raise_on_missing_state=True,
        )
    assert "no-such-uuid" in str(excinfo.value)
    assert flow.state.counter == 0

    flow2 = AsyncStreamStrictFlow(persistence=persistence)
    flow2.stream = True
    with pytest.raises(ValueError):
        await flow2.kickoff_async(
            restore_from_state_id="no-such-uuid",
            raise_on_missing_state=True,
        )
    assert flow2.state.counter == 0


def test_last_restore_succeeded_signal(tmp_path):
    """last_restore_succeeded is None when no restore was requested, False after a
    silent miss, and True after a successful restore."""
    db_path = os.path.join(tmp_path, "test_flows.db")
    persistence = SQLiteFlowPersistence(db_path)

    class SignalFlow(Flow[TestState]):
        @start()
        @persist(persistence)
        def step(self):
            self.state.counter += 1

    flow1 = SignalFlow(persistence=persistence)
    assert flow1.last_restore_succeeded is None
    flow1.kickoff()
    # No restore requested: signal stays None.
    assert flow1.last_restore_succeeded is None
    source_uuid = flow1.state.id

    # Silent miss (default non-strict): the run proceeds but the signal is False.
    flow2 = SignalFlow(persistence=persistence)
    flow2.kickoff(restore_from_state_id="no-such-uuid")
    assert flow2.state.counter == 1
    assert flow2.last_restore_succeeded is False

    # Successful restore: the signal is True.
    flow3 = SignalFlow(persistence=persistence)
    flow3.kickoff(restore_from_state_id=source_uuid)
    assert flow3.state.counter == 2
    assert flow3.last_restore_succeeded is True

    # A later kickoff without a restore request resets the signal to None.
    flow3.kickoff()
    assert flow3.last_restore_succeeded is None


def test_restore_from_state_id_none_is_no_op(tmp_path):
    """restore_from_state_id=None (default) preserves baseline kickoff behavior."""
    db_path = os.path.join(tmp_path, "test_flows.db")
    persistence = SQLiteFlowPersistence(db_path)

    class BaselineFlow(Flow[TestState]):
        @start()
        @persist(persistence)
        def step(self):
            self.state.counter += 1

    flow = BaselineFlow(persistence=persistence)
    flow.kickoff(restore_from_state_id=None)
    assert flow.state.counter == 1


def test_fork_conflict_with_from_checkpoint_raises():
    """Passing both from_checkpoint and restore_from_state_id raises ValueError, naming
    both parameters."""
    from crewai.state import CheckpointConfig

    class ConflictFlow(Flow[TestState]):
        @start()
        def step(self):
            pass

    flow = ConflictFlow()
    with pytest.raises(ValueError) as excinfo:
        flow.kickoff(
            from_checkpoint=CheckpointConfig(),
            restore_from_state_id="some-uuid",
        )
    msg = str(excinfo.value)
    assert "from_checkpoint" in msg
    assert "restore_from_state_id" in msg


@pytest.mark.asyncio
async def test_fork_via_kickoff_async(tmp_path):
    """kickoff_async honors restore_from_state_id: hydrates from source, mints fresh
    state.id, persists under the new id, source history preserved."""
    db_path = os.path.join(tmp_path, "test_flows.db")
    persistence = SQLiteFlowPersistence(db_path)

    class AsyncForkableFlow(Flow[TestState]):
        @start()
        @persist(persistence)
        def step(self):
            self.state.counter += 1

    flow1 = AsyncForkableFlow(persistence=persistence)
    await flow1.kickoff_async()
    source_uuid = flow1.state.id
    assert flow1.state.counter == 1

    flow2 = AsyncForkableFlow(persistence=persistence)
    await flow2.kickoff_async(restore_from_state_id=source_uuid)

    assert flow2.state.id != source_uuid
    assert flow2.state.counter == 2
    assert persistence.load_state(source_uuid)["counter"] == 1
    assert persistence.load_state(flow2.state.id)["counter"] == 2


@pytest.mark.asyncio
async def test_fork_via_akickoff(tmp_path):
    """akickoff is the public async alias and must accept restore_from_state_id with
    the same semantics as kickoff_async."""
    db_path = os.path.join(tmp_path, "test_flows.db")
    persistence = SQLiteFlowPersistence(db_path)

    class AkickoffForkableFlow(Flow[TestState]):
        @start()
        @persist(persistence)
        def step(self):
            self.state.counter += 1

    flow1 = AkickoffForkableFlow(persistence=persistence)
    await flow1.akickoff()
    source_uuid = flow1.state.id
    assert flow1.state.counter == 1

    flow2 = AkickoffForkableFlow(persistence=persistence)
    await flow2.akickoff(restore_from_state_id=source_uuid)

    assert flow2.state.id != source_uuid
    assert flow2.state.counter == 2
    assert persistence.load_state(source_uuid)["counter"] == 1
    assert persistence.load_state(flow2.state.id)["counter"] == 2


@pytest.mark.asyncio
async def test_akickoff_pinned_fork(tmp_path):
    """akickoff with both inputs.id and restore_from_state_id pins state.id while
    hydrating from the source."""
    db_path = os.path.join(tmp_path, "test_flows.db")
    persistence = SQLiteFlowPersistence(db_path)

    class PinnableAsyncFlow(Flow[TestState]):
        @start()
        @persist(persistence)
        def step(self):
            self.state.counter += 1

    flow1 = PinnableAsyncFlow(persistence=persistence)
    await flow1.akickoff()
    source_uuid = flow1.state.id

    pinned_uuid = "pinned-akickoff-fork-uuid"
    flow2 = PinnableAsyncFlow(persistence=persistence)
    await flow2.akickoff(
        inputs={"id": pinned_uuid},
        restore_from_state_id=source_uuid,
    )

    assert flow2.state.id == pinned_uuid
    assert flow2.state.counter == 2
    assert persistence.load_state(source_uuid)["counter"] == 1
    assert persistence.load_state(pinned_uuid)["counter"] == 2


@pytest.mark.asyncio
async def test_restore_from_state_id_not_found_raises_when_strict_async(tmp_path):
    """kickoff_async and akickoff honor raise_on_missing_state on a lookup miss."""
    db_path = os.path.join(tmp_path, "test_flows.db")
    persistence = SQLiteFlowPersistence(db_path)

    class AsyncStrictFlow(Flow[TestState]):
        @start()
        @persist(persistence)
        def step(self):
            self.state.counter += 1

    flow = AsyncStrictFlow(persistence=persistence)
    with pytest.raises(ValueError) as excinfo:
        await flow.kickoff_async(
            restore_from_state_id="no-such-uuid",
            raise_on_missing_state=True,
        )
    assert "no-such-uuid" in str(excinfo.value)
    assert flow.state.counter == 0
    assert flow.last_restore_succeeded is False

    flow_alias = AsyncStrictFlow(persistence=persistence)
    with pytest.raises(ValueError):
        await flow_alias.akickoff(
            restore_from_state_id="no-such-uuid",
            raise_on_missing_state=True,
        )
    assert flow_alias.last_restore_succeeded is False


@pytest.mark.asyncio
async def test_restore_from_state_id_strict_with_valid_id_restores_async(tmp_path):
    """kickoff_async with raise_on_missing_state=True restores normally when the
    source state exists and reports success."""
    db_path = os.path.join(tmp_path, "test_flows.db")
    persistence = SQLiteFlowPersistence(db_path)

    class AsyncStrictForkFlow(Flow[TestState]):
        @start()
        @persist(persistence)
        def step(self):
            self.state.counter += 1

    flow1 = AsyncStrictForkFlow(persistence=persistence)
    await flow1.kickoff_async()
    source_uuid = flow1.state.id

    flow2 = AsyncStrictForkFlow(persistence=persistence)
    await flow2.kickoff_async(
        restore_from_state_id=source_uuid,
        raise_on_missing_state=True,
    )

    assert flow2.state.id != source_uuid
    assert flow2.state.counter == 2
    assert flow2.last_restore_succeeded is True


@pytest.mark.asyncio
async def test_akickoff_fork_conflict_with_from_checkpoint_raises():
    """akickoff must raise the same conflict ValueError as kickoff/kickoff_async when
    both from_checkpoint and restore_from_state_id are set."""
    from crewai.state import CheckpointConfig

    class AsyncConflictFlow(Flow[TestState]):
        @start()
        def step(self):
            pass

    flow = AsyncConflictFlow()
    with pytest.raises(ValueError) as excinfo:
        await flow.akickoff(
            from_checkpoint=CheckpointConfig(),
            restore_from_state_id="some-uuid",
        )
    msg = str(excinfo.value)
    assert "from_checkpoint" in msg
    assert "restore_from_state_id" in msg