"""Sorted Parquet writer for benchmark data.

Reads TBL (pipe-delimited) files generated by TPC tools and writes them
as sorted Parquet files. Sorting is done via PyArrow's compute kernels
to support large datasets without loading everything into memory at once.
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import Any

import pyarrow as pa
import pyarrow.compute as pc
import pyarrow.csv as csv
import pyarrow.parquet as pq

from benchbox.core.data_organization.config import (
    DataOrganizationConfig,
    SortColumn,
    SortOrder,
)

logger = logging.getLogger(__name__)


class SortedParquetWriter:
    """Writes TBL benchmark data as sorted Parquet files.

    Reads pipe-delimited TBL files (as produced by dbgen/dsdgen), sorts
    them by the configured columns, and writes Parquet output with controlled
    row group sizes for zone map experiments.
    """

    def __init__(self, config: DataOrganizationConfig, schema_registry: dict[str, Any] | None = None) -> None:
        """Initialize the writer.

        Args:
            config: Data organization configuration specifying sort columns,
                row group size, and compression.
            schema_registry: Optional mapping of table names to schema dicts
                (with 'columns' key). Used for column name resolution when
                reading TBL files. If not provided, column names are inferred
                as col_0, col_1, etc.
        """
        self.config = config
        self.schema_registry = schema_registry or {}

    def write_sorted_parquet(
        self,
        table_name: str,
        source_files: list[Path],
        output_dir: Path,
    ) -> Path:
        """Read TBL files, sort, and write as Parquet.

        Args:
            table_name: Benchmark table name (e.g., 'lineitem').
            source_files: One or more TBL source files (may be sharded chunks).
            output_dir: Directory to write the output Parquet file.

        Returns:
            Path to the written Parquet file.

        Raises:
            FileNotFoundError: If any source file does not exist.
            ValueError: If sort columns reference non-existent columns.
        """
        for f in source_files:
            if not f.exists():
                raise FileNotFoundError(f"Source file not found: {f}")

        sort_columns = self.config.get_sort_columns_for_table(table_name)
        column_names = self._get_column_names(table_name)

        # Read all source TBL files into a single table
        arrow_table = self._read_tbl_files(source_files, column_names)

        # Sort if columns are specified
        if sort_columns:
            arrow_table = self._sort_table(arrow_table, sort_columns)

        # Write Parquet
        output_path = output_dir / f"{table_name}.parquet"
        output_dir.mkdir(parents=True, exist_ok=True)

        compression_map = {
            "snappy": "SNAPPY",
            "gzip": "GZIP",
            "zstd": "ZSTD",
            "none": None,
        }
        compression = compression_map.get(self.config.compression)

        pq.write_table(
            arrow_table,
            output_path,
            compression=compression,
            use_dictionary=True,
            write_statistics=True,
            row_group_size=self.config.row_group_size,
        )

        row_count = arrow_table.num_rows
        file_size = output_path.stat().st_size
        logger.info(
            "Wrote sorted Parquet: %s (%s rows, %.1f MB, row_group_size=%d)",
            output_path.name,
            f"{row_count:,}",
            file_size / (1024 * 1024),
            self.config.row_group_size,
        )

        return output_path

    def _get_column_names(self, table_name: str) -> list[str] | None:
        """Get column names for a table from the schema registry.

        Returns:
            List of column names, or None if schema not available.
        """
        schema = self.schema_registry.get(table_name)
        if schema and "columns" in schema:
            return [col["name"] for col in schema["columns"]]
        return None

    def _read_tbl_files(
        self,
        source_files: list[Path],
        column_names: list[str] | None,
    ) -> pa.Table:
        """Read one or more TBL files into a PyArrow table.

        TBL format: pipe-delimited, no header, trailing delimiter on each line.
        """
        read_options = csv.ReadOptions(
            autogenerate_column_names=column_names is None,
            column_names=column_names,
        )
        parse_options = csv.ParseOptions(delimiter="|")

        tables = []
        for source_file in source_files:
            table = csv.read_csv(source_file, read_options=read_options, parse_options=parse_options)
            tables.append(table)

        combined = pa.concat_tables(tables) if len(tables) > 1 else tables[0]

        # TBL files have a trailing pipe which creates an extra empty column.
        # Drop it if present.
        last_col = combined.column_names[-1]
        if last_col.startswith("col_") or last_col == "" or last_col.startswith("f"):
            # Check if the last column is all null/empty (trailing delimiter artifact)
            last_values = combined.column(last_col)
            if last_values.null_count == len(last_values) or (
                pa.types.is_string(last_values.type) and pc.all(pc.equal(last_values, "")).as_py()
            ):
                combined = combined.drop(last_col)

        return combined

    def _sort_table(self, table: pa.Table, sort_columns: list[SortColumn]) -> pa.Table:
        """Sort a PyArrow table by the specified columns.

        Args:
            table: Table to sort.
            sort_columns: Columns and directions to sort by.

        Returns:
            Sorted table.

        Raises:
            ValueError: If a sort column is not found in the table.
        """
        available = set(table.column_names)
        for sc in sort_columns:
            if sc.name not in available:
                raise ValueError(f"Sort column '{sc.name}' not found in table. Available columns: {sorted(available)}")

        sort_keys = [(sc.name, "ascending" if sc.order == SortOrder.ASC else "descending") for sc in sort_columns]
        indices = pc.sort_indices(table, sort_keys=sort_keys)
        return table.take(indices)
