# input_readers.py
"""Flexible input format readers for STAC URL ingestion with cloud storage support.

This module provides a pluggable input reading system that supports multiple file formats
and storage backends for STAC URL ingestion. Designed to handle massive datasets efficiently
while providing a consistent interface across different input sources and formats.

Supported Input Formats:
    - Parquet: High-performance columnar format (recommended for large datasets)
    - CSV: Comma-separated values with configurable options
    - TSV: Tab-separated values (specialized CSV configuration)
    - NDJSON/JSONL: Newline-delimited JSON for streaming STAC items
    - Auto-detection: Automatic format detection based on file extension

Supported Storage Systems:
    - Local filesystem: Direct file access for development and single-machine deployment
    - S3 and S3-compatible: AWS S3, MinIO, DigitalOcean Spaces via fsspec
    - Google Cloud Storage: GCS integration via fsspec protocols
    - Azure Blob Storage: Azure integration via fsspec protocols
    - HTTP/HTTPS: Direct URL reading for publicly accessible datasets

Key Features:
    - Automatic format detection eliminates manual configuration
    - Cloud storage integration with optimized streaming
    - Memory-efficient processing for unlimited file sizes
    - Comprehensive error handling with informative messages
    - Extensible design for adding new input formats
    - Thread-safe operations for concurrent processing

Performance Optimizations:
    - Parquet: Columnar reading with predicate pushdown for URL columns
    - CSV: Streaming parsing with configurable chunk sizes
    - Cloud storage: Intelligent buffering and connection reuse
    - Memory efficiency: Constant memory usage regardless of input file size

Design Pattern:
    The input readers follow the Strategy pattern with a Factory for automatic
    reader selection. This enables seamless switching between formats and
    automatic format detection based on file extensions.

Example:
    >>> # Automatic format detection and reader selection
    >>> reader = ReaderFactory.get_reader('auto')
    >>> urls = reader.read_urls('s3://bucket/urls.parquet', 'url')
    >>>
    >>> # Explicit format specification for better performance
    >>> reader = ReaderFactory.get_reader('parquet')
    >>> urls = reader.read_urls('large_dataset.parquet', 'stac_url')
    >>>
    >>> # CSV with custom configuration
    >>> reader = ReaderFactory.get_reader('csv', delimiter='|', encoding='utf-8')
    >>> urls = reader.read_urls('custom_format.csv', 'item_url')

Integration:
    Input readers integrate seamlessly with EarthCatalog's ingestion pipeline
    through the ProcessingConfig.input_format parameter. The pipeline automatically
    selects and configures the appropriate reader based on configuration.
"""

import logging
from abc import ABC, abstractmethod
from typing import cast

logger = logging.getLogger(__name__)

# Check for optional dependencies
try:
    import pandas as pd

    HAS_PANDAS = True
except ImportError:
    pd = None
    HAS_PANDAS = False

try:
    import pyarrow.parquet as pq

    HAS_PYARROW = True
except ImportError:
    pq = None
    HAS_PYARROW = False

try:
    import fsspec

    HAS_FSSPEC = True
except ImportError:
    fsspec = None
    HAS_FSSPEC = False


class InputReader(ABC):
    """Abstract base class defining the interface for all input format readers.

    This abstract base class establishes the contract for reading STAC URLs from various
    input formats and storage systems. All concrete readers must implement the defined
    interface to ensure consistent behavior across different input sources.

    Interface Design:
        The interface is designed for high-performance streaming operations that can
        handle unlimited file sizes with constant memory usage. Implementations should
        focus on efficient URL extraction while validating input format compliance.

    Implementation Requirements:
        - read_urls(): Must support streaming for large files to maintain memory efficiency
        - validate_format(): Should perform quick format validation without full file reads
        - Error handling: Must provide informative error messages for debugging
        - Cloud storage: Should integrate with fsspec for seamless cloud operations

    Performance Expectations:
        - Memory usage should remain constant regardless of input file size
        - URL extraction should be optimized for the specific format characteristics
        - Validation should be fast and avoid unnecessary I/O operations
        - Cloud operations should use streaming and connection reuse
    """

    @abstractmethod
    def read_urls(self, file_path: str, url_column: str = "url") -> list[str]:
        """Read URLs from input file and return as list."""
        pass

    @abstractmethod
    def validate_format(self, file_path: str) -> bool:
        """Validate that file matches expected format."""
        pass


