import asyncio
import logging

from faststream.context import Context
from faststream.nats import NatsRouter
from pydantic import BaseModel

from spark_advisor_analyzer.config import ContextKey
from spark_advisor_analyzer.orchestrator import AdviceOrchestrator
from spark_advisor_models.defaults import (
    NATS_ANALYSIS_RESULT_SUBJECT,
    NATS_ANALYSIS_RUN_AGENT_SUBJECT,
    NATS_ANALYSIS_RUN_SUBJECT,
)
from spark_advisor_models.model import AnalysisMode, AnalysisResult, JobAnalysis

router = NatsRouter()
logger = logging.getLogger(__name__)


class ErrorResponse(BaseModel):
    error: str


@router.subscriber(NATS_ANALYSIS_RUN_SUBJECT)
@router.publisher(NATS_ANALYSIS_RESULT_SUBJECT)
async def handle_analyze(
    job: JobAnalysis,
    orchestrator: AdviceOrchestrator = Context(ContextKey.ORCHESTRATOR),  # type: ignore[assignment]  # noqa: B008
) -> AnalysisResult | ErrorResponse:
    try:
        return await asyncio.to_thread(orchestrator.run, job)
    except Exception as e:
        logger.exception("Analysis failed for %s", job.app_id)
        return ErrorResponse(error=str(e))


@router.subscriber(NATS_ANALYSIS_RUN_AGENT_SUBJECT)
@router.publisher(NATS_ANALYSIS_RESULT_SUBJECT)
async def handle_agent_analyze(
    job: JobAnalysis,
    orchestrator: AdviceOrchestrator = Context(ContextKey.ORCHESTRATOR),  # type: ignore[assignment]  # noqa: B008
) -> AnalysisResult | ErrorResponse:
    try:
        return await asyncio.to_thread(orchestrator.run, job, mode=AnalysisMode.AGENT)
    except Exception as e:
        logger.exception("Agent analysis failed for %s", job.app_id)
        return ErrorResponse(error=str(e))
