# vecta_sdk/vecta/connectors/file_store_connector.py
"""File store connector for ingesting documents from file systems."""

from typing import List, Optional
import os
import re
import uuid

from vecta.connectors.base import BaseFileStoreConnector
from vecta.core.schemas import ChunkData


class FileStoreConnector(BaseFileStoreConnector):
    """
    Connector for file stores that ingests documents using markitdown.

    This connector creates ChunkData objects from files on disk,
    splitting the converted Markdown output into page-level chunks.
    """

    def __init__(
        self,
        file_paths: List[str],
        base_path: str = "/mnt/file_stores",
    ):
        """
        Initialize file store connector.

        Args:
            file_paths: List of file paths to ingest (relative to base_path)
            base_path: Base directory where files are stored
        """
        super().__init__()
        self.file_paths = file_paths
        self.base_path = base_path
        self._chunks_cache: Optional[List[ChunkData]] = None

    @staticmethod
    def _split_pages(markdown_text: str) -> List[str]:
        """
        Split converted Markdown into page-level chunks.

        Uses form-feed characters (\\f) inserted by markitdown for PDF
        page breaks. If no form-feeds are found, falls back to splitting
        on top-level headings (# …). If neither delimiter is present the
        entire text is returned as a single chunk.
        """
        # markitdown inserts \f between PDF pages
        if "\f" in markdown_text:
            pages = [p.strip() for p in markdown_text.split("\f") if p.strip()]
            if pages:
                return pages

        # Fallback: split on top-level Markdown headings
        parts = re.split(r"(?m)^(?=# )", markdown_text)
        pages = [p.strip() for p in parts if p.strip()]
        return pages if pages else [markdown_text]

    def get_all_chunks(self) -> List[ChunkData]:
        """
        Retrieve all chunks by ingesting files with markitdown.

        Returns:
            List of ChunkData objects containing ingested content
        """
        if self._chunks_cache is not None:
            return self._chunks_cache

        try:
            from markitdown import MarkItDown
        except ImportError:
            raise ImportError(
                "markitdown is required for file store ingestion. "
                "Install with: pip install 'markitdown[pdf,docx,pptx,xlsx]'"
            )

        md = MarkItDown()
        all_chunks: List[ChunkData] = []

        for file_path in self.file_paths:
            full_path = os.path.join(self.base_path, file_path)

            if not os.path.exists(full_path):
                print(f"Warning: File not found: {full_path}")
                continue

            try:
                # Convert file to Markdown using markitdown
                result = md.convert(full_path)
                text = result.text_content or ""

                # Split into page-level chunks
                pages = self._split_pages(text)

                # Convert pages to ChunkData
                for idx, page_text in enumerate(pages):
                    # Generate unique chunk ID
                    chunk_id = f"{file_path}_{idx}_{uuid.uuid4().hex[:8]}"

                    # Extract document name from path
                    doc_name = os.path.basename(file_path)

                    # Create metadata
                    metadata = {
                        "source_path": doc_name,
                        # Chroma only accepts primitive metadata values. Using a plain
                        # integer ensures the chunk can be stored without errors while
                        # still allowing the SDK to normalize the value back into a
                        # list of page numbers when reading results.
                        "page_nums": idx,
                        "file_path": file_path,
                        "chunk_index": idx,
                    }

                    # Create ChunkData directly
                    chunk_data = ChunkData(
                        id=chunk_id,
                        content=page_text,
                        metadata=metadata,
                        source_path=doc_name,
                        page_nums=[idx],
                    )

                    all_chunks.append(chunk_data)

            except Exception as e:
                print(f"Error ingesting file {full_path}: {e}")
                continue

        self._chunks_cache = all_chunks
        return all_chunks

    def get_chunk_by_id(self, chunk_id: str) -> ChunkData:
        """
        Retrieve a specific chunk by its ID.

        Args:
            chunk_id: The unique identifier for the chunk

        Returns:
            ChunkData object for the specified chunk
        """
        all_chunks = self.get_all_chunks()

        for chunk in all_chunks:
            if chunk.id == chunk_id:
                return chunk

        raise ValueError(f"Chunk with ID '{chunk_id}' not found")

    def clear_cache(self) -> None:
        """Clear the cached chunks to force re-ingestion."""
        self._chunks_cache = None
