from __future__ import annotations

import logging

from app.core.api.infrastructure.streaming.sse_response_builder import build_sse_response
from app.core.rag.module import RagModule
from app.core.shared.messaging import EventBus
from app.schemas.rag_sessions import RagSessionJobResponse

LOGGER = logging.getLogger(__name__)


class RagPublicController:
    def __init__(self, rag: RagModule) -> None:
        self._rag = rag

    def get_job(self, rag_session_id: str, index_job_id: str) -> RagSessionJobResponse:
        job = self._rag.get_session_job(rag_session_id, index_job_id)
        LOGGER.warning(
            "rag job polled: rag_session_id=%s job_id=%s status=%s indexed=%s failed=%s",
            rag_session_id,
            index_job_id,
            job.status.value if hasattr(job.status, "value") else str(job.status),
            job.indexed_files,
            job.failed_files,
        )
        return RagSessionJobResponse(
            rag_session_id=rag_session_id,
            index_job_id=job.index_job_id,
            status=job.status,
            indexed_files=job.indexed_files,
            failed_files=job.failed_files,
            cache_hit_files=job.cache_hit_files,
            cache_miss_files=job.cache_miss_files,
            error=job.error.model_dump(mode="json") if job.error else None,
        )

    async def stream_job_events(self, rag_session_id: str, index_job_id: str):
        channel_id, queue = await self._rag.subscribe_session_job_events(rag_session_id, index_job_id)
        LOGGER.warning("rag job events subscribed: rag_session_id=%s job_id=%s", rag_session_id, index_job_id)
        return build_sse_response(
            queue,
            encoder=EventBus.as_sse,
            unsubscribe=lambda: self._rag.unsubscribe_job_events(channel_id, queue),
            stop_on_event="terminal",
        )
