"""Tests for ChatKit server workflow execution pathways."""

from __future__ import annotations
from datetime import UTC, datetime
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from chatkit.errors import CustomStreamError
from chatkit.types import ThreadMetadata
from pydantic import BaseModel
from orcheo.models.workflow import WorkflowDraftAccess
from orcheo_backend.app.chatkit import OrcheoChatKitServer
from orcheo_backend.app.repository import InMemoryWorkflowRepository
from tests.backend.chatkit_test_utils import (
    create_chatkit_test_server,
    create_workflow_with_graph,
)


@pytest.mark.asyncio
async def test_chatkit_server_run_workflow_end_to_end() -> None:
    repository = InMemoryWorkflowRepository()
    workflow = await repository.create_workflow(
        name="Test workflow",
        slug=None,
        description=None,
        tags=None,
        draft_access=WorkflowDraftAccess.PERSONAL,
        actor="tester",
    )

    graph_config = {
        "format": "langgraph-script",
        "source": """
from langgraph.graph import END, START, StateGraph

def build_graph():
    graph = StateGraph(dict)

    def respond(state):
        message = state.get("message", "")
        return {"reply": f"Echo: {message}"}

    graph.add_node("respond", respond)
    graph.add_edge(START, "respond")
    graph.add_edge("respond", END)
    return graph
""",
        "entrypoint": "build_graph",
    }

    await repository.create_version(
        workflow.id,
        graph=graph_config,
        metadata={},
        notes=None,
        created_by="tester",
    )

    server = create_chatkit_test_server(repository)
    reply, state, run = await server._run_workflow(workflow.id, {"message": "Test"})

    assert reply == "Echo: Test"
    assert isinstance(state, dict)
    assert run is not None


@pytest.mark.asyncio
async def test_chatkit_server_run_workflow_without_reply() -> None:
    repository = InMemoryWorkflowRepository()
    workflow = await repository.create_workflow(
        name="Test workflow",
        slug=None,
        description=None,
        tags=None,
        draft_access=WorkflowDraftAccess.PERSONAL,
        actor="tester",
    )

    graph_config = {
        "format": "langgraph-script",
        "source": """
from langgraph.graph import END, START, StateGraph

def build_graph():
    graph = StateGraph(dict)

    def no_reply(state):
        return {"output": "something else"}

    graph.add_node("no_reply", no_reply)
    graph.add_edge(START, "no_reply")
    graph.add_edge("no_reply", END)
    return graph
""",
        "entrypoint": "build_graph",
    }

    await repository.create_version(
        workflow.id,
        graph=graph_config,
        metadata={},
        notes=None,
        created_by="tester",
    )

    server = create_chatkit_test_server(repository)

    with pytest.raises(CustomStreamError, match="without producing a reply"):
        await server._run_workflow(workflow.id, {})


@pytest.mark.asyncio
async def test_chatkit_server_run_workflow_with_basemodel_state() -> None:
    repository = InMemoryWorkflowRepository()
    workflow = await create_workflow_with_graph(repository)

    server = create_chatkit_test_server(repository)

    class TestState(BaseModel):
        reply: str

    mock_compiled = MagicMock()
    mock_compiled.ainvoke = AsyncMock(return_value=TestState(reply="Test reply"))

    with patch(
        "orcheo_backend.app.chatkit.workflow_executor.build_graph"
    ) as mock_build:
        mock_graph = MagicMock()
        mock_graph.compile.return_value = mock_compiled
        mock_build.return_value = mock_graph

        reply, state, run = await server._run_workflow(workflow.id, {"message": "Test"})

    assert reply == "Test reply"
    assert isinstance(state, dict)
    assert run is not None


def test_chatkit_server_records_run_metadata_without_run() -> None:
    thread = ThreadMetadata(
        id="thr_no_run",
        created_at=datetime.now(UTC),
        metadata={},
    )

    OrcheoChatKitServer._record_run_metadata(thread, None)

    assert "last_run_at" in thread.metadata
    assert "last_run_id" not in thread.metadata


@pytest.mark.asyncio
async def test_chatkit_server_run_workflow_with_repository_create_run_failure() -> None:
    repository = InMemoryWorkflowRepository()
    workflow = await create_workflow_with_graph(repository)

    server = create_chatkit_test_server(repository)

    original_create_run = server._repository.create_run
    server._repository.create_run = AsyncMock(side_effect=Exception("DB error"))

    reply, state, run = await server._run_workflow(workflow.id, {"message": "Test"})

    assert reply == "Echo: Test"
    assert isinstance(state, dict)
    assert run is None

    server._repository.create_run = original_create_run


def test_chatkit_server_backfills_workflow_id_from_context() -> None:
    repository = InMemoryWorkflowRepository()
    server = create_chatkit_test_server(repository)

    thread = ThreadMetadata(
        id="thr_missing_workflow",
        created_at=datetime.now(UTC),
        metadata={},
    )
    context = {"workflow_id": "wf-123"}

    server._ensure_workflow_metadata(thread, context)

    assert thread.metadata["workflow_id"] == "wf-123"
