import asyncio
import threading

from fastapi import WebSocketDisconnect

from src.web import task_manager as task_manager_module
from src.web.routes import websocket as websocket_routes
from src.web.task_manager import TaskManager


class DummyWebSocket:
    def __init__(self):
        self.sent = []
        self.accepted = False
        self._incoming: asyncio.Queue = asyncio.Queue()

    async def accept(self):
        self.accepted = True

    async def send_json(self, message):
        self.sent.append(message)

    async def receive_json(self):
        item = await self._incoming.get()
        if isinstance(item, Exception):
            raise item
        return item

    async def disconnect(self):
        await self._incoming.put(WebSocketDisconnect())


def _reset_task_manager_state():
    task_manager_module._log_queues.clear()
    task_manager_module._log_locks.clear()
    task_manager_module._ws_connections.clear()
    task_manager_module._ws_sent_index.clear()
    task_manager_module._task_status.clear()
    task_manager_module._task_cancelled.clear()
    task_manager_module._batch_status.clear()
    task_manager_module._batch_logs.clear()
    task_manager_module._batch_locks.clear()


def test_update_status_broadcasts_terminal_failed_with_error_when_loop_available():
    async def scenario():
        _reset_task_manager_state()
        manager = TaskManager()
        manager.set_loop(asyncio.get_running_loop())
        ws = DummyWebSocket()
        manager.register_websocket("task-failed", ws)

        manager.update_status("task-failed", "failed")
        await asyncio.sleep(0.05)

        assert ws.sent, "expected websocket status push for failed status"
        failed_messages = [m for m in ws.sent if m.get("type") == "status" and m.get("status") == "failed"]
        assert failed_messages, "expected failed status message"
        assert failed_messages[-1].get("error"), "failed status must include error"
        assert manager.get_status("task-failed").get("error")

    asyncio.run(scenario())


def test_update_status_broadcasts_terminal_completed_payload_when_loop_available():
    async def scenario():
        _reset_task_manager_state()
        manager = TaskManager()
        manager.set_loop(asyncio.get_running_loop())
        ws = DummyWebSocket()
        manager.register_websocket("task-completed", ws)

        manager.update_status("task-completed", "completed", email="done@example.com")
        await asyncio.sleep(0.05)

        completed_messages = [m for m in ws.sent if m.get("type") == "status" and m.get("status") == "completed"]
        assert completed_messages, "expected completed status message"
        assert completed_messages[-1].get("email") == "done@example.com"
        assert completed_messages[-1].get("updated_at")

    asyncio.run(scenario())


def test_task_websocket_sends_initial_status_snapshot():
    async def scenario():
        _reset_task_manager_state()
        manager = TaskManager()
        manager.update_status("task-initial", "running", message="booting")
        ws = DummyWebSocket()

        original = websocket_routes.task_manager
        websocket_routes.task_manager = manager
        try:
            task = asyncio.create_task(websocket_routes.task_websocket(ws, "task-initial"))
            await asyncio.sleep(0.01)
            await ws.disconnect()
            await task
        finally:
            websocket_routes.task_manager = original

        assert ws.accepted
        status_messages = [m for m in ws.sent if m.get("type") == "status"]
        assert status_messages, "expected initial status snapshot on connect"
        assert status_messages[0].get("status") == "running"
        assert status_messages[0].get("message") == "booting"

    asyncio.run(scenario())


def test_task_websocket_receives_terminal_update_after_connect():
    async def scenario():
        _reset_task_manager_state()
        manager = TaskManager()
        manager.set_loop(asyncio.get_running_loop())
        ws = DummyWebSocket()

        original = websocket_routes.task_manager
        websocket_routes.task_manager = manager
        try:
            task = asyncio.create_task(websocket_routes.task_websocket(ws, "task-live"))
            await asyncio.sleep(0.01)
            manager.update_status("task-live", "failed", error="boom")
            await asyncio.sleep(0.05)
            await ws.disconnect()
            await task
        finally:
            websocket_routes.task_manager = original

        terminal_messages = [m for m in ws.sent if m.get("type") == "status" and m.get("status") == "failed"]
        assert terminal_messages, "expected live failed status over websocket"
        assert terminal_messages[-1].get("error") == "boom"

    asyncio.run(scenario())


def test_update_status_broadcasts_from_worker_thread_when_loop_available():
    async def scenario():
        _reset_task_manager_state()
        manager = TaskManager()
        manager.set_loop(asyncio.get_running_loop())
        ws = DummyWebSocket()
        manager.register_websocket("task-threaded", ws)

        worker = threading.Thread(
            target=lambda: manager.update_status(
                "task-threaded",
                "completed",
                email="thread@example.com"
            )
        )
        worker.start()
        worker.join()
        await asyncio.sleep(0.05)

        completed_messages = [
            m for m in ws.sent
            if m.get("type") == "status" and m.get("status") == "completed"
        ]
        assert completed_messages, "expected threaded status push over websocket"
        assert completed_messages[-1].get("email") == "thread@example.com"
        assert completed_messages[-1].get("message")

    asyncio.run(scenario())


def test_update_status_terminal_defaults_include_useful_fields():
    _reset_task_manager_state()
    manager = TaskManager()

    manager.update_status("task-term", "completed")
    completed = manager.get_status("task-term")
    assert completed["message"]

    manager.update_status("task-term", "failed")
    failed = manager.get_status("task-term")
    assert failed["message"]
    assert failed["error"]

    manager.update_status("task-term", "cancelled")
    cancelled = manager.get_status("task-term")
    assert cancelled["message"]

    manager.update_status("task-term", "cancelling")
    cancelling = manager.get_status("task-term")
    assert cancelling["message"]