"""Tests for the FastAPI WebUI server."""

from __future__ import annotations

import asyncio
import json

import pytest
try:
    from fastapi.testclient import TestClient
except ImportError:  # pragma: no cover - optional dependency path
    TestClient = None

from agentbyte.session_store import FileSessionStore
from agentbyte.workflow import FunctionStep, StepMetadata, Workflow, WorkflowConfig
from agentbyte.webui.server import AgentbyteWebUIServer, _infer_entity_id

from .helpers import DummyAgent, DummyOrchestrator

pytestmark = pytest.mark.skipif(TestClient is None, reason="fastapi not installed")


def test_health_endpoint() -> None:
    server = AgentbyteWebUIServer()
    client = TestClient(server.create_app())

    response = client.get("/api/health")
    assert response.status_code == 200
    assert response.json()["status"] == "healthy"


def test_entities_endpoint_lists_registered_entities() -> None:
    server = AgentbyteWebUIServer()
    server.registry.register_entity("agent-1", DummyAgent())
    server.registry.register_entity("orch-1", DummyOrchestrator(DummyAgent()))
    client = TestClient(server.create_app())

    response = client.get("/api/entities")
    data = response.json()

    assert response.status_code == 200
    assert len(data) == 2
    assert {item["type"] for item in data} == {"agent", "orchestrator"}


def test_create_and_fetch_session() -> None:
    server = AgentbyteWebUIServer()
    client = TestClient(server.create_app())

    created = client.post(
        "/api/sessions",
        json={"entity_id": "agent-1", "entity_type": "agent", "user_id": "user-1"},
    )
    assert created.status_code == 200
    session_id = created.json()["id"]
    assert created.json()["user_id"] == "user-1"

    fetched = client.get(f"/api/sessions/{session_id}")
    assert fetched.status_code == 200
    assert fetched.json()["metadata"]["entity_id"] == "agent-1"
    assert fetched.json()["user_id"] == "user-1"


def test_sessions_endpoint_can_filter_by_user_id_and_entity_id() -> None:
    server = AgentbyteWebUIServer()
    client = TestClient(server.create_app())

    first = client.post(
        "/api/sessions",
        json={"entity_id": "agent-1", "entity_type": "agent", "user_id": "alice"},
    )
    second = client.post(
        "/api/sessions",
        json={"entity_id": "agent-2", "entity_type": "agent", "user_id": "alice"},
    )
    third = client.post(
        "/api/sessions",
        json={"entity_id": "agent-1", "entity_type": "agent", "user_id": "bob"},
    )

    assert first.status_code == 200
    assert second.status_code == 200
    assert third.status_code == 200

    response = client.get("/api/sessions", params={"user_id": "alice", "entity_id": "agent-1"})

    assert response.status_code == 200
    payload = response.json()
    assert len(payload) == 1
    assert payload[0]["user_id"] == "alice"
    assert payload[0]["entity_id"] == "agent-1"


def test_run_endpoint_returns_conflict_for_entity_mismatch() -> None:
    server = AgentbyteWebUIServer()
    server.registry.register_entity("agent-1", DummyAgent())
    server.registry.register_entity("agent-2", DummyAgent())
    client = TestClient(server.create_app())

    created = client.post(
        "/api/sessions",
        json={"entity_id": "agent-1", "entity_type": "agent", "user_id": "alice"},
    )
    session_id = created.json()["id"]

    response = client.post(
        "/api/entities/agent-2/run",
        json={
            "session_id": session_id,
            "user_id": "alice",
            "messages": [{"role": "user", "content": "hello", "source": "user"}],
        },
    )

    assert response.status_code == 409
    assert "bound to entity agent-1" in response.json()["detail"]


def test_stream_endpoint_returns_conflict_for_user_mismatch() -> None:
    server = AgentbyteWebUIServer()
    server.registry.register_entity("agent-1", DummyAgent())
    client = TestClient(server.create_app())

    created = client.post(
        "/api/sessions",
        json={"entity_id": "agent-1", "entity_type": "agent", "user_id": "alice"},
    )
    session_id = created.json()["id"]

    response = client.post(
        "/api/entities/agent-1/run/stream",
        json={
            "session_id": session_id,
            "user_id": "bob",
            "messages": [{"role": "user", "content": "hello", "source": "user"}],
        },
    )

    assert response.status_code == 409
    assert "owned by user alice" in response.json()["detail"]


def test_stream_endpoint_uses_fastapi_sse_event_stream() -> None:
    server = AgentbyteWebUIServer()
    server.registry.register_entity("agent-1", DummyAgent())
    client = TestClient(server.create_app())

    with client.stream(
        "POST",
        "/api/entities/agent-1/run/stream",
        json={
            "messages": [{"role": "user", "content": "hello", "source": "user"}],
        },
    ) as response:
        assert response.status_code == 200
        assert response.headers["content-type"].startswith("text/event-stream")

        first_data_line = next(
            line for line in response.iter_lines() if isinstance(line, str) and line.startswith("data: ")
        )

    payload = json.loads(first_data_line.removeprefix("data: "))
    assert payload["event"]["role"] == "assistant"
    assert payload["event"]["content"] == "echo: hello"


def test_server_uses_injected_persistent_session_store(tmp_path) -> None:
    storage_dir = tmp_path / "sessions"

    first_server = AgentbyteWebUIServer(session_store=FileSessionStore(str(storage_dir)))
    first_client = TestClient(first_server.create_app())

    created = first_client.post(
        "/api/sessions",
        json={"entity_id": "agent-1", "entity_type": "agent", "user_id": "alice"},
    )
    assert created.status_code == 200
    session_id = created.json()["id"]

    second_server = AgentbyteWebUIServer(session_store=FileSessionStore(str(storage_dir)))
    second_client = TestClient(second_server.create_app())

    fetched = second_client.get(f"/api/sessions/{session_id}")

    assert fetched.status_code == 200
    assert fetched.json()["metadata"]["entity_id"] == "agent-1"
    assert fetched.json()["user_id"] == "alice"

    stored = asyncio.run(second_server.session_manager.get(session_id))
    assert stored is not None


def test_session_messages_include_orchestrator_final_result_and_stop_reason() -> None:
    server = AgentbyteWebUIServer()
    server.registry.register_entity("orch-1", DummyOrchestrator(DummyAgent()))
    client = TestClient(server.create_app())

    created = client.post(
        "/api/sessions",
        json={"entity_id": "orch-1", "entity_type": "orchestrator", "user_id": "alice"},
    )
    session_id = created.json()["id"]

    run_response = client.post(
        "/api/entities/orch-1/run",
        json={
            "session_id": session_id,
            "user_id": "alice",
            "messages": [{"role": "user", "content": "hello", "source": "user"}],
        },
    )
    assert run_response.status_code == 200

    response = client.get(f"/api/sessions/{session_id}/messages")
    assert response.status_code == 200
    payload = response.json()
    assert payload["session_id"] == session_id
    assert payload["messages"]
    assert payload["final_result"] == run_response.json()["final_result"]
    assert payload["stop_reason"] == run_response.json()["stop_message"]["content"]


def test_infer_entity_id_uses_workflow_id_property() -> None:
    workflow = Workflow(
        config=WorkflowConfig(
            name="request_triage_workflow",
            description="Demo workflow",
        )
    )
    workflow.add_step(
        FunctionStep(
            step_id="step_1",
            metadata=StepMetadata(name="Step 1"),
            input_type=dict,
            output_type=dict,
            func=lambda input_data, context: input_data,
        )
    )

    assert _infer_entity_id(workflow, 4) == "request_triage_workflow"
