"""Base classes for PyArrow-based writers (Parquet, ORC)."""

import logging
import math
import os
import pathlib
from abc import abstractmethod
from collections.abc import Callable
from pathlib import Path
from typing import Any, BinaryIO, TextIO

from transmog.exceptions import MissingDependencyError, OutputError
from transmog.writers.base import DataWriter, StreamingWriter, _collect_field_names

logger = logging.getLogger(__name__)

try:
    import pyarrow as pa

    PYARROW_AVAILABLE = True
    _ARROW_WRITE_ERRORS: tuple[type[Exception], ...] = (OSError, pa.lib.ArrowException)
except ImportError:
    pa = None
    PYARROW_AVAILABLE = False
    _ARROW_WRITE_ERRORS: tuple[type[Exception], ...] = (OSError,)  # type: ignore[no-redef]


def _is_valid_float_for_inference(value: Any) -> bool:
    """Check if a float value is valid for type inference.

    NaN and Infinity values should not be used for type inference as they
    don't represent typical float data patterns.

    Args:
        value: Value to check

    Returns:
        True if value is a valid float for inference purposes
    """
    if not isinstance(value, float):
        return False
    return not (math.isnan(value) or math.isinf(value))


def _convert_bool(value: Any) -> bool | None:
    """Convert a value to bool for Arrow column."""
    try:
        return bool(value)
    except (ValueError, TypeError):
        return None


def _convert_int(value: Any) -> int | None:
    """Convert a value to int for Arrow column."""
    try:
        return int(value)
    except (ValueError, TypeError):
        return None


def _convert_float(value: Any) -> float | None:
    """Convert a value to float for Arrow column."""
    try:
        return float(value)
    except (ValueError, TypeError):
        return None


def _convert_str(value: Any) -> str | None:
    """Convert a value to str for Arrow column."""
    try:
        return str(value)
    except (ValueError, TypeError):
        return None


_TYPE_CONVERTERS: dict[Any, Callable] = {}


def _get_type_converters() -> dict[Any, Callable]:
    """Get the mapping of PyArrow types to converter functions.

    Lazily initialized to avoid import-time dependency on PyArrow.
    """
    if not _TYPE_CONVERTERS and pa is not None:
        _TYPE_CONVERTERS[pa.bool_()] = _convert_bool
        _TYPE_CONVERTERS[pa.int64()] = _convert_int
        _TYPE_CONVERTERS[pa.float64()] = _convert_float
    return _TYPE_CONVERTERS


class PyArrowWriter(DataWriter):
    """Base writer for PyArrow-based formats (Parquet, ORC)."""

    def __init__(self, compression: str, **options: Any) -> None:
        """Initialize the PyArrow writer.

        Args:
            compression: Compression format
            **options: Additional writer options
        """
        self.compression = compression
        self.options = options

    @abstractmethod
    def _get_format_name(self) -> str:
        """Get the format name for error messages."""
        pass

    @abstractmethod
    def _write_table(
        self,
        table: Any,
        destination: str | BinaryIO | TextIO,
        compression: str,
        **options: Any,
    ) -> None:
        """Write a PyArrow table to destination."""
        pass

    def write(
        self,
        data: list[dict[str, Any]],
        destination: str | BinaryIO | TextIO,
        **options: Any,
    ) -> str | BinaryIO | TextIO:
        """Write data to a file.

        Args:
            data: Data to write
            destination: Path or file-like object to write to
            **options: Format-specific options (compression, etc.)

        Returns:
            Path to the written file or file-like object

        Raises:
            MissingDependencyError: If PyArrow is not available
            OutputError: If writing fails
        """
        if pa is None:
            format_name = self._get_format_name()
            raise MissingDependencyError(
                f"PyArrow is required for {format_name} support. "
                "Install with: pip install pyarrow"
            )

        try:
            compression_val = options.get("compression", self.compression)

            if not data:
                return destination

            field_names = _collect_field_names(data)

            columns: dict[str, list[Any]] = {field: [] for field in field_names}
            for record in data:
                for field in field_names:
                    columns[field].append(record.get(field))

            table = pa.table(columns)

            if isinstance(destination, (str, pathlib.Path)):
                path = pathlib.Path(destination)
                path.parent.mkdir(parents=True, exist_ok=True)
                self._write_table(table, str(path), compression_val, **options)
                return str(path) if isinstance(destination, str) else path
            else:
                mode = getattr(destination, "mode", "")
                if mode and "b" not in mode:
                    format_name = self._get_format_name()
                    raise OutputError(
                        f"{format_name} format requires binary streams, "
                        "text streams not supported"
                    )
                self._write_table(table, destination, compression_val, **options)
                return destination
        except _ARROW_WRITE_ERRORS as exc:
            format_name = self._get_format_name()
            raise OutputError(f"Failed to write {format_name} file: {exc}") from exc


