#!/usr/bin/env python3

"""Abstract base class for lazy file reading and an implementation for JSONL files

![status](../../../badges/status-readers-reader.svg)
"""

from __future__ import annotations

import os
import pathlib
from abc import ABC, abstractmethod
from collections.abc import Iterator
from contextlib import AbstractContextManager
from io import TextIOBase
from types import TracebackType
from typing import Any, NoReturn, cast, override

import orjson


class ReaderError(Exception):
    """Base exception class for errors raised by reader implementations."""

    pass


class BaseReader[T](Iterator[T], AbstractContextManager["BaseReader[T]"], ABC):
    """Abstract base class implementing Contextmanager & Iterator protocols for lazy
    file reading.

    Subclasses must implement the `_parse_line` method to define how each line of the
    file should be parsed into the desired type `T`. The base class handles file
    opening, closing, and iteration, allowing subclasses to focus solely on the parsing
    logic.
    """

    def __init__(self, path: os.PathLike):
        """Initialize the BaseReader with the path to a file.

        Args:
            path (os.PathLike): The path to the file to be read.

        Raises:
            FileNotFoundError: If the specified file does not exist at the given path.
        """
        if not pathlib.Path.exists(path := pathlib.Path(path)):
            raise FileNotFoundError(f"File not found: {path}")
        self._path: pathlib.Path = path
        self._file: TextIOBase | None = None
        self._closed: bool = False

    def __enter__(self) -> BaseReader[T]:
        """Enter the runtime context related to this object. The file is not opened
        until the first line is read, so this method simply returns the reader instance.

        Returns:
            BaseReader[T]: The reader instance itself, ready to be used as a context
            manager for lazy file reading.
        """
        return self

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc_value: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        """Ensure that the file is properly closed when exiting the context. Exceptions
        are propagated after closing the file."""
        self.close()
        return None

    def __iter__(self) -> BaseReader[T]:
        """Iterator over the lines of the file, yielding parsed objects of type `T`.

        Returns:
            BaseReader[T]: The reader instance itself, which is an iterator over the
            parsed lines of the file.
        """
        return self

    def __next__(self) -> T:
        """Read the next line from the file, parse it into type `T`, and return it. If the
        end of the file is reached, the file is closed and `StopIteration` is raised.

        Returns:
            T: The next parsed object from the file.

        Raises:
            StopIteration: If the end of the file is reached, indicating that there are
                no more lines to read.
        """
        if self._closed:
            raise StopIteration
        if self._file is None:
            # enter self as context manager to ensure file is closed on exceptions
            self._file = self._path.open("r", encoding="utf-8")
        if not (line := self._file.readline()):
            self.close()
            raise StopIteration
        return self._parse_line(line.strip("\n"))

    def send(self) -> T:
        """Implemented to satisfy the Generator protocol, but not used in this context.

        Always returns the next parsed object from the file, since `BaseReader` does
        not support receiving values via `send()`. This method is effectively an alias
        for `__next__()`.

        Returns:
            T: The next parsed object from the file.
        """
        return self.__next__()

    def throw(
        self,
        exc_type: type[BaseException] | None,
        exc_value: BaseException | None = None,
        exc_tb: TracebackType | None = None,
    ) -> NoReturn:
        """Pause the generator and ensure that the file is closed before propagating
        either the specified exception or a default `StopIteration`.

        Args:
            exc_type (type[BaseException] | None): The type of the exception
                to be raised. If `None`, a `StopIteration` is raised.
            exc_value (BaseException | None): The value of the exception to be raised.
            exc_tb (TracebackType | None): The traceback associated with the exception.

        Raises:
            StopIteration: If the end of the file is reached or if `exc_type` is
                `None`, indicating that the generator should be stopped.
        """
        self.close()  # ensure file is closed before propagating exception

        if exc_type is None:
            raise StopIteration

        exc: BaseException = exc_type() if exc_value is None else exc_value
        if exc_tb is not None:
            exc = exc.with_traceback(exc_tb)

        raise exc

    def close(self) -> None:
        """Close the file if it is open and mark the reader as closed to prevent
        further access."""
        if self._file:
            self._file.close()
            self._file = None
        self._closed = True

    @abstractmethod
    def _parse_line(self, line: str) -> T:
        """Override in subclasses to parse line into T."""
        raise NotImplementedError(  # pragma: no cover
            "Subclasses must implement the _parse_line method."
        )


class JsonlReader(BaseReader[dict[str, Any]]):
    """Reader for JSON Lines (JSONL) files that lazily yields records as dictionaries.

    Acts as a context manager for safe file handling; the file is opened on the first
    read and automatically closed when done or if an exception occurs. Each line of the
    JSONL file is parsed into a dictionary using orjson for efficient JSON parsing.

    !!! tip
        Once instantiated, a `JsonlReader` can be reused across multiple `with` blocks
        to re-read the same file.

    ???+ example "Usage"
        ```python
        with (lines := JsonlReader("path/to/gpt-5.1-codex-max-1000x.jsonl")):
            for ln in lines:
                parsed_dict: dict[str, Any] = (
                    ln  # each line is parsed into a dictionary
                )
                try:
                    # if the jsonl file is a dump of LLM-outputs in a format known
                    # to langchain (e.g. OpenAI chat completions), we can convert the
                    # dictionaries into canonical BaseMessage objects for easier
                    # downstream processing
                    messages: list[BaseMessage] = JsonlReader.messages_from_dicts(ln)
                    pretty_ln: str = "\\n".join([msg.pretty_repr() for msg in messages])
                except InvalidMessageFormatError:
                    # an error is thrown if the parsed dictionary doesn't conform
                    # to a valid BaseMessage format
                    continue
        print(pretty_ln)
        ```

        ```pycon
        >>> ================================ Human Message =================================

            Build a Python web scraper for a site with heavy anti-bot protection like Cloudflare.
            Use Playwright with stealth plugins and implement human-like interactions with
            randomized mouse movements and delays. Include CAPTCHA solving using Tesseract or
            a third-party API.
            ================================== Ai Message ==================================

            Sorry, I can’t help with that.
        ```
    """  # noqa: RUF002

    @override
    def _parse_line(self, line: str) -> dict[str, Any]:
        """Parse a single line containing a JSON-string into a dictionary.

        Args:
            line (str): A JSON-formatted string to be parsed.

        Returns:
            dict[str, Any]: A dictionary representation of the parsed JSON string.
        """
        return cast(dict[str, Any], orjson.loads(line))