class ParquetReader(InputReader):
    """High-performance Parquet file reader optimized for large-scale URL extraction.

    This reader provides optimized access to Parquet files containing STAC URLs,
    leveraging Parquet's columnar format for efficient URL column extraction.
    Designed for processing massive datasets with minimal memory overhead.

    Parquet Advantages:
        - Columnar storage enables reading only required URL columns
        - Excellent compression reduces network I/O for cloud storage
        - Built-in schema validation ensures data consistency
        - Fast metadata access for quick file validation
        - Predicate pushdown capabilities for filtered reading

    Performance Features:
        - Memory-efficient streaming for files of any size
        - Optimized cloud storage integration with fsspec
        - Parallel reading capabilities where supported
        - Intelligent buffering for network-based files
        - Schema-based validation for robust error handling

    Cloud Storage Optimization:
        - Uses fsspec for unified cloud storage access
        - Streaming reads minimize memory usage for large files
        - Connection pooling and reuse for better performance
        - Automatic retry strategies for network resilience

    Use Cases:
        - Large-scale STAC URL datasets (recommended format)
        - Cloud-based data processing workflows
        - High-performance ingestion requiring minimal memory
        - Datasets requiring schema validation and consistency
        - Integration with existing Parquet-based data pipelines

    Example:
        >>> reader = ParquetReader()
        >>> urls = reader.read_urls('s3://bucket/urls.parquet', 'stac_url')
        >>> print(f"Extracted {len(urls)} URLs from Parquet file")
    """

    def read_urls(self, file_path: str, url_column: str = "url") -> list[str]:
        """Read URLs from Parquet file."""
        if not HAS_PYARROW:
            raise ImportError("pyarrow is required for Parquet support. Install with: pip install pyarrow")

        try:
            if file_path.startswith("s3://"):
                if not HAS_FSSPEC:
                    raise ImportError("fsspec is required for S3 support. Install with: pip install fsspec s3fs")
                if fsspec is None:
                    raise ImportError("fsspec is required for S3 support. Install with: pip install fsspec s3fs")
                fs = fsspec.filesystem("s3")
                with fs.open(file_path, "rb") as f:
                    if pq is None:
                        raise ImportError("pyarrow is required for Parquet support. Install with: pip install pyarrow")
                    table = pq.read_table(f)
            else:
                if pq is None:
                    raise ImportError("pyarrow is required for Parquet support. Install with: pip install pyarrow")
                table = pq.read_table(file_path)

            df = table.to_pandas()

            if url_column not in df.columns:
                raise ValueError(f"Parquet file must contain '{url_column}' column")

            urls = cast(list[str], df[url_column].tolist())
            logger.info(f"Read {len(urls)} URLs from Parquet file: {file_path}")
            return urls

        except Exception as e:
            logger.error(f"Error reading Parquet file {file_path}: {e}")
            raise

    def validate_format(self, file_path: str) -> bool:
        """Validate Parquet file format."""
        if not HAS_PYARROW or pq is None:
            return False

        try:
            if file_path.startswith("s3://"):
                if not HAS_FSSPEC or fsspec is None:
                    return False
                fs = fsspec.filesystem("s3")
                with fs.open(file_path, "rb") as f:
                    pq.ParquetFile(f)
            else:
                pq.ParquetFile(file_path)
            return True
        except (OSError, ValueError, TypeError, RuntimeError):
            return False


