potato/data_sources/partial_reader.py

"""
Partial/incremental loading for data sources.

This module provides functionality for loading data in chunks, enabling:
- Initial loading of first K items
- Batch loading of additional items as annotation progresses
- Auto-loading when annotation reaches a threshold
- State persistence for resumption after restart
"""

import json
import logging
import os
import threading
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, List, Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from potato.data_sources.base import DataSource

logger = logging.getLogger(__name__)


@dataclass
class PartialReadState:
    """
    Tracks the read state for a data source.

    This dataclass maintains information about how much data has been
    read from a source, enabling incremental loading.

    Attributes:
        source_id: Identifier of the data source
        items_loaded: Number of items loaded so far
        total_estimate: Estimated total items (None if unknown)
        file_position: Byte position for file-based sources
        line_number: Line number for line-based sources
        is_complete: Whether all data has been loaded
        last_loaded_at: Unix timestamp of last load
        metadata: Additional source-specific state
    """
    source_id: str
    items_loaded: int = 0
    total_estimate: Optional[int] = None
    file_position: int = 0
    line_number: int = 0
    is_complete: bool = False
    last_loaded_at: Optional[float] = None
    metadata: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for JSON serialization."""
        return {
            "source_id": self.source_id,
            "items_loaded": self.items_loaded,
            "total_estimate": self.total_estimate,
            "file_position": self.file_position,
            "line_number": self.line_number,
            "is_complete": self.is_complete,
            "last_loaded_at": self.last_loaded_at,
            "metadata": self.metadata,
        }

    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> "PartialReadState":
        """Create from dictionary."""
        return cls(
            source_id=data["source_id"],
            items_loaded=data.get("items_loaded", 0),
            total_estimate=data.get("total_estimate"),
            file_position=data.get("file_position", 0),
            line_number=data.get("line_number", 0),
            is_complete=data.get("is_complete", False),
            last_loaded_at=data.get("last_loaded_at"),
            metadata=data.get("metadata", {}),
        )


@dataclass
class PartialLoadingConfig:
    """
    Configuration for partial/incremental loading.

    Attributes:
        enabled: Whether partial loading is enabled
        initial_count: Number of items to load initially
        batch_size: Number of items per incremental load
        auto_load_threshold: Auto-load when this fraction is annotated (0.0-1.0)
        auto_load_enabled: Whether auto-loading is enabled
    """
    enabled: bool = False
    initial_count: int = 1000
    batch_size: int = 500
    auto_load_threshold: float = 0.8
    auto_load_enabled: bool = True

    @classmethod
    def from_dict(cls, config: Dict[str, Any]) -> "PartialLoadingConfig":
        """Create from configuration dictionary."""
        partial_config = config.get("partial_loading", {})
        return cls(
            enabled=partial_config.get("enabled", False),
            initial_count=partial_config.get("initial_count", 1000),
            batch_size=partial_config.get("batch_size", 500),
            auto_load_threshold=partial_config.get("auto_load_threshold", 0.8),
            auto_load_enabled=partial_config.get("auto_load_enabled", True),
        )

    def validate(self) -> List[str]:
        """Validate the configuration."""
        errors = []

        if self.initial_count < 1:
            errors.append("partial_loading.initial_count must be at least 1")

        if self.batch_size < 1:
            errors.append("partial_loading.batch_size must be at least 1")

        if not 0.0 <= self.auto_load_threshold <= 1.0:
            errors.append(
                "partial_loading.auto_load_threshold must be between 0.0 and 1.0"
            )

        return errors


class PartialReader:
    """
    Manages partial/incremental loading of data sources.

    This class coordinates loading data in chunks across multiple sources,
    tracking state for each source and providing auto-loading when
    annotation progress reaches a threshold.

    Thread Safety:
        All public methods are thread-safe. Internal state is protected
        by a lock.

    Attributes:
        config: Partial loading configuration
        state_file: Path to state persistence file
    """

    STATE_FILENAME = ".data_source_state.json"

    def __init__(
        self,
        config: PartialLoadingConfig,
        output_dir: str
    ):
        """
        Initialize the partial reader.

        Args:
            config: Partial loading configuration
            output_dir: Directory for state persistence
        """
        self.config = config
        self.output_dir = Path(output_dir)
        self.state_file = self.output_dir / self.STATE_FILENAME

        self._states: Dict[str, PartialReadState] = {}
        self._lock = threading.RLock()

        # Load existing state
        self._load_state()

        logger.info(
            f"PartialReader initialized: initial={config.initial_count}, "
            f"batch={config.batch_size}, auto_threshold={config.auto_load_threshold}"
        )

    def _load_state(self) -> None:
        """Load persisted state from disk."""
        if not self.state_file.exists():
            return

        try:
            with open(self.state_file, 'r', encoding='utf-8') as f:
                data = json.load(f)

            for state_data in data.get("sources", []):
                state = PartialReadState.from_dict(state_data)
                self._states[state.source_id] = state

            logger.debug(f"Loaded partial read state for {len(self._states)} sources")

        except (json.JSONDecodeError, KeyError) as e:
            logger.warning(f"Failed to load partial read state: {e}")
            self._states = {}

    def _save_state(self) -> None:
        """Save state to disk."""
        # Ensure output directory exists
        self.output_dir.mkdir(parents=True, exist_ok=True)

        data = {
            "version": 1,
            "sources": [state.to_dict() for state in self._states.values()]
        }

        try:
            with open(self.state_file, 'w', encoding='utf-8') as f:
                json.dump(data, f, indent=2)
        except IOError as e:
            logger.error(f"Failed to save partial read state: {e}")

    def get_state(self, source_id: str) -> Optional[PartialReadState]:
        """
        Get the current state for a source.

        Args:
            source_id: The source identifier

        Returns:
            PartialReadState or None if not tracked
        """
        with self._lock:
            return self._states.get(source_id)

    def get_or_create_state(self, source_id: str) -> PartialReadState:
        """
        Get or create state for a source.

        Args:
            source_id: The source identifier

        Returns:
            PartialReadState for the source
        """
        with self._lock:
            if source_id not in self._states:
                self._states[source_id] = PartialReadState(source_id=source_id)
                self._save_state()
            return self._states[source_id]

    def update_state(
        self,
        source_id: str,
        items_added: int,
        file_position: Optional[int] = None,
        line_number: Optional[int] = None,
        is_complete: bool = False,
        total_estimate: Optional[int] = None
    ) -> PartialReadState:
        """
        Update the state after loading items.

        Args:
            source_id: The source identifier
            items_added: Number of items added in this batch
            file_position: New file position (for file sources)
            line_number: New line number (for line-based sources)
            is_complete: Whether all data has been loaded
            total_estimate: Updated total estimate

        Returns:
            Updated PartialReadState
        """
        import time

        with self._lock:
            state = self.get_or_create_state(source_id)
            state.items_loaded += items_added
            state.last_loaded_at = time.time()

            if file_position is not None:
                state.file_position = file_position
            if line_number is not None:
                state.line_number = line_number
            if total_estimate is not None:
                state.total_estimate = total_estimate
            if is_complete:
                state.is_complete = True

            self._save_state()
            return state

    def mark_complete(self, source_id: str) -> None:
        """
        Mark a source as completely loaded.

        Args:
            source_id: The source identifier
        """
        with self._lock:
            state = self.get_or_create_state(source_id)
            state.is_complete = True
            self._save_state()

    def should_load_more(
        self,
        source_id: str,
        annotated_count: int,
        total_loaded: int
    ) -> bool:
        """
        Check if more data should be loaded based on annotation progress.

        This method implements the auto-load threshold logic.

        Args:
            source_id: The source identifier
            annotated_count: Number of items annotated
            total_loaded: Total number of items loaded

        Returns:
            True if more data should be loaded
        """
        if not self.config.auto_load_enabled:
            return False

        with self._lock:
            state = self._states.get(source_id)
            if state and state.is_complete:
                return False

            if total_loaded == 0:
                return False

            progress = annotated_count / total_loaded
            return progress >= self.config.auto_load_threshold

    def get_load_count(
        self,
        source_id: str,
        is_initial: bool = False
    ) -> int:
        """
        Get the number of items to load.

        Args:
            source_id: The source identifier
            is_initial: Whether this is the initial load

        Returns:
            Number of items to load
        """
        if is_initial:
            return self.config.initial_count
        return self.config.batch_size

    def get_start_position(self, source_id: str) -> int:
        """
        Get the starting position for the next load.

        Args:
            source_id: The source identifier

        Returns:
            Number of items already loaded (start position for next batch)
        """
        with self._lock:
            state = self._states.get(source_id)
            if state:
                return state.items_loaded
            return 0

    def reset_state(self, source_id: str) -> None:
        """
        Reset the state for a source.

        Args:
            source_id: The source identifier
        """
        with self._lock:
            if source_id in self._states:
                del self._states[source_id]
                self._save_state()

    def clear_all_state(self) -> None:
        """Clear state for all sources."""
        with self._lock:
            self._states.clear()
            self._save_state()

    def get_all_states(self) -> Dict[str, PartialReadState]:
        """
        Get all source states.

        Returns:
            Dictionary mapping source_id to PartialReadState
        """
        with self._lock:
            return dict(self._states)

    def get_stats(self) -> Dict[str, Any]:
        """
        Get statistics about partial loading.

        Returns:
            Dictionary with loading statistics
        """
        with self._lock:
            total_loaded = sum(s.items_loaded for s in self._states.values())
            complete_count = sum(1 for s in self._states.values() if s.is_complete)

            return {
                "enabled": self.config.enabled,
                "initial_count": self.config.initial_count,
                "batch_size": self.config.batch_size,
                "auto_load_threshold": self.config.auto_load_threshold,
                "sources_tracked": len(self._states),
                "sources_complete": complete_count,
                "total_items_loaded": total_loaded,
            }