"""Efficient file reader for large JSONL files."""

import json
from pathlib import Path
from typing import Iterator, Dict, Any, Optional, Tuple
import logging
from tqdm import tqdm

logger = logging.getLogger(__name__)


class FileReader:
    """Efficient reader for large JSONL files with chunking support."""

    def __init__(self, file_path: Path, cache_manager=None):
        """Initialize file reader.

        Args:
            file_path: Path to the JSONL file
            cache_manager: Optional cache manager for line counts
        """
        self.file_path = Path(file_path)
        if not self.file_path.exists():
            raise FileNotFoundError(f"File not found: {file_path}")

        self.cache_manager = cache_manager
        self._line_count = None

    def get_line_count(self) -> int:
        """Get total line count of the file.

        Returns:
            Number of lines in the file
        """
        if self._line_count is not None:
            return self._line_count

        # Try to get from cache
        if self.cache_manager:
            cached_count = self.cache_manager.get_line_count(self.file_path)
            if cached_count is not None:
                self._line_count = cached_count
                return cached_count

        # Count lines
        logger.info(f"Counting lines in {self.file_path.name}...")
        count = 0
        with open(self.file_path, "r") as f:
            for _ in tqdm(f, desc="Counting lines"):
                count += 1

        self._line_count = count

        # Save to cache
        if self.cache_manager:
            self.cache_manager.save_line_count(self.file_path, count)

        logger.info(f"Total lines: {count:,}")
        return count

    def read_records(
        self, start: int = 0, end: Optional[int] = None, show_progress: bool = True
    ) -> Iterator[Dict[str, Any]]:
        """Read records from the file.

        Args:
            start: Starting line number (0-indexed)
            end: Ending line number (exclusive). None means read to end.
            show_progress: Whether to show progress bar

        Yields:
            Parsed JSON records
        """
        with open(self.file_path, "r") as f:
            # Skip to start line
            for _ in range(start):
                f.readline()

            # Calculate number of lines to read
            if end is None:
                total_lines = None
                desc = f"Reading from line {start:,}"
            else:
                total_lines = end - start
                desc = f"Reading lines {start:,}-{end:,}"

            # Read records
            line_iter = f
            if show_progress:
                line_iter = tqdm(line_iter, total=total_lines, desc=desc)

            for i, line in enumerate(line_iter, start=start):
                if end is not None and i >= end:
                    break

                line = line.strip()
                if not line:
                    continue

                try:
                    yield json.loads(line)
                except json.JSONDecodeError as e:
                    logger.warning(f"Failed to parse line {i+1}: {e}")
                    continue

    def read_chunks(
        self, chunk_size: int = 1_000_000, num_chunks: Optional[int] = None
    ) -> Iterator[Tuple[int, int]]:
        """Generate chunk boundaries for parallel processing.

        Args:
            chunk_size: Number of lines per chunk
            num_chunks: Optional number of chunks to create (overrides chunk_size)

        Yields:
            Tuples of (start_line, end_line) for each chunk
        """
        total_lines = self.get_line_count()

        if num_chunks is not None:
            # Divide into specified number of chunks
            chunk_size = (total_lines + num_chunks - 1) // num_chunks

        start = 0
        while start < total_lines:
            end = min(start + chunk_size, total_lines)
            yield (start, end)
            start = end

    def read_sample(self, n: int = 100, seed: Optional[int] = None) -> list[Dict[str, Any]]:
        """Read a random sample of records.

        Args:
            n: Number of records to sample
            seed: Random seed for reproducibility

        Returns:
            List of sampled records
        """
        import random

        total_lines = self.get_line_count()

        if n >= total_lines:
            # Read all records
            return list(self.read_records(show_progress=True))

        # Generate random line numbers
        if seed is not None:
            random.seed(seed)

        line_numbers = sorted(random.sample(range(total_lines), n))

        records = []
        with open(self.file_path, "r") as f:
            current_line = 0
            line_idx = 0

            for line in tqdm(f, total=total_lines, desc=f"Sampling {n} records"):
                if line_idx < len(line_numbers) and current_line == line_numbers[line_idx]:
                    line = line.strip()
                    if line:
                        try:
                            records.append(json.loads(line))
                            line_idx += 1
                        except json.JSONDecodeError:
                            pass

                current_line += 1

                if line_idx >= len(line_numbers):
                    break

        return records

    def extract_ids(self, show_progress: bool = True) -> set:
        """Extract all IDs from the file.

        Args:
            show_progress: Whether to show progress bar

        Returns:
            Set of all IDs in the file
        """
        ids = set()

        for record in self.read_records(show_progress=show_progress):
            if "id" in record:
                ids.add(record["id"])

        return ids