python/src/server/services/storage/document_storage.py

"""
Document Storage Facade

Provides the main entry point for document and code storage operations.
Consolidates responsibilities previously split between document_storage_service.py
and crawling/document_storage_operations.py.
"""

import asyncio
import os
from collections.abc import Callable
from typing import Any
from urllib.parse import urlparse

from src.server.config.logfire_config import safe_logfire_error, safe_logfire_info, safe_span, search_logger
from src.server.repositories.base_repository import BaseRepository
from src.server.services.credential_service import credential_service
from src.server.services.embeddings.contextual_embedding_service import generate_contextual_embeddings_batch
from src.server.services.embeddings.embedding_service import create_embeddings_batch
from src.server.services.source_management_service import extract_source_summary, update_source_info

from .chunking_utils import ChunkingUtils
from .progress_tracker import ProgressTracker
from .repositories.document_repo import DocumentRepository


class DocumentStorageFacade(BaseRepository):
    """Facade for all document storage and processing operations."""

    def __init__(self, supabase_client=None):
        super().__init__(supabase_client)
        self.repo = DocumentRepository(self.supabase_client)
        self.chunking_utils = ChunkingUtils(self.supabase_client)
        from src.server.services.crawling.code_extraction_service import CodeExtractionService

        self.code_extraction_service = CodeExtractionService(self.supabase_client)

    async def store_documents(
        self,
        crawl_results: list[dict],
        request: dict[str, Any],
        crawl_type: str,
        original_source_id: str,
        progress_callback: Callable | None = None,
        cancellation_check: Callable | None = None,
        source_url: str | None = None,
        source_display_name: str | None = None,
    ) -> dict[str, Any]:
        """
        Process crawled documents, chunk them, generate embeddings, and store them.
        """
        # 1. Chunking Phase (Pure Logic)
        (
            all_urls,
            all_chunk_numbers,
            all_contents,
            all_metadatas,
            source_word_counts,
            url_to_full_document,
            processed_docs,
        ) = await self.chunking_utils.prepare_document_chunks(
            crawl_results, request, crawl_type, original_source_id, cancellation_check
        )

        # 2. Source Record Phase
        if all_contents and all_metadatas:
            await self._create_source_records(
                all_metadatas, all_contents, source_word_counts, request, source_url, source_display_name
            )

        avg_chunks = (len(all_contents) / processed_docs) if processed_docs > 0 else 0.0
        safe_logfire_info(
            f"Document storage | processed={processed_docs}/{len(crawl_results)} | chunks={len(all_contents)} | avg_chunks={avg_chunks:.1f}"
        )

        # 3. Storage and Embeddings Phase
        storage_stats = await self._add_documents_to_supabase(
            urls=all_urls,
            chunk_numbers=all_chunk_numbers,
            contents=all_contents,
            metadatas=all_metadatas,
            url_to_full_document=url_to_full_document,
            progress_callback=progress_callback,
            cancellation_check=cancellation_check,
        )

        return {
            "chunk_count": len(all_contents),
            "chunks_stored": storage_stats.get("chunks_stored", 0),
            "total_word_count": sum(source_word_counts.values()),
            "url_to_full_document": url_to_full_document,
            "source_id": original_source_id,
        }

    async def store_code_examples(
        self,
        crawl_results: list[dict],
        url_to_full_document: dict[str, str],
        source_id: str,
        progress_callback: Callable | None = None,
        start_progress: int = 85,
        end_progress: int = 95,
        cancellation_check: Callable | None = None,
    ) -> int:
        """
        Delegates to CodeExtractionService to extract and store code snippets.
        """
        return await self.code_extraction_service.extract_and_store_code_examples(
            crawl_results,
            url_to_full_document,
            source_id,
            progress_callback,
            start_progress,
            end_progress,
            cancellation_check,
        )

    async def _add_documents_to_supabase(
        self,
        urls: list[str],
        chunk_numbers: list[int],
        contents: list[str],
        metadatas: list[dict[str, Any]],
        url_to_full_document: dict[str, str],
        batch_size: int | None = None,
        progress_callback: Callable | None = None,
        cancellation_check: Callable | None = None,
    ) -> dict[str, int]:
        """Internal method handling DB ingestion, deletion, and embeddings."""
        tracker = ProgressTracker(progress_callback)

        with safe_span("add_documents_to_supabase", total_documents=len(contents), batch_size=batch_size) as span:
            try:
                rag_settings = await credential_service.get_credentials_by_category("rag_strategy")
                if batch_size is None:
                    batch_size = int(rag_settings.get("DOCUMENT_STORAGE_BATCH_SIZE", "50"))
                delete_batch_size = int(rag_settings.get("DELETE_BATCH_SIZE", "50"))
            except Exception as e:
                search_logger.warning(f"Failed to load storage settings: {e}, using defaults")
                batch_size = batch_size or 50
                delete_batch_size = 50

            # Delete existing records using Repository
            await self.repo.delete_existing_urls_in_batches(urls, delete_batch_size, cancellation_check)

            # Check contextual embeddings
            try:
                use_contextual_embeddings = await credential_service.get_credential(
                    "USE_CONTEXTUAL_EMBEDDINGS", "false", decrypt=True
                )
                if isinstance(use_contextual_embeddings, str):
                    use_contextual_embeddings = use_contextual_embeddings.lower() == "true"
            except Exception:
                use_contextual_embeddings = os.getenv("USE_CONTEXTUAL_EMBEDDINGS", "false") == "true"

            completed_batches = 0
            total_batches = (len(contents) + batch_size - 1) // batch_size
            total_chunks_stored = 0

            for batch_num, i in enumerate(range(0, len(contents), batch_size), 1):
                if cancellation_check:
                    cancellation_check()

                batch_end = min(i + batch_size, len(contents))
                batch_urls = urls[i:batch_end]
                batch_chunk_numbers = chunk_numbers[i:batch_end]
                batch_contents = contents[i:batch_end]
                batch_metadatas = metadatas[i:batch_end]

                current_progress = int((completed_batches / total_batches) * 100)

                # Contextual Embeddings Phase
                contextual_contents = batch_contents
                max_workers = 1
                if use_contextual_embeddings:
                    max_workers = int(os.getenv("CONTEXTUAL_EMBEDDINGS_MAX_WORKERS", "4"))
                    full_documents = [url_to_full_document.get(u, "") for u in batch_urls]
                    try:
                        sub_results = await generate_contextual_embeddings_batch(full_documents, batch_contents)
                        contextual_contents = []
                        for idx, (contextual_text, success) in enumerate(sub_results):
                            contextual_contents.append(contextual_text)
                            if success:
                                batch_metadatas[idx]["contextual_embedding"] = True
                    except Exception as e:
                        search_logger.error(f"Contextual embedding error: {e}")

                # Standard Embedding Phase
                async def progress_wrap(msg, pct, cp=current_progress, bn=batch_num):
                    await tracker.embedding_progress_wrapper(msg, pct, cp, bn)

                result = await create_embeddings_batch(
                    contextual_contents, progress_callback=progress_wrap if progress_callback else None
                )

                if not result.embeddings:
                    raise Exception(f"Skipping batch {batch_num} - no successful embeddings")

                batch_data = []
                for _j, (embedding, text) in enumerate(zip(result.embeddings, result.texts_processed, strict=False)):
                    try:
                        orig_idx = contextual_contents.index(text)
                    except ValueError:
                        continue

                    source_id = batch_metadatas[orig_idx].get("source_id")
                    if not source_id:
                        parsed = urlparse(batch_urls[orig_idx])
                        source_id = parsed.netloc or parsed.path

                    batch_data.append(
                        {
                            "url": batch_urls[orig_idx],
                            "chunk_number": batch_chunk_numbers[orig_idx],
                            "content": text,
                            "metadata": {"chunk_size": len(text), **batch_metadatas[orig_idx]},
                            "source_id": source_id,
                            "embedding": embedding,
                        }
                    )

                # DB Insertion Phase using Repository
                max_retries = 3
                retry_delay = 1.0
                for retry in range(max_retries):
                    if cancellation_check:
                        cancellation_check()
                    try:
                        self.repo.insert_document_batch(batch_data)
                        total_chunks_stored += len(batch_data)
                        completed_batches += 1
                        new_progress = (
                            100
                            if completed_batches == total_batches
                            else int((completed_batches / total_batches) * 100)
                        )

                        await tracker.report_progress(
                            f"Completed batch {batch_num}/{total_batches} ({len(batch_data)} chunks)",
                            new_progress,
                            {
                                "current_batch": batch_num,
                                "total_batches": total_batches,
                                "completed_batches": completed_batches,
                                "active_workers": max_workers,
                            },
                        )
                        break
                    except Exception as e:
                        if retry < max_retries - 1:
                            await asyncio.sleep(retry_delay)
                            retry_delay *= 2
                        else:
                            search_logger.error(f"Failed to insert batch: {e}")

                if i + batch_size < len(contents):
                    await asyncio.sleep(0.1)

            await tracker.report_final_progress(len(contents), total_batches)

            span.set_attribute("success", True)
            span.set_attribute("total_processed", len(contents))
            span.set_attribute("total_stored", total_chunks_stored)
            return {"chunks_stored": total_chunks_stored}

    async def _create_source_records(
        self,
        all_metadatas: list[dict],
        all_contents: list[str],
        source_word_counts: dict[str, int],
        request: dict[str, Any],
        source_url: str | None = None,
        source_display_name: str | None = None,
    ):
        """Internal helper to create source records via Repositories and external services."""
        unique_source_ids = set()
        source_id_contents: dict[str, list[str]] = {}

        for i, metadata in enumerate(all_metadatas):
            source_id = metadata["source_id"]
            unique_source_ids.add(source_id)
            if source_id not in source_id_contents:
                source_id_contents[source_id] = []
            source_id_contents[source_id].append(all_contents[i])

        for source_id in unique_source_ids:
            source_contents = source_id_contents[source_id]
            combined_content = ""
            for chunk in source_contents[:3]:
                if len(combined_content) + len(chunk) < 15000:
                    combined_content += " " + chunk
                else:
                    break

            try:
                summary = await extract_source_summary(source_id, combined_content)
            except Exception as e:
                search_logger.error(f"Failed to generate AI summary for '{source_id}'", exc_info=True)
                safe_logfire_error(f"Failed to generate AI summary for '{source_id}': {str(e)}, using fallback")
                summary = f"Documentation from {source_id} - {len(source_contents)} pages crawled"

            try:
                await update_source_info(
                    client=self.supabase_client,
                    source_id=source_id,
                    summary=summary,
                    word_count=source_word_counts[source_id],
                    content=combined_content,
                    knowledge_type=request.get("knowledge_type", "documentation"),
                    tags=request.get("tags", []),
                    update_frequency=0,
                    original_url=request.get("url"),
                    source_url=source_url,
                    source_display_name=source_display_name,
                )
                safe_logfire_info(f"Successfully created/updated source record for '{source_id}'")
            except Exception as e:
                search_logger.error(f"Failed to create/update source record for '{source_id}'", exc_info=True)
                safe_logfire_error(f"Failed to create/update source record for '{source_id}': {str(e)}")
                safe_logfire_info(f"Attempting fallback source creation for '{source_id}'")
                fallback_data = {
                    "source_id": source_id,
                    "title": source_id,
                    "summary": summary,
                    "total_word_count": source_word_counts[source_id],
                    "metadata": {
                        "knowledge_type": request.get("knowledge_type", "documentation"),
                        "tags": request.get("tags", []),
                        "auto_generated": True,
                        "fallback_creation": True,
                        "original_url": request.get("url"),
                    },
                }
                if source_url:
                    fallback_data["source_url"] = source_url
                if source_display_name:
                    fallback_data["source_display_name"] = source_display_name

                success, _ = self.repo.upsert_source_fallback(source_id, fallback_data)
                if not success:
                    raise Exception(f"Failed fallback creation for {source_id}") from None

        for source_id in unique_source_ids:
            success, _ = self.repo.verify_source_exists(source_id)
            if not success:
                raise Exception(f"Source verification failed for {source_id}")