from __future__ import annotations

import uuid
from datetime import datetime
from typing import Any

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.models.company_enrichment import CompanyEnrichment
from app.models.enums import EnrichmentStatus


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

    async def get_for_company(self, company_id: uuid.UUID) -> CompanyEnrichment | None:
        result = await self.db.execute(
            select(CompanyEnrichment).where(CompanyEnrichment.company_id == company_id)
        )
        return result.scalar_one_or_none()

    async def upsert(
        self,
        company_id: uuid.UUID,
        *,
        status: EnrichmentStatus,
        data: dict[str, Any],
        errors: dict[str, str],
        credits_spent: int | None,
        fetched_at: datetime,
    ) -> CompanyEnrichment:
        """One row per company (unique FK) - the second-ever call for a
        company would only happen via a manual re-run, never automatically
        (enrichment fires once, at onboarding - see enrichment_service)."""
        existing = await self.get_for_company(company_id)
        if existing is not None:
            existing.status = status
            existing.data = data
            existing.errors = errors
            existing.credits_spent = credits_spent
            existing.fetched_at = fetched_at
            await self.db.flush()
            return existing

        enrichment = CompanyEnrichment(
            company_id=company_id,
            status=status,
            data=data,
            errors=errors,
            credits_spent=credits_spent,
            fetched_at=fetched_at,
        )
        self.db.add(enrichment)
        await self.db.flush()
        return enrichment