"""Unit tests for streaming functionality in multi-serve app."""

import asyncio
import hashlib
import tempfile
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import AsyncMock, patch

import pytest
from asgi_lifespan import LifespanManager
from httpx import ASGITransport, AsyncClient
from lfx.cli.serve_app import FlowMeta, StreamRequest, create_multi_serve_app
from lfx.interface.components import component_cache


def _make_settings_service(*, allow_custom_components: bool = False):
    return SimpleNamespace(
        settings=SimpleNamespace(
            allow_custom_components=allow_custom_components,
        )
    )


def _blocked_raw_graph() -> dict:
    return {
        "nodes": [
            {
                "id": "node-1",
                "data": {
                    "id": "node-1",
                    "type": "TotallyCustom",
                    "node": {
                        "display_name": "Blocked Node",
                        "template": {
                            "code": {"value": "print('blocked')"},
                        },
                    },
                },
            }
        ],
        "edges": [],
    }


class MockNode:
    """Mock node for testing graph structure."""

    def __init__(self, node_id: str, node_type: str = "TestComponent", display_name: str | None = None):
        self.data = {
            "id": node_id,
            "type": node_type,
            "display_name": display_name or node_type,
            "description": f"Mock {node_type} component",
            "template": {
                "input_field": {"type": "str", "value": "default_value"},
                "output_field": {"type": "str", "value": ""},
            },
        }


class MockEdge:
    """Mock edge for testing graph structure."""

    def __init__(self, source: str, target: str):
        self.source = source
        self.target = target


class MockGraph:
    """Mock graph for testing."""

    def __init__(self, nodes=None, edges=None):
        self.nodes = nodes or {
            "input_node": MockNode("input_node", "ChatInput", "Chat Input"),
            "output_node": MockNode("output_node", "ChatOutput", "Chat Output"),
        }
        self.edges = edges or [MockEdge("input_node", "output_node")]
        self.raw_graph_data = {"nodes": [], "edges": []}


@pytest.fixture
def mock_graphs():
    """Create mock graphs for testing."""
    return {
        "flow1": MockGraph(),
        "flow2": MockGraph(
            nodes={
                "text_input": MockNode("text_input", "TextInput", "Text Input"),
                "processor": MockNode("processor", "Processor", "Text Processor"),
                "text_output": MockNode("text_output", "TextOutput", "Text Output"),
            },
            edges=[MockEdge("text_input", "processor"), MockEdge("processor", "text_output")],
        ),
    }


@pytest.fixture
def mock_metas():
    """Create mock metadata for testing."""
    return {
        "flow1": FlowMeta(
            id="flow1", relative_path="flow1.json", title="Test Flow 1", description="A simple test flow for chat"
        ),
        "flow2": FlowMeta(
            id="flow2", relative_path="flow2.json", title="Test Flow 2", description="A test flow with text processing"
        ),
    }


@pytest.fixture
def multi_serve_app(mock_graphs, mock_metas, monkeypatch):
    """Create a multi-serve app for testing."""
    from lfx.services.deps import get_settings_service

    # Set required environment variable
    monkeypatch.setenv("LANGFLOW_API_KEY", "test-api-key")
    monkeypatch.setattr(get_settings_service().settings, "allow_custom_components", True)

    with patch("lfx.cli.serve_app.execute_graph_with_capture") as mock_execute:
        # Mock successful execution
        mock_execute.return_value = (
            [{"result": "Test response", "type": "message"}],
            "Execution completed successfully",
        )

        with tempfile.TemporaryDirectory() as temp_dir:
            app = create_multi_serve_app(
                root_dir=Path(temp_dir), graphs=mock_graphs, metas=mock_metas, verbose_print=lambda _: None
            )

            # Override the dependency after app creation
            def mock_verify_api_key(query_param: str | None = None, header_param: str | None = None) -> str:  # noqa: ARG001
                return "test-api-key"

            # Import the original dependency
            from lfx.cli.serve_app import verify_api_key

            app.dependency_overrides[verify_api_key] = mock_verify_api_key

            yield app

            # Clean up
            app.dependency_overrides.clear()


