"""Read H3-indexed tables by hex ranges."""

from __future__ import annotations

import asyncio
import time
from typing import TYPE_CHECKING, Any, Dict, List, Optional

if TYPE_CHECKING:
    import pyarrow as pa

# Check for nest_asyncio availability
try:
    import nest_asyncio

    HAS_NEST_ASYNCIO = True
except ImportError:
    HAS_NEST_ASYNCIO = False


def _require_job2(func_name: str) -> None:
    """Raise RuntimeError if job2 is not available."""
    try:
        import job2  # noqa: F401
    except ImportError as e:
        raise RuntimeError(
            f"{func_name} requires the job2 module. "
            "This function is only available in the Fused execution environment."
        ) from e


class _S3RangeFile:
    """File-like object that implements efficient range reads from S3 using boto3.

    This wrapper allows PyArrow to read only what it needs (footer + specific row groups)
    instead of downloading the entire file. Uses a shared boto3 client with connection
    pooling for efficiency.
    """

    def __init__(
        self,
        bucket: str,
        key: str,
        timing: Dict[str, float],
        data_sizes: Dict[str, int],
        verbose: bool = False,
    ):
        from job2.fasttortoise._reconstruction import _get_s3_client

        self.bucket = bucket
        self.key = key
        self.timing = timing
        self.data_sizes = data_sizes
        self.verbose = verbose
        self.position = 0
        self.total_bytes_read = 0
        self.num_requests = 0
        self.all_requests: list[dict[str, Any]] = []
        self._closed = False

        # Get shared S3 client (LRU cached with connection pooling)
        self.s3_client = _get_s3_client()

        # Get file size via HEAD
        response = self.s3_client.head_object(Bucket=bucket, Key=key)
        self.size = int(response["ContentLength"])

        if self.verbose:
            print(
                f"  S3RangeFile: File size = {self.size / 1024 / 1024:.2f} MB ({self.size} bytes)"
            )

    def read(self, size: int = -1) -> bytes:
        """Read up to size bytes from current position."""
        import time as time_module

        if size == -1:
            size = self.size - self.position

        if size <= 0 or self.position >= self.size:
            return b""

        # Don't read past end of file
        size = min(size, self.size - self.position)

        # Track this request
        request_info: Dict[str, Any] = {
            "request_num": self.num_requests + 1,
            "position": self.position,
            "size": size,
            "end_position": self.position + size - 1,
        }

        t_start = time_module.perf_counter()

        # Use boto3 range request
        range_header = f"bytes={self.position}-{self.position + size - 1}"
        response = self.s3_client.get_object(
            Bucket=self.bucket, Key=self.key, Range=range_header
        )
        data = response["Body"].read()

        elapsed_ms = (time_module.perf_counter() - t_start) * 1000
        request_info["elapsed_ms"] = elapsed_ms
        request_info["actual_bytes"] = len(data)

        # Track download timing
        self.timing["s3_download_ms"] = (
            self.timing.get("s3_download_ms", 0.0) + elapsed_ms
        )

        self.position += len(data)
        self.total_bytes_read += len(data)
        self.num_requests += 1
        self.all_requests.append(request_info)

        if self.verbose:
            distance_from_end = self.size - request_info["end_position"]
            print(
                f"  S3 Request #{request_info['request_num']}: "
                f"pos={request_info['position']:,} "
                f"size={request_info['size']:,} ({request_info['size'] / 1024 / 1024:.2f}MB) "
                f"end={request_info['end_position']:,} "
                f"(from_end={distance_from_end:,}) "
                f"{elapsed_ms:.0f}ms"
            )

        return data

    def seek(self, offset: int, whence: int = 0) -> int:
        """Seek to a position in the file."""
        if whence == 0:  # SEEK_SET
            self.position = offset
        elif whence == 1:  # SEEK_CUR
            self.position += offset
        elif whence == 2:  # SEEK_END
            self.position = self.size + offset

        # Clamp to valid range
        self.position = max(0, min(self.position, self.size))
        return self.position

    def tell(self) -> int:
        """Return current file position."""
        return self.position

    def close(self) -> None:
        """Close the file and update statistics."""
        if self._closed:
            return
        self._closed = True

        # Track actual bytes downloaded as "metadata_bytes_downloaded"
        self.data_sizes["metadata_bytes_downloaded"] = (
            self.data_sizes.get("metadata_bytes_downloaded", 0) + self.total_bytes_read
        )
        self.data_sizes["metadata_files_count"] = (
            self.data_sizes.get("metadata_files_count", 0) + 1
        )

        if self.verbose and self.all_requests:
            print("\n  S3 Range Request Summary:")
            print(f"    Total requests: {len(self.all_requests)}")
            print(
                f"    Total bytes: {self.total_bytes_read:,} ({self.total_bytes_read / 1024 / 1024:.2f} MB)"
            )
            print(f"    File size: {self.size:,} ({self.size / 1024 / 1024:.2f} MB)")
            print(
                f"    Downloaded: {self.total_bytes_read / self.size * 100:.1f}% of file"
            )

    def __enter__(self) -> "_S3RangeFile":
        return self

    def __exit__(self, *args: Any) -> None:
        self.close()

    @property
    def closed(self) -> bool:
        # PyArrow checks this attribute on file-like objects
        return self._closed


