from __future__ import annotations

import logging
from datetime import datetime, timezone

from app.core.api.application.session_service import SessionService
from app.core.api.domain.models.agent_request import AgentRequest
from app.core.api.domain.models.agent_session import AgentSession
from app.core.api.infrastructure.stores.in_memory_request_store import InMemoryRequestStore
from app.core.agent.runtime.execution_context import RuntimeExecutionContext
from app.core.agent.runtime.process_registry import ProcessRegistry
from app.core.agent.runtime.process_runner import ProcessRunner
from app.core.agent.runtime.publisher import RuntimeEventPublisher
from app.infra.exceptions import AppError
from app.infra.observability.module_trace import RequestTraceContext
from app.infra.observability.request_trace_logger import RequestTraceLogger
from app.schemas.common import ErrorPayload, ModuleName
from app.schemas.orchestration import RequestExecutionStatus

LOGGER = logging.getLogger(__name__)


class AgentRuntime:
    def __init__(
        self,
        request_store: InMemoryRequestStore,
        sessions: SessionService,
        process_registry: ProcessRegistry,
        process_runner: ProcessRunner,
        publisher: RuntimeEventPublisher,
        trace_logger: RequestTraceLogger,
    ) -> None:
        self._request_store = request_store
        self._sessions = sessions
        self._process_registry = process_registry
        self._process_runner = process_runner
        self._publisher = publisher
        self._trace_logger = trace_logger

    async def run(self, request: AgentRequest, session: AgentSession) -> None:
        LOGGER.warning(
            "runtime run started: request_id=%s process_version=%s active_rag_session_id=%s",
            request.request_id,
            request.process_version,
            session.active_rag_session_id,
        )
        try:
            process = self._resolve_process(request.process_version)
            self._start_request(request, session)
            context = RuntimeExecutionContext(
                request=request,
                session=session,
                publisher=self._publisher,
                trace=RequestTraceContext(request_id=request.request_id, logger=self._trace_logger),
            )
            await self._announce_start(request.request_id, process.version)
            result = await self._process_runner.run(context, process)
            request.answer = result.answer
            request.changeset = list(result.changeset)
            request.apply_changeset = bool(result.apply_changeset)
            await self._publish_result(request)
            self._complete_request(request, session)
            LOGGER.warning(
                "runtime run completed: request_id=%s status=%s changeset_items=%s apply_changeset=%s",
                request.request_id,
                request.status,
                len(request.changeset),
                request.apply_changeset,
            )
        except Exception as exc:
            LOGGER.exception(
                "runtime run failed: request_id=%s process_version=%s",
                request.request_id,
                request.process_version,
            )
            await self._fail_request(request, exc)

    def _resolve_process(self, version: str):
        process = self._process_registry.get(version)
        if process is None:
            raise AppError("process_not_found", f"Unsupported process version: {version}", ModuleName.AGENT)
        return process

    def _start_request(self, request: AgentRequest, session: AgentSession) -> None:
        request.status = RequestExecutionStatus.RUNNING
        self._request_store.save(request)
        self._trace_logger.start_request(request, session)

    async def _announce_start(self, request_id: str, process_version: str) -> None:
        await self._safe_publish_status(request_id, "runtime", "Запрос принят и поставлен в обработку.")
        await self._safe_publish_status(
            request_id,
            "runtime",
            f"Запускаю процесс {process_version}.",
            {"process_version": process_version},
        )

    async def _publish_result(self, request: AgentRequest) -> None:
        try:
            await self._publisher.publish_user(request.request_id, "agent", request.answer or "")
        except Exception:
            LOGGER.exception("failed to publish user event: request_id=%s", request.request_id)
        await self._safe_publish_status(request.request_id, "runtime", "Обработка запроса завершена.")

    def _complete_request(self, request: AgentRequest, session: AgentSession) -> None:
        session.append_turn(user_message=request.message, assistant_message=request.answer or "")
        self._sessions.save(session)
        request.status = RequestExecutionStatus.DONE
        request.completed_at = datetime.now(timezone.utc)
        self._request_store.save(request)
        self._trace_logger.complete_request(request)

    async def _fail_request(self, request: AgentRequest, exc: Exception) -> None:
        request.status = RequestExecutionStatus.ERROR
        request.completed_at = datetime.now(timezone.utc)
        request.error = self._build_error_payload(exc)
        self._request_store.save(request)
        self._trace_logger.fail_request(request)
        await self._safe_publish_status(
            request.request_id,
            "runtime",
            "Во время обработки запроса произошла ошибка.",
            {"code": request.error.code},
        )

    def _build_error_payload(self, exc: Exception) -> ErrorPayload:
        if isinstance(exc, AppError):
            return ErrorPayload(code=exc.code, desc=exc.desc, module=exc.module)
        return ErrorPayload(
            code="api_runtime_error",
            desc="Agent request failed unexpectedly.",
            module=ModuleName.AGENT,
        )

    async def _safe_publish_status(self, request_id: str, source: str, text: str, payload: dict | None = None) -> None:
        try:
            await self._publisher.publish_status(request_id, source, text, payload)
        except Exception:
            LOGGER.exception(
                "failed to publish status event: request_id=%s source=%s text=%s",
                request_id,
                source,
                text,
            )
