CIAgent/apps/api/app/repositories/company_repository.py

from __future__ import annotations

import uuid
from datetime import datetime

from sqlalchemy import func, select
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import selectinload

from app.models.company import Company, CompanyAlias, Competitor
from app.models.monitor_configuration import MonitorConfiguration


def _with_relations(stmt):
    return stmt.options(
        selectinload(Company.aliases),
        selectinload(Company.competitors),
        selectinload(Company.monitor_configuration),
        selectinload(Company.enrichment),
        # Loaded eagerly so ORM-level cascade-delete (SQLite doesn't enforce
        # FK ON DELETE CASCADE without a pragma this app doesn't set) can
        # actually see the children to remove when a Company is deleted.
        selectinload(Company.notification_links),
    )


class CompanyRepository:
    def __init__(self, db: AsyncSession) -> None:
        self.db = db

    async def get_for_user(self, company_id: uuid.UUID, user_id: uuid.UUID) -> Company | None:
        stmt = _with_relations(
            select(Company).where(Company.id == company_id, Company.user_id == user_id)
        )
        result = await self.db.execute(stmt)
        return result.scalar_one_or_none()

    async def get_by_id(self, company_id: uuid.UUID) -> Company | None:
        """Unscoped lookup for trusted internal callers (Celery tasks) that
        already have the company_id from a source they control - not for
        anything reachable from an HTTP request without an ownership check."""
        stmt = _with_relations(select(Company).where(Company.id == company_id))
        result = await self.db.execute(stmt)
        return result.scalar_one_or_none()

    async def list_for_user(self, user_id: uuid.UUID) -> list[Company]:
        stmt = _with_relations(
            select(Company).where(Company.user_id == user_id).order_by(Company.created_at.desc())
        )
        result = await self.db.execute(stmt)
        return list(result.scalars().all())

    async def count_for_user(self, user_id: uuid.UUID) -> int:
        result = await self.db.execute(
            select(func.count()).select_from(Company).where(Company.user_id == user_id)
        )
        return int(result.scalar_one())

    async def slug_exists_for_user(self, user_id: uuid.UUID, slug: str) -> bool:
        result = await self.db.execute(
            select(func.count())
            .select_from(Company)
            .where(Company.user_id == user_id, Company.slug == slug)
        )
        return int(result.scalar_one()) > 0

    async def name_exists_for_user(self, user_id: uuid.UUID, name: str) -> bool:
        result = await self.db.execute(
            select(func.count())
            .select_from(Company)
            .where(Company.user_id == user_id, func.lower(Company.name) == name.lower())
        )
        return int(result.scalar_one()) > 0

    async def create(
        self,
        *,
        user_id: uuid.UUID,
        name: str,
        slug: str,
        official_website: str | None,
        description: str | None,
        monitoring_focus: str | None,
        industry: str | None,
        country: str | None,
        region: str | None,
        headquarters: str | None = None,
        public_identifiers: dict[str, str] | None = None,
        alias_names: list[str],
        competitor_names: list[str],
    ) -> Company:
        company = Company(
            user_id=user_id,
            name=name,
            slug=slug,
            official_website=official_website,
            description=description,
            monitoring_focus=monitoring_focus,
            industry=industry,
            country=country,
            region=region,
            headquarters=headquarters,
            public_identifiers=public_identifiers or {},
        )
        company.aliases = [CompanyAlias(alias=a) for a in alias_names if a.strip()]
        company.competitors = [Competitor(name=c) for c in competitor_names if c.strip()]
        self.db.add(company)
        await self.db.flush()
        return company

    async def delete(self, company: Company) -> None:
        await self.db.delete(company)
        await self.db.flush()

    async def replace_aliases(self, company: Company, alias_names: list[str]) -> None:
        company.aliases = [CompanyAlias(alias=a) for a in alias_names if a.strip()]

    async def replace_competitors(self, company: Company, competitor_names: list[str]) -> None:
        company.competitors = [Competitor(name=c) for c in competitor_names if c.strip()]


class MonitorConfigurationRepository:
    def __init__(self, db: AsyncSession) -> None:
        self.db = db

    async def create_default(
        self,
        *,
        company_id: uuid.UUID,
        frequency_type,
        interval_minutes: int | None,
        cron_expression: str | None,
        timezone: str,
        severity_threshold,
        next_run,
    ) -> MonitorConfiguration:
        config = MonitorConfiguration(
            company_id=company_id,
            frequency_type=frequency_type,
            interval_minutes=interval_minutes,
            cron_expression=cron_expression,
            timezone=timezone,
            severity_threshold=severity_threshold,
            next_run=next_run,
        )
        self.db.add(config)
        await self.db.flush()
        return config

    async def list_due(self, now: datetime) -> list[MonitorConfiguration]:
        """Enabled schedules whose next_run has arrived - what Celery Beat's
        sync_schedules task polls instead of requiring a static per-company
        beat_schedule entry (see ARCHITECTURE.md)."""
        result = await self.db.execute(
            select(MonitorConfiguration).where(
                MonitorConfiguration.enabled.is_(True),
                MonitorConfiguration.next_run.is_not(None),
                MonitorConfiguration.next_run <= now,
            )
        )
        return list(result.scalars().all())

    async def list_enabled(self) -> list[MonitorConfiguration]:
        """Every enabled schedule regardless of next_run - since a company
        can now also become due purely via a per-source frequency override
        even when its own default next_run isn't due yet (see
        SourceRepository.list_due_for_company), the scheduler needs this
        broader set to check, not just the ones already due on the
        company-level clock."""
        result = await self.db.execute(
            select(MonitorConfiguration).where(MonitorConfiguration.enabled.is_(True))
        )
        return list(result.scalars().all())