def read_hex_table_slow(
    dataset_path: str,
    hex_ranges_list: List[List[int]],
    columns: Optional[List[str]] = None,
    base_url: Optional[str] = None,
    verbose: bool = False,
    return_timing_info: bool = False,
    metadata_batch_size: int = 50,
) -> "pa.Table" | tuple["pa.Table", Dict[str, Any]]:
    """
    Read data from an H3-indexed dataset by querying hex ranges.

    This function queries the dataset index for row groups that match the given
    H3 hex ranges, downloads them in parallel with optimized batching, and returns
    a concatenated table.

    Adjacent row groups from the same file are combined into single downloads
    for better S3 performance. The batch size is controlled by
    `fused.options.row_group_batch_size` (default: 32KB).

    Args:
        dataset_path: Path to the H3-indexed dataset (e.g., "s3://bucket/dataset/")
        hex_ranges_list: List of [min_hex, max_hex] pairs as integers.
            Example: [[622236719905341439, 622246719905341439]]
        columns: Optional list of column names to read. If None, reads all columns.
        base_url: Base URL for API. If None, uses current environment.
        verbose: If True, print progress information. Default is False.
        return_timing_info: If True, return a tuple of (table, timing_info) instead of just the table.
            Default is False for backward compatibility.
        metadata_batch_size: Maximum number of row group metadata requests to batch together
            in a single API call. Larger batches reduce API overhead. Default is 50.
            Consider MongoDB's 16KB document limit when adjusting this value.

    Returns:
        PyArrow Table containing the concatenated data from all matching row groups.
        If return_timing_info is True, returns a tuple of (table, timing_info dict).

    Example:
        import fused

        # Read data for a specific H3 hex range
        table = fused.h3.read_hex_table_slow(
            dataset_path="s3://my-bucket/my-h3-dataset/",
            hex_ranges_list=[[622236719905341439, 622246719905341439]]
        )
        df = table.to_pandas()
    """
    import pyarrow as pa

    from fused._fasttortoise._api import get_row_groups_for_dataset
    from fused._options import options as OPTIONS

    # Use current environment's base URL if not specified
    if base_url is None:
        base_url = OPTIONS.base_url

    if not hex_ranges_list:
        return pa.table({})

    # Convert integer hex ranges to the format expected by get_row_groups_for_dataset
    geographical_regions = []
    for hex_range in hex_ranges_list:
        if len(hex_range) != 2:
            raise ValueError(
                f"Each hex range must be a list of [min, max], got {hex_range}"
            )
        min_hex, max_hex = hex_range
        geographical_regions.append({"min": f"{min_hex:x}", "max": f"{max_hex:x}"})

    # Query the dataset index to find matching row groups
    t0 = time.perf_counter()
    row_groups = get_row_groups_for_dataset(
        dataset_path=dataset_path,
        geographical_regions=geographical_regions,
        base_url=base_url,
    )
    t_api = time.perf_counter()

    if verbose:
        print(f"  API query: {(t_api - t0) * 1000:.1f}ms")
        print(f"  Found {len(row_groups)} row groups matching geo query")

    if not row_groups:
        # No matching row groups for the given hex ranges
        # This is normal - just return an empty table
        return pa.table({})

    # Get the batch size from options
    batch_size = OPTIONS.row_group_batch_size

    if verbose:
        print(f"  Using batch size: {batch_size} bytes")
        print(f"  Using metadata batch size: {metadata_batch_size}")

    # Run the pipelined fetch and download
    _require_job2("read_hex_table_slow")
    from job2.fasttortoise._h3_read import _fetch_with_combining

    tables, timing_info = _run_async(
        _fetch_with_combining(
            row_groups, base_url, columns, batch_size, verbose, metadata_batch_size
        )
    )
    t_fetch = time.perf_counter()

    if verbose:
        print(f"  Metadata + download: {(t_fetch - t_api) * 1000:.1f}ms")
        if timing_info:
            print(
                f"    Metadata fetch: {timing_info.get('metadata_wall_ms', 0):.1f}ms wall-clock, "
                f"{timing_info.get('metadata_ms', 0):.1f}ms cumulative"
            )
            print(
                f"      Longest metadata fetch: {timing_info.get('longest_metadata_fetch_ms', 0):.1f}ms"
            )
            print(
                f"    Data download: {timing_info.get('download_wall_ms', 0):.1f}ms wall-clock, "
                f"{timing_info.get('download_ms', 0):.1f}ms cumulative"
            )
            print(
                f"      Longest download: {timing_info.get('longest_download_ms', 0):.1f}ms"
            )
            print(f"    Download groups: {timing_info.get('num_groups', 0)}")

    # Concatenate all tables into one
    if not tables:
        return pa.table({})
    if len(tables) == 1:
        return tables[0]

    t_concat_start = time.perf_counter()
    # Use promote_options to handle schema mismatches (e.g., float vs double)
    # "permissive" allows type promotion like float->double
    result = pa.concat_tables(tables, promote_options="permissive")
    t_concat = time.perf_counter()

    if verbose:
        print(f"  Concat tables: {(t_concat - t_concat_start) * 1000:.1f}ms")

    if return_timing_info:
        return result, timing_info
    return result


