nexus-3/backend/core/repositories/accounts/accounts.py

"""
Repository for Account model operations.
"""
from typing import Optional
from django.db.models import Q, Prefetch
from django.utils import timezone
from backend.core.models import Account, Service
from backend.core.repositories.base import BaseRepository


class AccountRepository(BaseRepository[Account]):
    """
    Repository for Account model operations.
    """
    model = Account

    @classmethod
    def get_by_customer(cls, customer_id: str):
        """
        Get accounts by customer.

        Args:
            customer_id: The customer ID

        Returns:
            QuerySet of accounts for the customer
        """
        return Account.objects.filter(customer_id=customer_id)

    @classmethod
    def get_active(cls):
        """
        Get all active accounts.

        Returns:
            QuerySet of active accounts
        """
        return Account.objects.filter(
            Q(end_date__isnull=True) | Q(end_date__gt=timezone.now().date())
        )

    @classmethod
    def get_active_by_customer(cls, customer_id: str):
        """
        Get active accounts by customer.

        Args:
            customer_id: The customer ID

        Returns:
            QuerySet of active accounts for the customer
        """
        return Account.objects.filter(
            customer_id=customer_id
        ).filter(
            Q(end_date__isnull=True) | Q(end_date__gt=timezone.now().date())
        )

    @classmethod
    def get_inactive(cls):
        """
        Get all inactive accounts.

        Returns:
            QuerySet of inactive accounts
        """
        return Account.objects.filter(
            end_date__isnull=False,
            end_date__lte=timezone.now().date()
        )

    @classmethod
    def mark_inactive(cls, id: str) -> Optional[Account]:
        """
        Mark an account as inactive.

        Args:
            id: The account ID

        Returns:
            The updated account or None if not found
        """
        return cls.update(id, {'end_date': timezone.now().date()})

    @classmethod
    def search(cls, search_term: str, include_inactive: bool = False):
        """
        Search accounts by name or contact info.

        Args:
            search_term: The search term
            include_inactive: Whether to include inactive accounts

        Returns:
            QuerySet of matching accounts
        """
        queryset = super().search(
            search_term,
            [
                'name',
                'primary_contact_first_name',
                'primary_contact_last_name',
                'primary_contact_email',
                'secondary_contact_first_name',
                'secondary_contact_last_name',
                'secondary_contact_email',
                'customer__name'
            ]
        )

        if not include_inactive:
            queryset = queryset.filter(
                Q(end_date__isnull=True) | Q(end_date__gt=timezone.now().date())
            )

        return queryset

    @classmethod
    def get_with_services(cls, id: str, upcoming_only: bool = False) -> Optional[Account]:
        """
        Get an account with prefetched services.

        Args:
            id: The account ID
            upcoming_only: Whether to include only upcoming services

        Returns:
            The account with services or None if not found
        """
        services_queryset = Service.objects.all()

        if upcoming_only:
            services_queryset = services_queryset.filter(date__gte=timezone.now().date())

        services_prefetch = Prefetch('services', queryset=services_queryset)

        try:
            return Account.objects.prefetch_related(services_prefetch).get(pk=id)
        except Account.DoesNotExist:
            return None

    @classmethod
    def get_with_revenues(cls, id: str) -> Optional[Account]:
        """
        Get an account with prefetched revenues.

        Args:
            id: The account ID

        Returns:
            The account with revenues or None if not found
        """
        try:
            return Account.objects.prefetch_related('revenues').get(pk=id)
        except Account.DoesNotExist:
            return None

    @classmethod
    def get_with_labors(cls, id: str) -> Optional[Account]:
        """
        Get an account with prefetched labors.

        Args:
            id: The account ID

        Returns:
            The account with labors or None if not found
        """
        try:
            return Account.objects.prefetch_related('labors').get(pk=id)
        except Account.DoesNotExist:
            return None

    @classmethod
    def get_with_schedules(cls, id: str) -> Optional[Account]:
        """
        Get an account with prefetched schedules

        Args:
            id: The account ID

        Returns:
            The account with schedules or None if not found
        """
        try:
            return Account.objects.prefetch_related('schedules').get(pk=id)
        except Account.DoesNotExist:
            return None

    @classmethod
    def get_with_projects(cls, id: str) -> Optional[Account]:
        """
        Get an account with prefetched projects.
        Args:
            id: The account ID
        Returns:
            The account with projects or None if not found
        """
        try:
            return Account.objects.prefetch_related('projects').get(pk=id)
        except Account.DoesNotExist:
            return None

    @classmethod
    def get_with_all_related(cls, id: str) -> Optional[Account]:
        """
        Get an account with all related data prefetched.

        Args:
            id: The account ID

        Returns:
            The account with all related data or None if not found
        """
        try:
            return Account.objects.prefetch_related(
                'services', 'revenues', 'labors', 'schedules.py', 'projects'
            ).select_related('customer').get(pk=id)
        except Account.DoesNotExist:
            return None