from __future__ import annotations

from datetime import datetime, timezone
from typing import Optional

from analogai.deepthink.application.value_objects.user_agent import UserAgent
from analogai.deepthink.features.notion_enrichment.enricher_service import NotionEnricherService
from analogai.deepthink.infrastructure.gateways.llm.gateway import LlmGateway
from analogai.deepthink.infrastructure.gateways.llm.llm_factory import LlmGatewayFactory
from analogai.deepthink.messaging.notion.notion_enrichment_repository import NotionEnrichmentSleepRepository
from analogai.deepthink.startup.factories.domain_service_factory import DomainServiceFactory


class NotionEnrichmentSleeper:
    def __init__(
        self,
        domain_service_factory: DomainServiceFactory,
        repository: NotionEnrichmentSleepRepository,
        llm_gateway_instance: Optional[LlmGateway] = None,
        max_notions_per_sleep: Optional[int] = None,
    ):
        self._domain_service_factory = domain_service_factory
        self._repository = repository
        self._llm_gateway_instance = llm_gateway_instance
        self._max_notions_per_sleep = max_notions_per_sleep

    def sleep_for_agent(self, user_agent: UserAgent) -> int:
        user_id = user_agent.user_id
        agent_id = user_agent.agent_id
        last_sleep_at = self._repository.get_last_sleep_at(user_id, agent_id)
        notions = self._repository.list_unenriched_notions_created_after(
            user_id=user_id,
            agent_id=agent_id,
            since=last_sleep_at,
            limit=self._max_notions_per_sleep,
        )
        enricher = self._build_enricher()
        processed = 0
        for notion in notions:
            enricher.handle_notion_created(notion)
            self._repository.mark_enriched(user_id, agent_id, notion)
            processed += 1
        self._repository.set_last_sleep_at(user_id, agent_id, datetime.now(timezone.utc))
        return processed

    def _build_enricher(self) -> NotionEnricherService:
        if self._llm_gateway_instance is not None:
            LlmGatewayFactory.set_default_gateway_instance(self._llm_gateway_instance)
        completion_provider = LlmGatewayFactory.create_notion_enrichment()
        statement_service = self._domain_service_factory.get_direct_statement_service()
        notion_service = self._domain_service_factory.get_notion_service()
        return NotionEnricherService(
            statement_service=statement_service,
            notion_service=notion_service,
            completion_provider=completion_provider,
        )