def _run_async(coro):
    """Run an async coroutine, handling existing event loops."""

    async def _run_with_cleanup():
        """Run coroutine and clean up shared S3 client afterwards.

        Only cleans up when we own the event loop (not when nested via nest_asyncio).
        Note: We no longer clean up the aiohttp session here because _fetch_with_combining
        creates and manages its own session.
        """
        try:
            return await coro
        finally:
            # Clean up shared S3 client (job2) to prevent "Unclosed" warnings
            # This is safe because we own this event loop
            try:
                from job2.fasttortoise._reconstruction import (
                    _async_s3_client_loop,
                    _shared_async_s3_client,
                )

                # Only clean up if the client belongs to this event loop
                current_loop = asyncio.get_running_loop()
                if (
                    _shared_async_s3_client is not None
                    and _async_s3_client_loop is current_loop
                ):
                    await _shared_async_s3_client.close()
            except Exception:
                pass

    # Handle running inside an existing event loop (e.g., Jupyter, UDF runner)
    try:
        asyncio.get_running_loop()
        # We're inside an event loop - apply nest_asyncio if available
        if HAS_NEST_ASYNCIO:
            nest_asyncio.apply()
            # DON'T clean up here - the outer loop owns the session
            return asyncio.run(coro)
        else:
            # Fallback: create a new thread to run the async code
            # Clean up since we own this isolated event loop
            import concurrent.futures

            with concurrent.futures.ThreadPoolExecutor() as executor:
                future = executor.submit(asyncio.run, _run_with_cleanup())
                return future.result()
    except RuntimeError:
        # No event loop running - use asyncio.run directly
        # Clean up since we own this event loop
        return asyncio.run(_run_with_cleanup())


def read_hex_table(
    dataset_path: str,
    hex_ranges_list: List[List[int]],
    columns: Optional[List[str]] = None,
    base_url: Optional[str] = None,
    verbose: bool = False,
    return_timing_info: bool = False,
    batch_size: Optional[int] = None,
    max_concurrent_downloads: Optional[int] = None,
) -> "pa.Table" | tuple["pa.Table", Dict[str, Any]]:
    """
    Read data from an H3-indexed dataset by querying hex ranges.

    This is an optimized version that assumes the server always provides full metadata
    (start_offset, end_offset, metadata, and row_group_bytes) for all row groups.
    If any row group is missing required metadata, this function will raise an error
    indicating that the dataset needs to be re-indexed.

    This function eliminates all metadata API calls by using prefetched metadata from
    the /datasets/items-with-metadata endpoint.

    This function imports the implementation from job2 at runtime,
    similar to how read_parquet_row_group works.

    Args:
        dataset_path: Path to the H3-indexed dataset (e.g., "s3://bucket/dataset/")
        hex_ranges_list: List of [min_hex, max_hex] pairs as integers.
            Example: [[622236719905341439, 622246719905341439]]
        columns: Optional list of column names to read. If None, reads all columns.
        base_url: Base URL for API. If None, uses current environment.
        verbose: If True, print progress information. Default is False.
        return_timing_info: If True, return a tuple of (table, timing_info) instead of just the table.
            Default is False for backward compatibility.
        batch_size: Target size in bytes for combining row groups. If None, uses
            `fused.options.row_group_batch_size` (default: 32KB).
        max_concurrent_downloads: Maximum number of simultaneous download operations. If None,
            uses a default based on the number of files. Default is None.

    Returns:
        PyArrow Table containing the concatenated data from all matching row groups.
        If return_timing_info is True, returns a tuple of (table, timing_info dict).

    Raises:
        ValueError: If any row group is missing required metadata (start_offset, end_offset,
            metadata, or row_group_bytes). This indicates the dataset needs to be re-indexed.

    Example:
        import fused

        # Read data for a specific H3 hex range
        table = fused.h3.read_hex_table(
            dataset_path="s3://my-bucket/my-h3-dataset/",
            hex_ranges_list=[[622236719905341439, 622246719905341439]]
        )
        df = table.to_pandas()
    """
    _require_job2("read_hex_table")
    from job2.fasttortoise import read_hex_table as _job2_read_hex_table

    return _job2_read_hex_table(
        dataset_path=dataset_path,
        hex_ranges_list=hex_ranges_list,
        columns=columns,
        base_url=base_url,
        verbose=verbose,
        return_timing_info=return_timing_info,
        batch_size=batch_size,
        max_concurrent_downloads=max_concurrent_downloads,
    )