@pytest.fixture
def mock_api_key(monkeypatch):
    """Mock API key for authentication."""
    # Set the required environment variable
    monkeypatch.setenv("LANGFLOW_API_KEY", "test-api-key")

    with patch("lfx.cli.serve_app.verify_api_key") as mock_verify:
        mock_verify.return_value = True
        yield "test-api-key"


class TestMultiServeStreaming:
    """Test cases for multi-serve streaming functionality."""

    @pytest.mark.asyncio
    async def test_stream_endpoint_exists(self, multi_serve_app, mock_api_key):
        """Test that streaming endpoints are properly created."""
        async with (
            LifespanManager(multi_serve_app, startup_timeout=None, shutdown_timeout=None) as manager,
            AsyncClient(transport=ASGITransport(app=manager.app), base_url="http://testserver/", http2=True) as client,
        ):
            # Test that stream endpoints exist for each flow
            response = await client.post(
                "/flows/flow1/stream", json={"input_value": "Hello, world!"}, headers={"x-api-key": mock_api_key}
            )
            # Should not be 404 (endpoint exists)
            assert response.status_code != 404

    @pytest.mark.asyncio
    async def test_stream_basic_functionality(self, multi_serve_app, mock_api_key):
        """Test basic streaming functionality."""
        with patch("lfx.cli.serve_app.run_flow_generator_for_serve") as mock_generator:
            # Mock the streaming generator
            async def mock_stream_generator(*args, **kwargs):  # noqa: ARG001
                event_manager = kwargs.get("event_manager")
                client_consumed_queue = kwargs.get("client_consumed_queue")
                if event_manager:
                    event_manager.on_end(data={"result": {"result": "Streamed response", "success": True}})
                    if client_consumed_queue:
                        await client_consumed_queue.get()
                    else:
                        msg = "client_consumed_queue is None"
                        raise RuntimeError(msg)
                    # Send the final None to close the stream
                    await event_manager.queue.put((None, None, 0))

            mock_generator.side_effect = mock_stream_generator

            async with (
                LifespanManager(multi_serve_app, startup_timeout=None, shutdown_timeout=None) as manager,
                AsyncClient(
                    transport=ASGITransport(app=manager.app), base_url="http://testserver/", http2=True
                ) as client,
            ):
                response = await client.post(
                    "/flows/flow1/stream",
                    json={"input_value": "Test streaming input"},
                    headers={"x-api-key": mock_api_key},
                )

                # Debug output removed to pass linting

                assert response.status_code == 200
                assert response.headers["content-type"] == "text/event-stream; charset=utf-8"

    @pytest.mark.asyncio
    async def test_stream_request_validation(self, multi_serve_app, mock_api_key):
        """Test StreamRequest model validation."""
        async with (
            LifespanManager(multi_serve_app, startup_timeout=None, shutdown_timeout=None) as manager,
            AsyncClient(transport=ASGITransport(app=manager.app), base_url="http://testserver/", http2=True) as client,
        ):
            # Test with minimal valid request
            response = await client.post(
                "/flows/flow1/stream", json={"input_value": "test"}, headers={"x-api-key": mock_api_key}
            )
            assert response.status_code == 200

            # Test with full request
            response = await client.post(
                "/flows/flow1/stream",
                json={
                    "input_value": "test input",
                    "input_type": "chat",
                    "output_type": "chat",
                    "session_id": "test-session-123",
                    "tweaks": {"component1": {"param1": "value1"}},
                },
                headers={"x-api-key": mock_api_key},
            )
            assert response.status_code == 200

    @pytest.mark.asyncio
    async def test_stream_authentication_required(self, multi_serve_app):
        """Test that streaming endpoints require authentication."""
        # Temporarily remove the auth override to test real auth behavior
        from lfx.cli.serve_app import verify_api_key

        override = multi_serve_app.dependency_overrides.pop(verify_api_key, None)

        try:
            async with (
                LifespanManager(multi_serve_app, startup_timeout=None, shutdown_timeout=None) as manager,
                AsyncClient(
                    transport=ASGITransport(app=manager.app), base_url="http://testserver/", http2=True
                ) as client,
            ):
                # Test without API key
                response = await client.post("/flows/flow1/stream", json={"input_value": "test"})
                # Should fail authentication
                assert response.status_code in [401, 403]
        finally:
            # Restore the override for other tests
            if override:
                multi_serve_app.dependency_overrides[verify_api_key] = override

    @pytest.mark.asyncio
    async def test_stream_flow_not_found(self, multi_serve_app, mock_api_key):
        """Test streaming with non-existent flow."""
        async with (
            LifespanManager(multi_serve_app, startup_timeout=None, shutdown_timeout=None) as manager,
            AsyncClient(transport=ASGITransport(app=manager.app), base_url="http://testserver/", http2=True) as client,
        ):
            response = await client.post(
                "/flows/nonexistent/stream", json={"input_value": "test"}, headers={"x-api-key": mock_api_key}
            )
            assert response.status_code == 404

    @pytest.mark.asyncio
    async def test_stream_error_handling(self, multi_serve_app, mock_api_key):
        """Test error handling in streaming endpoint."""
        with patch("lfx.cli.serve_app.run_flow_generator_for_serve") as mock_generator:
            # Mock an error in the generator that properly terminates the stream
            async def mock_error_generator(graph, input_request, flow_id, event_manager, client_consumed_queue):  # noqa: ARG001
                try:
                    msg = "Test error during streaming"
                    raise RuntimeError(msg)
                except Exception as e:
                    # Properly handle the error like the real function does
                    event_manager.on_error(data={"error": str(e)})
                finally:
                    # Always send termination signal
                    import time

                    await event_manager.queue.put((None, None, time.time()))

            mock_generator.side_effect = mock_error_generator

            async with (
                LifespanManager(multi_serve_app, startup_timeout=None, shutdown_timeout=None) as manager,
                AsyncClient(
                    transport=ASGITransport(app=manager.app), base_url="http://testserver/", http2=True
                ) as client,
            ):
                response = await client.post(
                    "/flows/flow1/stream", json={"input_value": "test"}, headers={"x-api-key": mock_api_key}
                )

                # Should still return 200 but with error stream
                assert response.status_code == 200
                assert response.headers["content-type"] == "text/event-stream; charset=utf-8"

    @pytest.mark.asyncio
    async def test_stream_multiple_flows(self, multi_serve_app, mock_api_key):
        """Test streaming with multiple flows."""
        async with (
            LifespanManager(multi_serve_app, startup_timeout=None, shutdown_timeout=None) as manager,
            AsyncClient(transport=ASGITransport(app=manager.app), base_url="http://testserver/", http2=True) as client,
        ):
            # Test streaming for flow1
            response1 = await client.post(
                "/flows/flow1/stream", json={"input_value": "test flow 1"}, headers={"x-api-key": mock_api_key}
            )
            assert response1.status_code == 200

            # Test streaming for flow2
            response2 = await client.post(
                "/flows/flow2/stream", json={"input_value": "test flow 2"}, headers={"x-api-key": mock_api_key}
            )
            assert response2.status_code == 200

    @pytest.mark.asyncio
    async def test_regular_run_endpoint_still_works(self, multi_serve_app, mock_api_key):
        """Test that regular run endpoints still work alongside streaming."""
        with patch("lfx.cli.serve_app.extract_result_data") as mock_extract:
            mock_extract.return_value = {
                "result": "Regular response",
                "success": True,
                "type": "message",
                "component": "test",
            }

            async with (
                LifespanManager(multi_serve_app, startup_timeout=None, shutdown_timeout=None) as manager,
                AsyncClient(
                    transport=ASGITransport(app=manager.app), base_url="http://testserver/", http2=True
                ) as client,
            ):
                response = await client.post(
                    "/flows/flow1/run", json={"input_value": "test regular run"}, headers={"x-api-key": mock_api_key}
                )

                assert response.status_code == 200
                assert response.headers["content-type"] == "application/json"
                data = response.json()
                assert data["result"] == "Regular response"
                assert data["success"] is True

    @pytest.mark.asyncio
    async def test_list_flows_endpoint(self, multi_serve_app):
        """Test that the flows listing endpoint works."""
        async with (
            LifespanManager(multi_serve_app, startup_timeout=None, shutdown_timeout=None) as manager,
            AsyncClient(transport=ASGITransport(app=manager.app), base_url="http://testserver/", http2=True) as client,
        ):
            response = await client.get("/flows")
            assert response.status_code == 200

            flows = response.json()
            assert len(flows) == 2
            assert any(flow["id"] == "flow1" for flow in flows)
            assert any(flow["id"] == "flow2" for flow in flows)

    def test_stream_request_model(self):
        """Test the StreamRequest model validation."""
        # Test minimal request
        request = StreamRequest(input_value="test")
        assert request.input_value == "test"
        assert request.input_type == "chat"  # default
        assert request.output_type == "chat"  # default
        assert request.session_id is None
        assert request.tweaks is None

        # Test full request
        request = StreamRequest(
            input_value="test input",
            input_type="text",
            output_type="debug",
            output_component="specific_component",
            session_id="session123",
            tweaks={"comp1": {"param1": "value1"}},
        )
        assert request.input_value == "test input"
        assert request.input_type == "text"
        assert request.output_type == "debug"
        assert request.output_component == "specific_component"
        assert request.session_id == "session123"
        assert request.tweaks == {"comp1": {"param1": "value1"}}

    @pytest.mark.asyncio
    async def test_concurrent_streaming(self, multi_serve_app, mock_api_key):
        """Test concurrent streaming requests."""
        async with (
            LifespanManager(multi_serve_app, startup_timeout=None, shutdown_timeout=None) as manager,
            AsyncClient(transport=ASGITransport(app=manager.app), base_url="http://testserver/", http2=True) as client,
        ):
            # Start multiple concurrent streaming requests
            tasks = []
            for i in range(3):
                task = asyncio.create_task(
                    client.post(
                        "/flows/flow1/stream",
                        json={"input_value": f"concurrent test {i}"},
                        headers={"x-api-key": mock_api_key},
                    )
                )
                tasks.append(task)

            # Wait for all requests to complete
            responses = await asyncio.gather(*tasks)

            # All should be successful
            for response in responses:
                assert response.status_code == 200
                assert response.headers["content-type"] == "text/event-stream; charset=utf-8"

    @pytest.mark.asyncio
    async def test_stream_blocks_custom_components_when_disabled(
        self,
        mock_graphs,
        mock_metas,
        monkeypatch,
    ):
        """Test that /stream returns an error stream when validation blocks the flow."""
        monkeypatch.setenv("LANGFLOW_API_KEY", "test-api-key")
        mock_graphs["flow1"].raw_graph_data = _blocked_raw_graph()

        app = create_multi_serve_app(
            root_dir=Path("/tmp"),
            graphs=mock_graphs,
            metas=mock_metas,
            verbose_print=lambda _: None,
        )

        from lfx.cli.serve_app import verify_api_key

        def mock_verify_api_key() -> str:
            return "test-api-key"

        app.dependency_overrides[verify_api_key] = mock_verify_api_key

        async with (
            LifespanManager(app, startup_timeout=None, shutdown_timeout=None) as manager,
            AsyncClient(
                transport=ASGITransport(app=manager.app),
                base_url="http://testserver/",
                http2=True,
            ) as client,
        ):
            with (
                patch(
                    "lfx.services.deps.get_settings_service",
                    return_value=_make_settings_service(allow_custom_components=False),
                ),
                patch(
                    "lfx.utils.flow_validation.ensure_component_hash_lookups_loaded",
                    new=AsyncMock(return_value={"ChatInput": {hashlib.sha256(b"known").hexdigest()[:12]}}),
                ),
                patch.object(
                    component_cache,
                    "type_to_current_hash",
                    {"ChatInput": {hashlib.sha256(b"known").hexdigest()[:12]}},
                ),
                patch(
                    "lfx.cli.serve_app.run_flow_generator_for_serve",
                    new=AsyncMock(),
                ) as mock_run_flow,
            ):
                response = await client.post(
                    "/flows/flow1/stream",
                    json={"input_value": "test"},
                    headers={"x-api-key": "test-api-key"},
                )

        assert response.status_code == 200
        assert "custom components are not allowed" in response.text
        mock_run_flow.assert_not_called()
