"""File-based data source implementations.

This module provides data source implementations for file-based data sources,
including JSON, JSONL, and CSV/TSV files. It integrates with the existing
chunker logic for efficient batch processing.
"""

from pathlib import Path
from typing import Iterator

from pydantic import field_validator

from graflo.architecture.onto import EncodingType
from graflo.data_source.base import AbstractDataSource, DataSourceType
from graflo.util.chunker import ChunkerFactory, ChunkerType


class FileDataSource(AbstractDataSource):
    """Base class for file-based data sources.

    This class provides a common interface for file-based data sources,
    integrating with the existing chunker system for batch processing.

    Attributes:
        path: Path to the file
        file_type: Type of file (json, jsonl, table)
        encoding: File encoding (default: UTF_8)
    """

    path: Path | str
    file_type: str | None = None
    encoding: EncodingType = EncodingType.UTF_8
    source_type: DataSourceType = DataSourceType.FILE

    @field_validator("path", mode="before")
    @classmethod
    def _path_to_path(cls, v: Path | str) -> Path:
        return Path(v) if isinstance(v, str) else v

    def iter_batches(
        self, batch_size: int = 1000, limit: int | None = None
    ) -> Iterator[list[dict]]:
        """Iterate over file data in batches.

        Args:
            batch_size: Number of items per batch
            limit: Maximum number of items to retrieve

        Yields:
            list[dict]: Batches of documents as dictionaries
        """
        # Determine chunker type
        chunker_type = None
        if self.file_type:
            chunker_type = ChunkerType(self.file_type.lower())

        # Create chunker using factory
        chunker_kwargs = {
            "resource": self.path,
            "type": chunker_type,
            "batch_size": batch_size,
            "limit": limit,
            "encoding": self.encoding,
        }
        # Only add sep for table files
        if chunker_type == ChunkerType.TABLE and hasattr(self, "sep"):
            chunker_kwargs["sep"] = self.sep

        chunker = ChunkerFactory.create_chunker(**chunker_kwargs)

        # Yield batches
        for batch in chunker:
            yield batch


class JsonFileDataSource(FileDataSource):
    """Data source for JSON files.

    JSON files are expected to contain hierarchical data structures,
    similar to REST API responses. The chunker handles nested structures
    and converts them to dictionaries.

    Attributes:
        path: Path to the JSON file
        encoding: File encoding (default: UTF_8)
    """

    file_type: str = ChunkerType.JSON.value


class JsonlFileDataSource(FileDataSource):
    """Data source for JSONL (JSON Lines) files.

    JSONL files contain one JSON object per line, making them suitable
    for streaming and batch processing.

    Attributes:
        path: Path to the JSONL file
        encoding: File encoding (default: UTF_8)
    """

    file_type: str = ChunkerType.JSONL.value


class TableFileDataSource(FileDataSource):
    """Data source for CSV/TSV files.

    Table files are converted to dictionaries with column headers as keys.
    Each row becomes a dictionary.

    Attributes:
        path: Path to the CSV/TSV file
        encoding: File encoding (default: UTF_8)
        sep: Field separator (default: ',')
    """

    sep: str = ","
    file_type: str = ChunkerType.TABLE.value


class ParquetFileDataSource(FileDataSource):
    """Data source for Parquet files.

    Parquet files are columnar storage format files that are read using pandas.
    Each row becomes a dictionary with column names as keys.

    Attributes:
        path: Path to the Parquet file
    """

    file_type: str = ChunkerType.PARQUET.value
