import asyncio
import os
import tempfile
import unittest
from unittest.mock import patch

_TEST_DATA_DIR = None
if "SOVEREIGNAI_DATA_DIR" not in os.environ:
    _TEST_DATA_DIR = tempfile.TemporaryDirectory(prefix="sovereignai-stream-tests-")
    os.environ["SOVEREIGNAI_DATA_DIR"] = _TEST_DATA_DIR.name

from app import main


class AgentStreamTests(unittest.IsolatedAsyncioTestCase):
    async def test_idle_model_wait_sends_heartbeat_then_final_result(self) -> None:
        async def slow_agent(request, report_progress=None, *, user_id):
            await report_progress("بدأ تحليل المهمة.")
            await asyncio.sleep(0.05)
            return {"task": request.task, "result": "اكتمل التحليل."}

        with (
            patch.object(main, "AGENT_STREAM_HEARTBEAT_SECONDS", 0.01),
            patch.object(main, "_execute_agent", side_effect=slow_agent),
        ):
            response = await main.run_agent_stream(
                main.AgentRequest(task="اختبار انتظار النموذج"), user_id="test-user"
            )
            chunks = [chunk async for chunk in response.body_iterator]

        body = b"".join(
            chunk if isinstance(chunk, bytes) else chunk.encode("utf-8")
            for chunk in chunks
        ).decode("utf-8")
        self.assertIn("event: progress", body)
        self.assertIn("event: heartbeat\ndata: {}", body)
        self.assertIn("event: done", body)
        self.assertIn("اكتمل التحليل", body)


if __name__ == "__main__":
    unittest.main()