"""Async data loading operations."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence

from ..io.records import AsyncRecords
from ..table.schema import ColumnDef
from .readers.async_csv_reader import read_csv, read_csv_stream
from .readers.async_json_reader import (
    read_json,
    read_json_stream,
    read_jsonl,
    read_jsonl_stream,
)
from .readers.async_text_reader import read_text, read_text_stream


# Lazy import for parquet (optional dependency)
def _get_parquet_readers() -> tuple[Any, Any]:
    """Lazy import for parquet readers."""
    try:
        from .readers.async_parquet_reader import read_parquet, read_parquet_stream

        return read_parquet, read_parquet_stream
    except ImportError:
        return None, None


if TYPE_CHECKING:
    from ..table.async_table import AsyncDatabase


class AsyncDataLoader:
    """Builder for loading data from files and tables as AsyncRecords."""

    def __init__(self, database: "AsyncDatabase"):
        self._database = database
        self._schema: Optional[Sequence[ColumnDef]] = None
        self._options: Dict[str, object] = {}

    def stream(self, enabled: bool = True) -> "AsyncDataLoader":
        """Enable or disable streaming mode (chunked reading for large files)."""
        self._options["stream"] = enabled
        return self

    def schema(self, schema: Sequence[ColumnDef]) -> "AsyncDataLoader":
        """Set an explicit schema for the data source."""
        self._schema = schema
        return self

    def option(self, key: str, value: object) -> "AsyncDataLoader":
        """Set a read option (e.g., header=True for CSV, multiline=True for JSON)."""
        self._options[key] = value
        return self

    async def table(self, name: str) -> AsyncRecords:
        """Read from a database table as AsyncRecords.

        Note: For SQL operations on tables, use await db.table(name).select() instead.
        This method materializes the table data into AsyncRecords.
        """
        table_handle = await self._database.table(name)
        df = table_handle.select()
        rows = await df.collect()
        if not isinstance(rows, list):
            raise TypeError("table() requires collect() to return a list, not an iterator")
        return AsyncRecords(_data=rows, _database=self._database)

    async def csv(self, path: str) -> AsyncRecords:
        """Read a CSV file asynchronously.

        Args:
            path: Path to the CSV file

        Returns:
            AsyncRecords containing the CSV data
        """
        stream = self._options.get("stream", False)
        if stream:
            return await read_csv_stream(path, self._database, self._schema, self._options)
        return await read_csv(path, self._database, self._schema, self._options)

    async def json(self, path: str) -> AsyncRecords:
        """Read a JSON file (array of objects) asynchronously.

        Args:
            path: Path to the JSON file

        Returns:
            AsyncRecords containing the JSON data
        """
        stream = self._options.get("stream", False)
        if stream:
            return await read_json_stream(path, self._database, self._schema, self._options)
        return await read_json(path, self._database, self._schema, self._options)

    async def jsonl(self, path: str) -> AsyncRecords:
        """Read a JSONL file (one JSON object per line) asynchronously.

        Args:
            path: Path to the JSONL file

        Returns:
            AsyncRecords containing the JSONL data
        """
        stream = self._options.get("stream", False)
        if stream:
            return await read_jsonl_stream(path, self._database, self._schema, self._options)
        # For non-streaming JSONL, use the regular read_jsonl

        return await read_jsonl(path, self._database, self._schema, self._options)

    async def parquet(self, path: str) -> AsyncRecords:
        """Read a Parquet file asynchronously.

        Args:
            path: Path to the Parquet file

        Returns:
            AsyncRecords containing the Parquet data

        Raises:
            RuntimeError: If pandas or pyarrow are not installed
        """
        read_parquet, read_parquet_stream = _get_parquet_readers()
        if read_parquet is None:
            raise ImportError("Parquet support requires pyarrow. Install with: pip install pyarrow")
        stream = self._options.get("stream", False)
        if stream:
            return await read_parquet_stream(path, self._database, self._schema, self._options)  # type: ignore[no-any-return]
        return await read_parquet(path, self._database, self._schema, self._options)  # type: ignore[no-any-return]

    async def text(self, path: str, column_name: str = "value") -> AsyncRecords:
        """Read a text file as a single column (one line per row) asynchronously.

        Args:
            path: Path to the text file
            column_name: Name of the column to create (default: "value")

        Returns:
            AsyncRecords containing the text file lines
        """
        stream = self._options.get("stream", False)
        if stream:
            return await read_text_stream(
                path, self._database, self._schema, self._options, column_name
            )
        return await read_text(path, self._database, self._schema, self._options, column_name)

    async def format(self, source: str) -> "AsyncFormatReader":
        """Specify the data source format.

        Args:
            source: Format name (e.g., "csv", "json", "parquet")

        Returns:
            AsyncFormatReader for the specified format
        """
        return AsyncFormatReader(self, source)


class AsyncFormatReader:
    """Builder for format-specific async reads."""

    def __init__(self, reader: AsyncDataLoader, source: str):
        self._reader = reader
        self._source = source.lower()

    async def load(self, path: str) -> AsyncRecords:
        """Load data from the specified path using the configured format.

        Args:
            path: Path to the data file

        Returns:
            AsyncRecords containing the data

        Raises:
            ValueError: If format is unsupported
        """
        if self._source == "csv":
            return await self._reader.csv(path)
        elif self._source == "json":
            return await self._reader.json(path)
        elif self._source == "jsonl":
            return await self._reader.jsonl(path)
        elif self._source == "parquet":
            read_parquet, _ = _get_parquet_readers()
            if read_parquet is None:
                raise ImportError(
                    "Parquet support requires pyarrow. Install with: pip install pyarrow"
                )
            return await self._reader.parquet(path)
        elif self._source == "text":
            return await self._reader.text(path)
        else:
            raise ValueError(f"Unsupported format: {self._source}")