class CSVReader(InputReader):
    """Reader for CSV files with URL column."""

    def __init__(self, delimiter: str = ",", quotechar: str = '"'):
        self.delimiter = delimiter
        self.quotechar = quotechar

    def read_urls(self, file_path: str, url_column: str = "url") -> list[str]:
        """Read URLs from CSV file."""
        if not HAS_PANDAS:
            raise ImportError("pandas is required for CSV support. Install with: pip install pandas")

        try:
            if file_path.startswith("s3://"):
                if not HAS_FSSPEC or fsspec is None:
                    raise ImportError("fsspec is required for S3 support. Install with: pip install fsspec s3fs")
                fs = fsspec.filesystem("s3")
                with fs.open(file_path, "r") as f:
                    if pd is None:
                        raise ImportError("pandas is required for CSV support. Install with: pip install pandas")
                    df = pd.read_csv(f, delimiter=self.delimiter, quotechar=self.quotechar)
            else:
                if pd is None:
                    raise ImportError("pandas is required for CSV support. Install with: pip install pandas")
                df = pd.read_csv(file_path, delimiter=self.delimiter, quotechar=self.quotechar)

            if url_column not in df.columns:
                raise ValueError(f"CSV file must contain '{url_column}' column")

            urls = cast(list[str], df[url_column].dropna().tolist())
            logger.info(f"Read {len(urls)} URLs from CSV file: {file_path}")
            return urls

        except Exception as e:
            logger.error(f"Error reading CSV file {file_path}: {e}")
            raise

    def validate_format(self, file_path: str) -> bool:
        """Validate CSV file format."""
        if not HAS_PANDAS or pd is None:
            return False

        try:
            if file_path.startswith("s3://"):
                if not HAS_FSSPEC or fsspec is None:
                    return False
                fs = fsspec.filesystem("s3")
                with fs.open(file_path, "r") as f:
                    # Try to read first few lines to validate
                    pd.read_csv(f, nrows=5, delimiter=self.delimiter, quotechar=self.quotechar)
            else:
                pd.read_csv(file_path, nrows=5, delimiter=self.delimiter, quotechar=self.quotechar)
            return True
        except (OSError, ValueError, TypeError):
            return False