def read_hex_table_with_persisted_metadata(
    dataset_path: str,
    hex_ranges_list: List[List[int]],
    columns: Optional[List[str]] = None,
    metadata_path: Optional[str] = None,
    verbose: bool = False,
    return_timing_info: bool = False,
    batch_size: Optional[int] = None,
    max_concurrent_downloads: Optional[int] = None,
    max_concurrent_metadata_reads: Optional[int] = None,
) -> "pa.Table" | tuple["pa.Table", Dict[str, Any]]:
    """
    Read data from an H3-indexed dataset using persisted metadata parquet.

    This function reads from per-file metadata parquet files instead of
    querying a server. Each source parquet file has a corresponding
    .metadata.parquet file stored at:
    {source_dir}/.fused/{source_filename}.metadata.parquet

    Or at the specified metadata_path if provided:
    {metadata_path}/{full_source_path}.metadata.parquet

    Supports subdirectory queries - if dataset_path points to a subdirectory,
    the function will look for metadata files for files in that subdirectory.

    Args:
        dataset_path: Path to the H3-indexed dataset (e.g., "s3://bucket/dataset/")
            Can also be a subdirectory path for filtering (e.g., "s3://bucket/dataset/year=2024/")
        hex_ranges_list: List of [min_hex, max_hex] pairs as integers.
            Example: [[622236719905341439, 622246719905341439]]
        columns: Optional list of column names to read. If None, reads all columns.
        metadata_path: Directory path where metadata files are stored.
                      If None, looks for metadata at {source_dir}/.fused/ for each source file.
                      If provided, reads metadata files from this location using full source paths.
                      This allows reading metadata from a different location when
                      you don't have access to the original dataset directory.
        verbose: If True, print progress information. Default is False.
        return_timing_info: If True, return a tuple of (table, timing_info) instead of just the table.
            Default is False for backward compatibility.
        batch_size: Target size in bytes for combining row groups. If None, uses
            `fused.options.row_group_batch_size` (default: 32KB).
        max_concurrent_downloads: Maximum number of simultaneous download operations. If None,
            uses a default based on the number of files. Default is None.
        max_concurrent_metadata_reads: Maximum number of simultaneous metadata file reads. If None,
                        defaults to min(10, len(source_files)).

    Returns:
        PyArrow Table containing the concatenated data from all matching row groups.
        If return_timing_info is True, returns a tuple of (table, timing_info dict).

    Raises:
        FileNotFoundError: If the metadata parquet file is not found.
        ValueError: If any row group is missing required metadata.

    Example:
        import fused

        # First, persist metadata (one-time operation)
        fused.h3.persist_hex_table_metadata("s3://my-bucket/my-dataset/")

        # Then read using persisted metadata (no server required)
        table = fused.h3.read_hex_table_with_persisted_metadata(
            dataset_path="s3://my-bucket/my-dataset/",
            hex_ranges_list=[[622236719905341439, 622246719905341439]]
        )
        df = table.to_pandas()

        # Read from subdirectory (filters by path prefix)
        table = fused.h3.read_hex_table_with_persisted_metadata(
            dataset_path="s3://my-bucket/my-dataset/year=2024/",
            hex_ranges_list=[[622236719905341439, 622246719905341439]]
        )

        # Read with metadata in alternate location
        table = fused.h3.read_hex_table_with_persisted_metadata(
            dataset_path="s3://my-bucket/my-dataset/",
            metadata_path="s3://my-bucket/metadata/",
            hex_ranges_list=[[622236719905341439, 622246719905341439]]
        )

    """
    import pyarrow as pa
    import pyarrow.parquet as pq

    from fused._options import options as OPTIONS

    # Lazy import from job2
    _require_job2("read_hex_table_with_persisted_metadata")
    from job2.fasttortoise._h3_index import (
        _get_metadata_file_path,
        _list_parquet_files,
        is_single_parquet_file,
    )

    if not hex_ranges_list:
        return pa.table({})

    t0 = time.perf_counter()

    import fsspec

    # Initialize timing breakdown and data size tracking
    timing_breakdown = {
        "file_listing_ms": 0.0,
        "s3_exists_check_ms": 0.0,
        "s3_download_ms": 0.0,
        "s3_open_ms": 0.0,
        "parquet_file_init_ms": 0.0,
        "row_group_filter_ms": 0.0,  # Time to filter row groups using statistics
        "read_table_ms": 0.0,
        "schema_extract_ms": 0.0,
        "json_decode_ms": 0.0,
        "filter_ms": 0.0,
        "convert_to_row_groups_ms": 0.0,
        "validation_ms": 0.0,
        "download_ms": 0.0,
        "concat_ms": 0.0,
        "row_group_stats": {
            "total": 0,
            "matched": 0,
        },  # Track row group filtering stats
    }

    # Track data sizes
    data_sizes = {
        "metadata_bytes_downloaded": 0,  # Total metadata files downloaded (full files)
        "metadata_bytes_read": 0,  # Bytes actually read from parquet
        "metadata_bytes_used": 0,  # Total metadata bytes actually used (after filtering)
        "data_bytes": 0,  # Total row group data downloaded
        "metadata_files_count": 0,
    }

    t_listing = time.perf_counter()
    # Determine source files to query
    if dataset_split := is_single_parquet_file(dataset_path):
        # Single file case
        dataset_root = dataset_split[0]
        source_files = [dataset_split[1]]
    else:
        # Directory case - list all parquet files
        dataset_root = dataset_path.rstrip("/")
        source_files = _list_parquet_files(dataset_path)
    timing_breakdown["file_listing_ms"] = (time.perf_counter() - t_listing) * 1000

    if not source_files:
        raise FileNotFoundError(f"No parquet files found in {dataset_path}")

    if verbose:
        print(f"Found {len(source_files)} source files to query")

    # Determine max concurrent metadata reads
    if max_concurrent_metadata_reads is None:
        max_concurrent_metadata_reads = min(10, len(source_files))

    from concurrent.futures import ThreadPoolExecutor, as_completed

    def _read_single_metadata_file(
        rel_path: str,
    ) -> tuple[list[dict], Dict[str, float], Dict[str, int], float]:
        """Read metadata and row-group info for a single source file.

        Returns:
            (row_groups_list, timing_delta, data_sizes_delta, metadata_read_time)
        """
        full_source_path = f"{dataset_root}/{rel_path}"

        # Local timing and size accumulators
        file_timing: Dict[str, float] = {
            "s3_exists_check_ms": 0.0,
            "s3_download_ms": 0.0,
            "s3_open_ms": 0.0,
            "parquet_file_init_ms": 0.0,
            "row_group_filter_ms": 0.0,
            "read_table_ms": 0.0,
            "schema_extract_ms": 0.0,
            "json_decode_ms": 0.0,
            "filter_ms": 0.0,
            "convert_to_row_groups_ms": 0.0,
        }
        file_data_sizes: Dict[str, int] = {
            "metadata_bytes_downloaded": 0,
            "metadata_bytes_read": 0,
            "metadata_bytes_used": 0,
            "metadata_files_count": 0,
        }
        file_row_group_stats = {"total": 0, "matched": 0}
        file_metadata_read_time = 0.0

        # Determine metadata file path
        metadata_file_path = _get_metadata_file_path(
            full_source_path, dataset_root, metadata_path
        )

        # S3 paths: use _S3RangeFile + boto3, not fsspec
        is_s3 = metadata_file_path.startswith("s3://")

        try:
            parquet_file = None

            if is_s3:
                # Direct S3 range-based path using boto3
                # Parse s3://bucket/key
                path_parts = metadata_file_path[5:].split("/", 1)
                bucket = path_parts[0]
                key = path_parts[1] if len(path_parts) > 1 else ""

                t_parquet_init = time.perf_counter()
                s3_file_handle = _S3RangeFile(
                    bucket,
                    key,
                    file_timing,
                    file_data_sizes,
                    verbose=verbose,
                )
                parquet_file = pq.ParquetFile(s3_file_handle)
                file_timing["parquet_file_init_ms"] = (
                    time.perf_counter() - t_parquet_init
                ) * 1000
            else:
                # fsspec-based path (local/other filesystems)
                t_exists = time.perf_counter()
                fs, fs_path = fsspec.core.url_to_fs(
                    metadata_file_path,
                    use_listings_cache=False,
                )

                if not fs.exists(fs_path):
                    if verbose:
                        print(
                            f"  Skipping {rel_path} - metadata file not found: {metadata_file_path}"
                        )
                    return [], {}, {}, 0.0

                # Get metadata file size if available
                try:
                    if hasattr(fs, "info"):
                        file_info = fs.info(fs_path)
                        if "size" in file_info:
                            file_data_sizes["metadata_bytes_downloaded"] += int(
                                file_info["size"]
                            )
                            file_data_sizes["metadata_files_count"] += 1
                except Exception:
                    pass

                file_timing["s3_exists_check_ms"] += (
                    time.perf_counter() - t_exists
                ) * 1000

                # Read metadata file with row group filtering
                # Note: we open file without `with` so handle stays open for parquet_file
                t_open = time.perf_counter()
                f = fs.open(fs_path, "rb")
                file_timing["s3_open_ms"] += (time.perf_counter() - t_open) * 1000

                t_parquet_init = time.perf_counter()
                parquet_file = pq.ParquetFile(f)
                file_timing["parquet_file_init_ms"] += (
                    time.perf_counter() - t_parquet_init
                ) * 1000

            # At this point parquet_file is initialized using either S3RangeFile or fsspec
            if parquet_file is None:
                return [], {}, {}, 0.0

            # Extract binary metadata fields directly from schema metadata
            t_schema = time.perf_counter()
            schema = parquet_file.schema_arrow

            metadata_dict: Dict[str, Any] = {}
            if b"parquet_version_bytes" in schema.metadata:
                metadata_dict["parquet_version_bytes"] = schema.metadata[
                    b"parquet_version_bytes"
                ]
            if b"parquet_schema_bytes" in schema.metadata:
                metadata_dict["schema_info"] = {
                    "parquet_schema_bytes": schema.metadata[b"parquet_schema_bytes"]
                }
            if b"key_value_metadata_bytes" in schema.metadata:
                metadata_dict["key_value_metadata_bytes"] = schema.metadata[
                    b"key_value_metadata_bytes"
                ]
            if b"created_by_bytes" in schema.metadata:
                metadata_dict["created_by_bytes"] = schema.metadata[b"created_by_bytes"]
            if b"column_orders_bytes" in schema.metadata:
                metadata_dict["column_orders_bytes"] = schema.metadata[
                    b"column_orders_bytes"
                ]

            if not metadata_dict:
                if verbose:
                    print(
                        f"  Warning: {rel_path} metadata file missing binary metadata in schema"
                    )
                return [], {}, {}, 0.0

            file_timing["schema_extract_ms"] += (time.perf_counter() - t_schema) * 1000

            # Filter row groups using statistics before reading
            read_start = time.perf_counter()
            t_rg_filter = time.perf_counter()
            matching_row_groups: List[int] = []
            metadata = parquet_file.metadata
            total_row_groups = metadata.num_row_groups

            for rg_idx in range(total_row_groups):
                rg = metadata.row_group(rg_idx)

                h3_min_col_idx = None
                h3_max_col_idx = None
                for col_idx in range(rg.num_columns):
                    col_name = schema.names[col_idx]
                    if col_name == "h3_min":
                        h3_min_col_idx = col_idx
                    elif col_name == "h3_max":
                        h3_max_col_idx = col_idx

                if h3_min_col_idx is not None and h3_max_col_idx is not None:
                    h3_min_col = rg.column(h3_min_col_idx)
                    h3_max_col = rg.column(h3_max_col_idx)

                    rg_h3_min = None
                    rg_h3_max = None
                    if h3_min_col.statistics and h3_min_col.statistics.min is not None:
                        rg_h3_min = (
                            h3_min_col.statistics.min.as_py()
                            if hasattr(h3_min_col.statistics.min, "as_py")
                            else int(h3_min_col.statistics.min)
                        )
                    if h3_max_col.statistics and h3_max_col.statistics.max is not None:
                        rg_h3_max = (
                            h3_max_col.statistics.max.as_py()
                            if hasattr(h3_max_col.statistics.max, "as_py")
                            else int(h3_max_col.statistics.max)
                        )

                    if rg_h3_min is not None and rg_h3_max is not None:
                        matches = False
                        for q_min, q_max in hex_ranges_list:
                            if rg_h3_min <= q_max and rg_h3_max >= q_min:
                                matches = True
                                break

                        if matches:
                            matching_row_groups.append(rg_idx)
                    else:
                        matching_row_groups.append(rg_idx)
                else:
                    matching_row_groups.append(rg_idx)

            file_timing["row_group_filter_ms"] += (
                time.perf_counter() - t_rg_filter
            ) * 1000
            file_row_group_stats["total"] = total_row_groups
            file_row_group_stats["matched"] = len(matching_row_groups)

            # Read only matching row groups
            t_read_table = time.perf_counter()
            if matching_row_groups:
                metadata_table = parquet_file.read_row_groups(matching_row_groups)
            else:
                if verbose:
                    print(
                        f"  No matching row groups in {rel_path} (checked {total_row_groups} row groups)"
                    )
                return [], {}, {}, 0.0

            file_timing["read_table_ms"] += (time.perf_counter() - t_read_table) * 1000
            file_data_sizes["metadata_bytes_read"] += metadata_table.nbytes

            # Filter by H3 ranges on the read data (in case row group stats weren't precise)
            t_filter = time.perf_counter()
            filtered_table = _filter_by_h3_ranges(metadata_table, hex_ranges_list)
            file_timing["filter_ms"] += (time.perf_counter() - t_filter) * 1000

            reference_table = metadata_table
            if len(filtered_table) == 0:
                return [], {}, {}, 0.0

            file_metadata_read_time += time.perf_counter() - read_start

            # Track actual bytes used (size of filtered table)
            try:
                filtered_size = filtered_table.nbytes
                file_data_sizes["metadata_bytes_used"] += filtered_size
            except Exception:
                if len(reference_table) > 0:
                    avg_row_size = reference_table.nbytes / len(reference_table)
                    file_data_sizes["metadata_bytes_used"] += int(
                        avg_row_size * len(filtered_table)
                    )

            # Convert to row group format
            t_convert = time.perf_counter()
            row_group_indices = filtered_table["row_group_index"].to_pylist()
            start_offsets = filtered_table["start_offset"].to_pylist()
            end_offsets = filtered_table["end_offset"].to_pylist()
            row_group_bytes_list = filtered_table["row_group_bytes"].to_pylist()

            new_row_groups = [
                {
                    "path": full_source_path,
                    "row_group_index": idx,
                    "start_offset": start,
                    "end_offset": end,
                    "metadata": metadata_dict,
                    "row_group_bytes": bytes_str,
                }
                for idx, start, end, bytes_str in zip(
                    row_group_indices, start_offsets, end_offsets, row_group_bytes_list
                )
            ]
            file_timing["convert_to_row_groups_ms"] += (
                time.perf_counter() - t_convert
            ) * 1000

            # Attach row_group_stats to timing dict
            file_timing["row_group_stats"] = file_row_group_stats

            return new_row_groups, file_timing, file_data_sizes, file_metadata_read_time

        except Exception as e:
            # Provide more informative error messages for common S3 errors
            error_msg = str(e)
            try:
                from botocore.exceptions import ClientError

                if isinstance(e, ClientError):
                    error_code = e.response.get("Error", {}).get("Code", "Unknown")
                    if error_code == "NoSuchKey":
                        error_msg = f"Metadata file not found: {metadata_file_path}"
                    elif error_code == "AccessDenied":
                        error_msg = (
                            f"Access denied to metadata file: {metadata_file_path}"
                        )
                    else:
                        error_msg = f"S3 error ({error_code}): {e}"
            except ImportError:
                pass

            if verbose:
                print(f"  Warning: Failed to read metadata for {rel_path}: {error_msg}")
            return [], {}, {}, 0.0

    # Read metadata for each source file and collect row groups (in parallel)
    all_row_groups: list[dict[str, Any]] = []
    metadata_read_time = 0.0

    with ThreadPoolExecutor(max_workers=max_concurrent_metadata_reads) as executor:
        future_to_rel = {
            executor.submit(_read_single_metadata_file, rel_path): rel_path
            for rel_path in source_files
        }

        for future in as_completed(future_to_rel):
            rel_path = future_to_rel[future]
            try:
                row_groups, file_timing, file_sizes, file_meta_time = future.result()
                all_row_groups.extend(row_groups)
                metadata_read_time += file_meta_time

                # Aggregate timing
                for key, val in file_timing.items():
                    if key == "row_group_stats":
                        stats = val or {}
                        timing_breakdown["row_group_stats"]["total"] += stats.get(
                            "total", 0
                        )
                        timing_breakdown["row_group_stats"]["matched"] += stats.get(
                            "matched", 0
                        )
                    else:
                        timing_breakdown[key] += val

                # Aggregate data sizes
                for key, val in file_sizes.items():
                    data_sizes[key] += val
            except Exception as e:
                if verbose:
                    print(f"  Warning: Failed to process metadata for {rel_path}: {e}")
                continue

    # Calculate total data size from row groups
    for rg in all_row_groups:
        start_offset = rg.get("start_offset")
        end_offset = rg.get("end_offset")
        if start_offset is not None and end_offset is not None:
            data_sizes["data_bytes"] += end_offset - start_offset

    t_read = time.perf_counter()

    # Calculate download rates
    metadata_download_rate = 0.0
    if (
        timing_breakdown["s3_download_ms"] > 0
        and data_sizes["metadata_bytes_downloaded"] > 0
    ):
        metadata_download_rate = (
            data_sizes["metadata_bytes_downloaded"] / 1024 / 1024
        ) / (timing_breakdown["s3_download_ms"] / 1000)  # MB/s

    if verbose:
        print(
            f"  Read metadata: {(t_read - t0) * 1000:.1f}ms ({len(all_row_groups)} row groups from {len(source_files)} files)"
        )
        print("    Data sizes:")
        if data_sizes["metadata_bytes_downloaded"] > 0:
            metadata_downloaded_mb = (
                data_sizes["metadata_bytes_downloaded"] / 1024 / 1024
            )
            metadata_read_mb = data_sizes["metadata_bytes_read"] / 1024 / 1024
            metadata_used_mb = data_sizes["metadata_bytes_used"] / 1024 / 1024
            print(
                f"      Metadata files downloaded: {metadata_downloaded_mb:.2f} MB ({data_sizes['metadata_files_count']} files)"
            )
            if metadata_read_mb > 0:
                read_efficiency = (
                    (metadata_read_mb / metadata_downloaded_mb * 100)
                    if metadata_downloaded_mb > 0
                    else 0
                )
                print(
                    f"      Metadata bytes read (row group filtered): {metadata_read_mb:.2f} MB ({read_efficiency:.1f}% of downloaded)"
                )
            if metadata_used_mb > 0:
                efficiency = (
                    (metadata_used_mb / metadata_downloaded_mb * 100)
                    if metadata_downloaded_mb > 0
                    else 0
                )
                print(
                    f"      Metadata bytes used (after filtering): {metadata_used_mb:.2f} MB ({efficiency:.1f}% of downloaded)"
                )
            if timing_breakdown["s3_download_ms"] > 0:
                print(
                    f"      Metadata download rate: {metadata_download_rate:.2f} MB/s"
                )
        data_mb = data_sizes["data_bytes"] / 1024 / 1024
        print(
            f"      Row group data: {data_mb:.2f} MB ({len(all_row_groups)} row groups)"
        )
        print("    Timing breakdown:")
        print(f"      File listing: {timing_breakdown['file_listing_ms']:.1f}ms")
        print(f"      S3 exists check: {timing_breakdown['s3_exists_check_ms']:.1f}ms")
        print(f"      S3 open: {timing_breakdown['s3_open_ms']:.1f}ms")
        print(
            f"      ParquetFile init: {timing_breakdown['parquet_file_init_ms']:.1f}ms"
        )
        if timing_breakdown["row_group_filter_ms"] > 0:
            # Get row group stats from timing_breakdown if available
            rg_stats = timing_breakdown.get("row_group_stats", {})
            total_rgs = rg_stats.get("total", 0)
            matched_rgs = rg_stats.get("matched", 0)
            if total_rgs > 0:
                efficiency = (matched_rgs / total_rgs * 100) if total_rgs > 0 else 0
                skipped_rgs = total_rgs - matched_rgs
                print(
                    f"      Row group filter (stats): {timing_breakdown['row_group_filter_ms']:.1f}ms"
                )
                print(
                    f"        Total row groups: {total_rgs:,}, Matched: {matched_rgs:,}, Skipped: {skipped_rgs:,} ({efficiency:.1f}% match rate)"
                )
            else:
                print(
                    f"      Row group filter (stats): {timing_breakdown['row_group_filter_ms']:.1f}ms"
                )
        print(f"      Read table: {timing_breakdown['read_table_ms']:.1f}ms")
        print(f"      Schema extract: {timing_breakdown['schema_extract_ms']:.1f}ms")
        print(f"      JSON decode: {timing_breakdown['json_decode_ms']:.1f}ms")
        print(f"      Filter: {timing_breakdown['filter_ms']:.1f}ms")
        print(
            f"      Convert to row groups: {timing_breakdown['convert_to_row_groups_ms']:.1f}ms"
        )

    if len(all_row_groups) == 0:
        if return_timing_info:
            timing_info = {
                "metadata_ms": (t_read - t0) * 1000,
                "download_ms": 0,
                "num_groups": 0,
                "timing_breakdown": timing_breakdown,
                "data_sizes": data_sizes,
                "metadata_download_rate_mb_s": metadata_download_rate,
                "data_download_rate_mb_s": 0.0,
                "data_download_wall_rate_mb_s": 0.0,
            }
            return pa.table({}), timing_info
        return pa.table({})

    # Validate that all row groups have required metadata
    t_validate = time.perf_counter()
    for rg in all_row_groups:
        missing_fields = []
        if rg.get("start_offset") is None:
            missing_fields.append("start_offset")
        if rg.get("end_offset") is None:
            missing_fields.append("end_offset")
        if not rg.get("metadata"):
            missing_fields.append("metadata")
        if not rg.get("row_group_bytes"):
            missing_fields.append("row_group_bytes")

        if missing_fields:
            raise ValueError(
                f"Row group {rg.get('row_group_index')} in {rg.get('path')} is missing required metadata: "
                f"{', '.join(missing_fields)}. The dataset needs to be re-indexed with persist_hex_table_metadata()."
            )
    timing_breakdown["validation_ms"] = (time.perf_counter() - t_validate) * 1000

    # Get the batch size from options if not provided
    if batch_size is None:
        batch_size = OPTIONS.row_group_batch_size

    if verbose:
        print(f"  Using batch size: {batch_size} bytes")
        print("  Using persisted metadata (no API calls)")

    # Run the prefetched fetch and download
    async def _run_with_cleanup():
        """Run coroutine and clean up shared S3 client afterwards."""
        from job2.fasttortoise._h3_read import _fetch_with_combining_prefetched

        try:
            return await _fetch_with_combining_prefetched(
                all_row_groups,
                "",  # base_url not used for persisted metadata
                columns,
                batch_size,
                verbose,
                max_concurrent_downloads,
            )
        finally:
            # Clean up shared S3 client to prevent "Unclosed" warnings
            try:
                from job2.fasttortoise._reconstruction import _shared_async_s3_client

                if _shared_async_s3_client is not None:
                    await _shared_async_s3_client.close()
            except Exception:
                pass

    tables, download_timing_info = _run_async(_run_with_cleanup())
    t_fetch = time.perf_counter()
    timing_breakdown["download_ms"] = (t_fetch - t_read) * 1000

    # Calculate data download rate
    data_download_rate = 0.0
    data_download_wall_rate = 0.0
    if download_timing_info:
        download_wall_ms = download_timing_info.get("download_wall_ms", 0)
        if download_wall_ms > 0 and data_sizes["data_bytes"] > 0:
            data_download_wall_rate = (data_sizes["data_bytes"] / 1024 / 1024) / (
                download_wall_ms / 1000
            )  # MB/s
        download_cumulative_ms = download_timing_info.get("download_ms", 0)
        if download_cumulative_ms > 0 and data_sizes["data_bytes"] > 0:
            data_download_rate = (data_sizes["data_bytes"] / 1024 / 1024) / (
                download_cumulative_ms / 1000
            )  # MB/s

    if verbose:
        print(f"  Download: {(t_fetch - t_read) * 1000:.1f}ms")
        if data_sizes["data_bytes"] > 0:
            data_mb = data_sizes["data_bytes"] / 1024 / 1024
            print(f"    Data downloaded: {data_mb:.2f} MB")
            if data_download_wall_rate > 0:
                print(
                    f"    Effective download rate: {data_download_wall_rate:.2f} MB/s (wall-clock)"
                )
            if data_download_rate > 0 and data_download_rate != data_download_wall_rate:
                print(f"    Cumulative download rate: {data_download_rate:.2f} MB/s")
        if download_timing_info:
            print(
                f"    Data download: {download_timing_info.get('download_wall_ms', 0):.1f}ms wall-clock, "
                f"{download_timing_info.get('download_ms', 0):.1f}ms cumulative"
            )
            print(f"    Download groups: {download_timing_info.get('num_groups', 0)}")

    # Concatenate all tables into one
    if not tables:
        if return_timing_info:
            timing_info = {
                "metadata_ms": (t_read - t0) * 1000,
                "download_ms": timing_breakdown["download_ms"],
                "num_groups": download_timing_info.get("num_groups", 0)
                if download_timing_info
                else 0,
                "timing_breakdown": timing_breakdown,
                "data_sizes": data_sizes,
                "metadata_download_rate_mb_s": metadata_download_rate,
                "data_download_rate_mb_s": data_download_rate,
                "data_download_wall_rate_mb_s": data_download_wall_rate,
            }
            if download_timing_info:
                timing_info.update(download_timing_info)
            return pa.table({}), timing_info
        return pa.table({})
    if len(tables) == 1:
        if return_timing_info:
            timing_info = {
                "metadata_ms": (t_read - t0) * 1000,
                "download_ms": timing_breakdown["download_ms"],
                "num_groups": download_timing_info.get("num_groups", 0)
                if download_timing_info
                else 0,
                "timing_breakdown": timing_breakdown,
                "data_sizes": data_sizes,
                "metadata_download_rate_mb_s": metadata_download_rate,
                "data_download_rate_mb_s": data_download_rate,
                "data_download_wall_rate_mb_s": data_download_wall_rate,
            }
            if download_timing_info:
                timing_info.update(download_timing_info)
            return tables[0], timing_info
        return tables[0]

    t_concat_start = time.perf_counter()
    result = pa.concat_tables(tables, promote_options="permissive")
    t_concat = time.perf_counter()
    timing_breakdown["concat_ms"] = (t_concat - t_concat_start) * 1000

    if verbose:
        print(f"  Concat tables: {(t_concat - t_concat_start) * 1000:.1f}ms")

    if return_timing_info:
        timing_info = {
            "metadata_ms": (t_read - t0) * 1000,
            "download_ms": timing_breakdown["download_ms"],
            "num_groups": download_timing_info.get("num_groups", 0)
            if download_timing_info
            else 0,
            "timing_breakdown": timing_breakdown,
            "total_ms": (t_concat - t0) * 1000,
            "data_sizes": data_sizes,
            "metadata_download_rate_mb_s": metadata_download_rate,
            "data_download_rate_mb_s": data_download_rate,
            "data_download_wall_rate_mb_s": data_download_wall_rate,
        }
        if download_timing_info:
            timing_info.update(download_timing_info)
        return result, timing_info
    return result


def _filter_by_h3_ranges(table: "pa.Table", hex_ranges: List[List[int]]) -> "pa.Table":
    """Filter metadata table by H3 range overlap.

    For each query range [q_min, q_max], finds rows where:
    h3_min <= q_max AND h3_max >= q_min (range overlap)

    Args:
        table: PyArrow table with h3_min and h3_max columns
        hex_ranges: List of [min, max] integer pairs

    Returns:
        Filtered table with only matching row groups
    """
    import pyarrow.compute as pc

    if not hex_ranges:
        return table

    masks = []
    for q_min, q_max in hex_ranges:
        # Range overlap: h3_min <= q_max AND h3_max >= q_min
        mask = pc.and_(
            pc.less_equal(table["h3_min"], q_max),
            pc.greater_equal(table["h3_max"], q_min),
        )
        masks.append(mask)

    # Combine all masks with OR
    combined = masks[0]
    for m in masks[1:]:
        combined = pc.or_(combined, m)

    return table.filter(combined)
