from datetime import UTC, datetime

from home_agent.contracts.protocols import EmbeddingProvider, ItemTextRenderer, VectorBackend
from home_agent.contracts.types import (
    BackfillResult,
    EmbeddingRecord,
    Item,
    ItemKind,
    VectorDocument,
    VectorDocumentKey,
)
from home_agent.storage.sqlite_store import SqliteStore


def run_embeddings_backfill(
    store: SqliteStore,
    embedding_provider: EmbeddingProvider,
    item_text_renderer: ItemTextRenderer,
    vector_backend: VectorBackend | None = None,
    *,
    batch_size: int,
    limit: int | None,
    since: datetime | None,
    kind: ItemKind | None,
    dry_run: bool,
    rebuild: bool,
) -> BackfillResult:
    items = store.list_items_for_embedding(
        embedding_provider=embedding_provider.provider_name,
        embedding_model=embedding_provider.model_name,
        limit=limit,
        since=since,
        kind=kind,
        rebuild=rebuild or vector_backend is not None,
    )
    skipped = 0
    if vector_backend is not None and not rebuild:
        filtered_items: list[Item] = []
        for item in items:
            key = VectorDocumentKey(
                provider=item.provider,
                external_id=item.external_id,
                embedding_provider=embedding_provider.provider_name,
                embedding_model=embedding_provider.model_name,
            )
            if vector_backend.has_document(key):
                skipped += 1
                continue
            filtered_items.append(item)
        items = filtered_items
    if dry_run:
        return BackfillResult(processed=0, skipped=len(items) + skipped, failed=0)

    processed = 0
    failed = 0
    for offset in range(0, len(items), batch_size):
        batch = items[offset : offset + batch_size]
        rendered_batch = [item_text_renderer.render(item) for item in batch]
        try:
            vectors = embedding_provider.embed_texts(
                [rendered.embedding_text for rendered in rendered_batch]
            )
        except Exception:
            failed += len(batch)
            continue

        if len(vectors) != len(batch):
            failed += len(batch)
            continue

        backend_documents: list[VectorDocument] = []
        for item, rendered, vector in zip(batch, rendered_batch, vectors, strict=True):
            store.save_item_text(rendered)
            store.save_item_embedding(
                EmbeddingRecord(
                    provider=item.provider,
                    external_id=item.external_id,
                    embedding_provider=embedding_provider.provider_name,
                    embedding_model=embedding_provider.model_name,
                    vector=vector,
                    dimension=len(vector),
                    created_at=datetime.now(tz=UTC),
                )
            )
            backend_documents.append(
                VectorDocument(
                    key=VectorDocumentKey(
                        provider=item.provider,
                        external_id=item.external_id,
                        embedding_provider=embedding_provider.provider_name,
                        embedding_model=embedding_provider.model_name,
                    ),
                    kind=item.kind,
                    received_at=item.received_at,
                    urgency_label=None,
                    summary_text=rendered.summary_text,
                    source_version=rendered.source_version,
                    vector=vector,
                    metadata={},
                )
            )
            processed += 1
        if vector_backend is not None and backend_documents:
            vector_backend.upsert_documents(backend_documents)
    return BackfillResult(processed=processed, skipped=skipped, failed=failed)