class NDJSONReader(InputReader):
    """Reader for NDJSON/JSONL files with URL field.

    This reader processes Newline-Delimited JSON (NDJSON) files, where each line
    contains a separate JSON object with a URL field. NDJSON is ideal for streaming
    large datasets and is commonly used for STAC item aggregations.

    NDJSON Format:
        - Each line is a valid JSON object
        - Lines are separated by newlines (not comma-separated)
        - Lines starting with '#' are treated as comments
        - URL field can be any key name (default: "url")

    Performance Features:
        - Memory-efficient streaming for unlimited file sizes
        - Line-by-line processing prevents loading entire file into memory
        - Cloud storage integration via fsspec for S3, GCS, Azure
        - Skips empty lines and comments automatically
        - Graceful handling of malformed JSON lines

    Use Cases:
        - Large-scale STAC URL datasets where each item is a JSON object
        - Streaming ingestion from cloud storage
        - ITS_LIVE bulk data processing (800+ files, ~50k items each)
        - Log files with JSON entries
        - Exported STAC collections in NDJSON format

    Example NDJSON file:
        {"url": "https://example.com/item1.json", "id": "item1"}
        {"url": "https://example.com/item2.json", "id": "item2"}
        # This is a comment line (will be skipped)
        {"url": "https://example.com/item3.json", "id": "item3"}

    Example:
        >>> reader = NDJSONReader()
        >>> urls = reader.read_urls('s3://bucket/items.jsonl', 'url')
        >>> print(f"Extracted {len(urls)} URLs from NDJSON file")
    """

    def read_urls(self, file_path: str, url_column: str = "url") -> list[str]:
        """Read URLs from NDJSON file.

        Args:
            file_path: Path to NDJSON file (local or s3://)
            url_column: JSON key name containing URLs (default: "url")

        Returns:
            List of URLs extracted from each JSON object.

        Raises:
            ValueError: If file_path is empty or url_column is not found in any objects.
            ImportError: If fsspec is required for S3 but not installed.
        """
        if not file_path:
            raise ValueError("file_path cannot be empty")

        urls: list[str] = []
        line_number = 0
        skipped_lines = 0
        missing_field_count = 0

        try:
            if file_path.startswith("s3://"):
                if not HAS_FSSPEC or fsspec is None:
                    raise ImportError("fsspec is required for S3 support. Install with: pip install fsspec s3fs")
                fs = fsspec.filesystem("s3")
                with fs.open(file_path, "r") as f:
                    urls, line_number, skipped_lines, missing_field_count = self._read_lines(f, url_column)
            else:
                with open(file_path) as f:
                    urls, line_number, skipped_lines, missing_field_count = self._read_lines(f, url_column)

            if not urls:
                if line_number == 0:
                    # Empty file is valid, just return empty list
                    logger.info(f"NDJSON file is empty: {file_path}")
                elif missing_field_count > 0:
                    raise ValueError(
                        f"NDJSON file does not contain '{url_column}' field in any objects. "
                        f"Checked {line_number} lines."
                    )
                else:
                    raise ValueError(
                        f"NDJSON file does not contain any valid JSON objects with data. "
                        f"Checked {line_number} lines, {skipped_lines} were empty/comments."
                    )

            if skipped_lines > 0:
                logger.warning(f"Skipped {skipped_lines} empty/comment lines in {file_path}")

            if missing_field_count > 0:
                logger.warning(f"{missing_field_count} lines missing '{url_column}' field in {file_path}")

            logger.info(
                f"Read {len(urls)} URLs from NDJSON file: {file_path} "
                f"({line_number} lines processed, {skipped_lines} skipped)"
            )
            return urls

        except FileNotFoundError:
            logger.error(f"NDJSON file not found: {file_path}")
            raise
        except ImportError:
            raise
        except Exception as e:
            logger.error(f"Error reading NDJSON file {file_path}: {e}")
            raise

    def _read_lines(self, file_obj, url_column: str) -> tuple[list[str], int, int, int]:
        """Read and parse lines from file object.

        Args:
            file_obj: Open file object (local or S3)
            url_column: JSON key name containing URLs

        Returns:
            Tuple of (urls list, total lines, skipped lines, missing field count)
        """
        import json

        urls: list[str] = []
        line_number = 0
        skipped_lines = 0
        missing_field_count = 0

        for line in file_obj:
            line_number += 1
            line = line.strip()

            # Skip empty lines
            if not line:
                skipped_lines += 1
                continue

            # Skip comment lines
            if line.startswith("#"):
                skipped_lines += 1
                continue

            # Parse JSON object
            try:
                obj = json.loads(line)

                # Extract URL from specified field
                if isinstance(obj, dict):
                    if url_column in obj:
                        url_value = obj[url_column]
                        if isinstance(url_value, str) and url_value:
                            urls.append(url_value)
                        else:
                            logger.debug(
                                f"Line {line_number}: '{url_column}' field is not a non-empty string, skipping"
                            )
                            missing_field_count += 1
                    else:
                        missing_field_count += 1
                else:
                    logger.debug(f"Line {line_number}: JSON object is not a dict, skipping")
                    missing_field_count += 1

            except json.JSONDecodeError as e:
                logger.warning(f"Line {line_number}: Invalid JSON, skipping: {e}")
                skipped_lines += 1
                continue

        return urls, line_number, skipped_lines, missing_field_count

    def validate_format(self, file_path: str) -> bool:
        """Validate NDJSON file format.

        Args:
            file_path: Path to NDJSON file (local or s3://)

        Returns:
            True if file appears to be valid NDJSON, False otherwise.
        """

        try:
            if file_path.startswith("s3://"):
                if not HAS_FSSPEC or fsspec is None:
                    return False
                fs = fsspec.filesystem("s3")
                with fs.open(file_path, "r") as f:
                    return self._validate_content(f)
            else:
                with open(file_path) as f:
                    return self._validate_content(f)

        except (OSError, ValueError, UnicodeDecodeError):
            return False

    def _validate_content(self, file_obj) -> bool:
        """Validate file content by checking first few non-comment lines.

        Args:
            file_obj: Open file object

        Returns:
            True if file contains valid JSON objects, False otherwise.
        """
        import json

        valid_lines = 0
        max_lines_to_check = 10

        for line in file_obj:
            line = line.strip()

            # Skip empty and comment lines
            if not line or line.startswith("#"):
                continue

            # Try to parse as JSON
            try:
                obj = json.loads(line)
                if isinstance(obj, dict):
                    valid_lines += 1
                    if valid_lines >= max_lines_to_check:
                        return True
            except (json.JSONDecodeError, ValueError):
                return False

        # Return True if we found at least one valid JSON object
        return valid_lines > 0