class PyArrowStreamingWriter(StreamingWriter):
    """Base streaming writer for PyArrow-based formats."""

    def __init__(
        self,
        destination: str | BinaryIO | TextIO | None = None,
        entity_name: str = "entity",
        compression: str = "snappy",
        batch_size: int = 10000,
        stringify_mode: bool = False,
        **options: Any,
    ) -> None:
        """Initialize the PyArrow streaming writer.

        Args:
            destination: Path or file-like object to write to
            entity_name: Name of the entity for output files
            compression: Compression algorithm
            batch_size: Number of records per batch
            stringify_mode: If True, all fields are strings (skip type inference)
            **options: Additional options for PyArrow
        """
        super().__init__(destination, entity_name, **options)

        if pa is None:
            format_name = self._get_format_name()
            raise MissingDependencyError(
                f"PyArrow is required for {format_name} streaming support. "
                "Install with: pip install pyarrow"
            )

        self.compression = compression
        self.batch_size = batch_size
        self.stringify_mode = stringify_mode
        self.writers: dict[str, Any] = {}
        self.schemas: dict[str, Any] = {}
        self.converters: dict[str, dict[str, Callable]] = {}
        self.buffers: dict[str, list[dict[str, Any]]] = {}
        self.base_dir: str | None = None
        self._column_buffers: dict[str, dict[str, list[Any]]] = {}
        self.file_paths: dict[str, str] = {}

        if isinstance(destination, str):
            self.base_dir = destination
            os.makedirs(self.base_dir, exist_ok=True)

    @abstractmethod
    def _get_format_name(self) -> str:
        """Get the format name."""
        pass

    @abstractmethod
    def _get_file_extension(self) -> str:
        """Get the file extension (e.g., '.parquet', '.orc')."""
        pass

    @abstractmethod
    def _create_writer(self, file_path: str, schema: Any) -> Any:
        """Create the format-specific writer instance."""
        pass

    @abstractmethod
    def _write_to_writer(self, writer: Any, table: Any) -> None:
        """Write table using the format-specific writer."""
        pass

    def _get_table_path(self, table_name: str) -> str | None:
        """Get the file path for a table.

        Args:
            table_name: Name of the table

        Returns:
            File path or None if using file-like object
        """
        if self.base_dir is None:
            return None

        if table_name in self.file_paths:
            return self.file_paths[table_name]

        extension = self._get_file_extension()
        if table_name == "main":
            file_path = os.path.join(self.base_dir, f"{self.entity_name}{extension}")
        else:
            safe_name = table_name.replace(".", "_").replace("/", "_")
            file_path = os.path.join(self.base_dir, f"{safe_name}{extension}")

        self.file_paths[table_name] = file_path
        return file_path

    def _create_schema(
        self, records: list[dict[str, Any]], stringify_mode: bool = False
    ) -> tuple[Any, dict[str, Callable]]:
        """Create PyArrow schema and field converters from records.

        Handles special float values (NaN, Inf) by skipping them during type
        inference, as they don't represent typical data patterns.

        Args:
            records: Records to infer schema from
            stringify_mode: If True, all fields are strings (skip type inference)

        Returns:
            Tuple of (PyArrow schema, dict mapping field names to converters)
        """
        if not records:
            return pa.schema([]), {}

        field_names = _collect_field_names(records)
        type_converters = _get_type_converters()

        # If stringify mode, all fields are strings - skip type inference
        if stringify_mode:
            fields = [pa.field(key, pa.string()) for key in field_names]
            converters = dict.fromkeys(field_names, _convert_str)
            logger.debug(
                "arrow schema created (stringify mode), fields=%d", len(fields)
            )
            return pa.schema(fields), converters

        fields = []
        converters = {}

        for key in field_names:
            value = None
            found_float = False

            for record in records:
                if key not in record:
                    continue

                val = record[key]
                if val is None:
                    continue

                # For floats, skip NaN/Inf for type inference but note that
                # the field contains float data
                if isinstance(val, float):
                    found_float = True
                    if _is_valid_float_for_inference(val):
                        value = val
                        break
                else:
                    value = val
                    break

            # Determine type based on found value or float flag
            if value is None and found_float:
                pa_type = pa.float64()
            elif value is None:
                pa_type = pa.string()
            elif isinstance(value, bool):
                pa_type = pa.bool_()
            elif isinstance(value, int):
                pa_type = pa.int64()
            elif isinstance(value, float):
                pa_type = pa.float64()
            else:
                pa_type = pa.string()

            fields.append(pa.field(key, pa_type))
            converters[key] = type_converters.get(pa_type, _convert_str)

        types = {f.name: str(f.type) for f in fields}
        logger.debug("arrow schema created, fields=%d, types=%s", len(fields), types)
        return pa.schema(fields), converters

    def _records_to_table(self, records: list[dict[str, Any]], table_name: str) -> Any:
        """Convert records to PyArrow table.

        Args:
            records: Records to convert
            table_name: Name of the table

        Returns:
            PyArrow table
        """
        if not records:
            return pa.table({})

        if table_name not in self.schemas:
            schema, converters = self._create_schema(
                records, stringify_mode=self.stringify_mode
            )
            self.schemas[table_name] = schema
            self.converters[table_name] = converters

        schema = self.schemas[table_name]
        converters = self.converters[table_name]

        if table_name in self._column_buffers:
            columns = self._column_buffers[table_name]
            for col in columns.values():
                col.clear()
        else:
            columns = {field.name: [] for field in schema}
            self._column_buffers[table_name] = columns

        for record in records:
            for field in schema:
                field_name = field.name
                value = record.get(field_name)

                if value is None:
                    columns[field_name].append(None)
                else:
                    columns[field_name].append(converters[field_name](value))

        arrays = [pa.array(columns[field.name], type=field.type) for field in schema]
        return pa.table(arrays, schema=schema)

    def _initialize_writer(self, table_name: str, schema: Any) -> None:
        """Initialize a writer for a table.

        Args:
            table_name: Name of the table
            schema: PyArrow schema
        """
        file_path = self._get_table_path(table_name)

        if file_path is None:
            return

        writer = self._create_writer(file_path, schema)
        self.writers[table_name] = writer
        self.schemas[table_name] = schema

    def _write_buffer(self, table_name: str) -> None:
        """Write buffered records to file.

        Args:
            table_name: Name of the table
        """
        if table_name not in self.buffers or not self.buffers[table_name]:
            return

        records = self.buffers[table_name]
        table = self._records_to_table(records, table_name)

        if table_name not in self.writers:
            self._initialize_writer(table_name, table.schema)

        writer = self.writers.get(table_name)
        if writer:
            self._write_to_writer(writer, table)

        self.buffers[table_name].clear()

    def write_main_records(self, records: list[dict[str, Any]]) -> None:
        """Write a batch of main records.

        Args:
            records: List of main table records to write
        """
        if not records:
            return

        table_name = "main"

        if table_name not in self.buffers:
            self.buffers[table_name] = []

        self.buffers[table_name].extend(records)

        if len(self.buffers[table_name]) >= self.batch_size:
            self._write_buffer(table_name)

    def write_child_records(
        self, table_name: str, records: list[dict[str, Any]]
    ) -> None:
        """Write a batch of child records.

        Args:
            table_name: Name of the child table
            records: List of child records to write
        """
        if not records:
            return

        if table_name not in self.buffers:
            self.buffers[table_name] = []

        self.buffers[table_name].extend(records)

        if len(self.buffers[table_name]) >= self.batch_size:
            self._write_buffer(table_name)

    def close(self) -> list[Path]:
        """Finalize output, flush buffered data, and clean up resources.

        Returns:
            List of file paths written, or empty list if writing to a stream.
        """
        if getattr(self, "_closed", False):
            return []

        for table_name in list(self.buffers.keys()):
            if self.buffers[table_name]:
                self._write_buffer(table_name)

        paths = [Path(p) for p in self.file_paths.values()]

        for writer in self.writers.values():
            if hasattr(writer, "close"):
                writer.close()

        self.writers.clear()
        self.schemas.clear()
        self.converters.clear()
        self.buffers.clear()
        self._column_buffers.clear()
        self._closed = True
        return paths


__all__ = ["PyArrowWriter", "PyArrowStreamingWriter", "PYARROW_AVAILABLE"]
