"""
Document ingestion utilities for RapidRAG.
"""

from pathlib import Path
from typing import Optional, List, Dict, Any, Generator
import hashlib


class DocumentIngester:
    """
    Ingest documents from various sources.

    Example:
        ingester = DocumentIngester(rag)
        ingester.from_directory("./docs/")
        ingester.from_url("https://example.com/doc.txt")
    """

    def __init__(self, rag: "RapidRAG"):
        """
        Initialize ingester.

        Args:
            rag: RapidRAG instance to add documents to
        """
        self.rag = rag
        self._stats = {"files": 0, "chunks": 0, "errors": 0}

    def from_directory(
        self,
        path: str,
        extensions: Optional[List[str]] = None,
        recursive: bool = True,
        chunk_size: int = 1000,
        chunk_overlap: int = 200
    ) -> Dict[str, int]:
        """
        Ingest all documents from a directory.

        Args:
            path: Directory path
            extensions: File extensions to process
            recursive: Include subdirectories
            chunk_size: Characters per chunk
            chunk_overlap: Overlap between chunks

        Returns:
            Stats dict with files, chunks, errors
        """
        self._stats = {"files": 0, "chunks": 0, "errors": 0}

        ids = self.rag.add_directory(
            path,
            extensions=extensions,
            recursive=recursive,
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap
        )

        self._stats["chunks"] = len(ids)
        return self._stats

    def from_texts(
        self,
        texts: List[str],
        ids: Optional[List[str]] = None,
        metadatas: Optional[List[Dict]] = None
    ) -> int:
        """
        Ingest a list of texts.

        Args:
            texts: List of text strings
            ids: Optional document IDs
            metadatas: Optional metadata dicts

        Returns:
            Number of documents added
        """
        added_ids = self.rag.add_texts(texts, ids, metadatas)
        return len(added_ids)

    def from_jsonl(
        self,
        path: str,
        content_field: str = "text",
        id_field: Optional[str] = None,
        metadata_fields: Optional[List[str]] = None
    ) -> int:
        """
        Ingest documents from a JSONL file.

        Args:
            path: Path to JSONL file
            content_field: Field containing text content
            id_field: Field containing document ID
            metadata_fields: Fields to include as metadata

        Returns:
            Number of documents added
        """
        import json

        texts = []
        ids = []
        metadatas = []

        with open(path, "r", encoding="utf-8") as f:
            for line in f:
                doc = json.loads(line)
                content = doc.get(content_field, "")
                if not content:
                    continue

                texts.append(content)

                if id_field and id_field in doc:
                    ids.append(str(doc[id_field]))
                else:
                    ids.append(hashlib.sha256(content.encode()).hexdigest()[:16])

                if metadata_fields:
                    meta = {k: doc.get(k) for k in metadata_fields if k in doc}
                else:
                    meta = {k: v for k, v in doc.items() if k != content_field}
                metadatas.append(meta)

        return self.from_texts(texts, ids, metadatas)

    @property
    def stats(self) -> Dict[str, int]:
        """Get ingestion stats."""
        return self._stats