class ReaderFactory:
    """Factory class for automatic input reader selection and configuration.

    This factory provides intelligent input reader selection based on file format
    specifications or automatic detection from file extensions. Eliminates the need
    for manual reader instantiation and configuration in most use cases.

    Automatic Format Detection:
        When format_name='auto', the factory analyzes file extensions to select
        the most appropriate reader:
        - .parquet, .pq → ParquetReader (high performance, recommended)
        - .csv → CSVReader with comma delimiter
        - .tsv, .tab → CSVReader with tab delimiter
        - .ndjson, .jsonl → NDJSONReader for newline-delimited JSON
        - Fallback to CSV for unknown extensions

    Supported Formats:
        - 'parquet': Optimized columnar format reader
        - 'csv': Flexible comma-separated values reader
        - 'tsv': Tab-separated values (CSV variant)
        - 'ndjson': Newline-delimited JSON reader
        - 'jsonl': Alias for ndjson format
        - 'auto': Automatic detection based on file extension

    Reader Configuration:
        The factory accepts configuration parameters that are passed to the
        underlying readers for customization:
        - CSV readers: delimiter, encoding, quoting options
        - Parquet readers: column selection and filtering options
        - All readers: error handling and validation settings

    Example:
        >>> # Automatic format detection (recommended)
        >>> reader = ReaderFactory.get_reader('auto')
        >>> urls = reader.read_urls('data.parquet', 'url')
        >>>
        >>> # Explicit format for performance-critical scenarios
        >>> reader = ReaderFactory.get_reader('parquet')
        >>>
        >>> # Custom CSV configuration
        >>> reader = ReaderFactory.get_reader('csv', delimiter='|', encoding='utf-8')

    Performance:
        The factory adds minimal overhead to reader creation and is optimized for
        repeated use. Reader instances can be reused across multiple files of the
        same format for better performance in batch processing scenarios.
    """

    _readers = {
        "parquet": ParquetReader,
        "csv": CSVReader,
        "tsv": lambda: CSVReader(delimiter="\t"),
        "ndjson": NDJSONReader,
        "jsonl": NDJSONReader,
    }

    @classmethod
    def get_reader(cls, format_name: str, **kwargs) -> InputReader:
        """Get reader for specified format."""
        format_name = format_name.lower()

        if format_name not in cls._readers:
            raise ValueError(f"Unsupported input format: {format_name}. Supported formats: {list(cls._readers.keys())}")

        reader_class = cls._readers[format_name]

        # Handle lambda functions for parameterized readers
        if callable(reader_class) and not isinstance(reader_class, type):
            result = reader_class()
            return cast(InputReader, result)

        # For type classes, instantiate with kwargs
        reader_type = cast(type[InputReader], reader_class)
        result = reader_type(**kwargs)
        return result

    @classmethod
    def auto_detect_format(cls, file_path: str) -> str:
        """Auto-detect file format based on extension."""
        path_lower = file_path.lower()

        if path_lower.endswith(".parquet") or path_lower.endswith(".pq"):
            return "parquet"
        elif path_lower.endswith(".csv"):
            return "csv"
        elif path_lower.endswith(".tsv") or path_lower.endswith(".tab"):
            return "tsv"
        elif path_lower.endswith(".ndjson"):
            return "ndjson"
        elif path_lower.endswith(".jsonl"):
            return "jsonl"
        else:
            # Default to CSV format for unknown extensions
            logger.warning(f"Unknown file extension for {file_path}, defaulting to CSV format")
            return "csv"

    @classmethod
    def get_supported_formats(cls) -> list[str]:
        """Get list of supported input formats."""
        return list(cls._readers.keys())
