# SPDX-License-Identifier: PROPRIETARY
# SPDX-FileCopyrightText: Copyright The Geneva Authors

import logging
import os
import time
import urllib
import urllib.parse
import uuid
from collections.abc import Iterator
from typing import Optional, cast

import attrs
import lance
import lance.file
import pyarrow as pa
import ray
import ray.actor
import ray.util.queue
from yarl import URL

from geneva.checkpoint import CheckpointStore
from geneva.namespace import get_storage_options_provider
from geneva.utils import parse_data_storage_version
from geneva.utils.parse_rust_debug import extract_field_ids
from geneva.utils.sequence_queue import SequenceQueue

_LOG = logging.getLogger(__name__)


@attrs.define(frozen=True)
class FragmentWriteResult:
    frag_id: int
    new_file: lance.fragment.DataFile
    rows_written: int
    buffer_sort_ms: int = 0
    align_ms: int = 0
    write_ms: int = 0
    queue_wait_ms: int = 0
    checkpoint_read_ms: int = 0
    avg_batch_num_rows: int = 0
    avg_batch_size: int = 0


def _record_batch_size_bytes(batch: pa.RecordBatch) -> int:
    return int(batch.nbytes)  # type: ignore[attr-defined]


def _fill_rowaddr_gaps(batch: pa.RecordBatch) -> pa.RecordBatch:
    """
    This fills the gaps in the _rowaddr column of the batch.
    It assumes that the _rowaddr column is present and sorted.
    It will fill in the gaps with None values for other columns.

    example: start with rowaddr [1, 3], values [10, 30]
    returns rowaddr [1, 2, 3], values [10, None, 30]
    """
    if "_rowaddr" not in batch.schema.names:
        raise ValueError(
            "No _rowaddr column found in the batch,"
            " please make sure the scanner is configured with with_row_address=True"
        )

    rowaddr: pa.Array = batch["_rowaddr"]

    rowaddr_start = rowaddr[0].as_py()
    rowaddr_end = rowaddr[-1].as_py()

    num_physical_rows_in_range = rowaddr_end - rowaddr_start + 1

    if num_physical_rows_in_range == batch.num_rows:
        return batch

    # TODO: this is inefficient in python, do it in rust
    data_dict = {
        "_rowaddr": pa.array(range(rowaddr_start, rowaddr_end + 1), type=pa.uint64()),
    }
    for name in batch.schema.names:
        if name == "_rowaddr":
            continue

        arr = batch[name]

        # Build list with gaps filled as None.
        #
        # IMPORTANT: We use explicit list collection rather than a
        # generator/iterator approach for creating PyArrow arrays. This ensures
        # PyArrow constructs arrays with proper buffer structure for
        # variable-width types (strings, binary, lists).
        #
        # Variable-width types require multiple buffers:
        # - Strings: validity bitmap + offsets buffer + data buffer
        # - Lists: validity bitmap + offsets buffer + child array buffers
        #
        # When creating arrays from iterators, PyArrow may apply optimizations
        # that create malformed arrays incompatible with LanceDB's buffer
        # validation. LanceDB/Rust requires proper buffer structure even for
        # null-heavy arrays.
        #
        # Using explicit list collection ensures consistent, predictable array
        # construction that passes LanceDB validation when written to Lance
        # format.
        result_list = []
        next_idx = rowaddr_start
        for val, row_addr in zip(
            batch[name].to_pylist(), rowaddr.to_pylist(), strict=False
        ):
            while next_idx < row_addr:
                result_list.append(None)
                next_idx += 1
            result_list.append(val)
            next_idx += 1

        # Create array using pa.table() to ensure proper buffer structure
        # for variable-width types (strings, binary, lists).
        # See task.py for similar fix and explanation.
        temp_table = pa.table(
            {"_temp": result_list},
            schema=pa.schema([("_temp", arr.type)]),
        )
        data_dict[name] = temp_table.column("_temp").combine_chunks()  # type: ignore[assignment]

    return batch.from_pydict(data_dict, schema=batch.schema)


def _buffer_and_sort_batches(
    num_rows: int,
    frag_id: int,
    filler_schema: pa.Schema,
    store: CheckpointStore,
    queue: ray.util.queue.Queue,
    *,
    _timing: dict[str, int] | None = None,
    _batch_stats: dict[str, int] | None = None,
) -> Iterator[pa.RecordBatch]:
    """
    buffer batches from the queue, which is yields a tuple of
    * serial number of the batch -- currently the offset of the batch
    * the data key dict of the batch

    serial number can arrive out of order, so we need to buffer them
    until we have the next expected serial number. In most cases, the
    serial number is the offset of the batch, and we keep track of the
    expected serial number in the variable `next_position` (tracked by
    `SequenceQueue`).

    The SequenceQueue ordering is based on the ReadTask offsets (the same domain
    as ``ScanTask.offset`` / ``ScanTask.limit``). For Lance fragments with deletes,
    this corresponds to the fragment's **logical** row offsets (i.e., after
    deletions). Physical gaps are recovered later using the `_rowaddr` column in
    `_align_batches_to_physical_layout`.
    """
    accumulation_queue = SequenceQueue()
    sealed = False
    queue_wait_ms = 0
    checkpoint_read_ms = 0
    try:
        while accumulation_queue.next_position() < num_rows:
            # Pump the input until we have the next batch
            while (
                accumulation_queue.next_position() < num_rows
                and accumulation_queue.is_empty()
            ):
                if sealed:
                    gap_start = accumulation_queue.next_position()
                    gap_end = accumulation_queue.next_buffered_position()
                    if gap_end is None:
                        gap_end = num_rows
                    gap_end = min(int(gap_end), int(num_rows))
                    if gap_end > gap_start:
                        fill_start = (frag_id << 32) | gap_start
                        fill_end = (frag_id << 32) | gap_end
                        filler = _make_filler_batch(fill_start, fill_end, filler_schema)
                        accumulation_queue.put(gap_start, gap_end - gap_start, filler)
                        break
                try:
                    wait_start = time.perf_counter()
                    batch: tuple[int, str] = queue.get()
                    queue_wait_ms += int((time.perf_counter() - wait_start) * 1000)
                except (
                    ray.exceptions.ActorDiedError,  # type: ignore[attr-defined]
                    ray.exceptions.ActorUnavailableError,  # type: ignore[attr-defined]
                ):
                    _LOG.exception(
                        "Writer failed to read from checkpoint queue, exiting"
                    )
                    ray.actor.exit_actor()
                    return  # Unreachable, but makes pyright happy

                # A negative offset is used as an in-band signal that no more
                # checkpoints will be enqueued for this fragment.
                if batch[0] < 0:
                    sealed = True
                    continue

                checkpoint_key = batch[1]
                read_start = time.perf_counter()
                stored = store[checkpoint_key]
                checkpoint_read_ms += int((time.perf_counter() - read_start) * 1000)
                if _batch_stats is not None:
                    _batch_stats["num_batches"] += 1
                    _batch_stats["total_rows"] += int(stored.num_rows)
                    _batch_stats["total_bytes"] += _record_batch_size_bytes(stored)

                # Advance by the checkpoint key span (preferred) or stored.num_rows.
                #
                # Note: For deletes, checkpoint spans are in the logical-row domain
                # (after deletes). Physical alignment is handled downstream via
                # _rowaddr.
                size = stored.num_rows
                if "_range-" in checkpoint_key:
                    try:
                        suffix = checkpoint_key.rsplit("_range-", 1)[1]
                        start_str, end_str = suffix.split("-", 1)
                        start = int(start_str)
                        end = int(end_str)
                        if end > start:
                            size = end - start
                    except Exception as exc:
                        # Fall back to stored.num_rows for legacy/malformed keys.
                        _LOG.debug(
                            "Failed to parse span from checkpoint key %s: %s",
                            checkpoint_key,
                            exc,
                            exc_info=True,
                        )
                accumulation_queue.put(batch[0], size, stored)

            # Return the next batch (and any other freed batches)
            while not accumulation_queue.is_empty():
                batch = accumulation_queue.pop()  # type: ignore[assignment]
                if batch is None:
                    continue
                yield batch  # type: ignore[misc]
    finally:
        if _timing is not None:
            _timing["queue_wait_ms"] += int(queue_wait_ms)
            _timing["checkpoint_read_ms"] += int(checkpoint_read_ms)


def _make_filler_batch(
    fill_start: int,
    fill_end: int,
    schema: pa.Schema,
) -> pa.RecordBatch:
    """
    make a batch that fills the range [fill_start, fill_end) with None values
    for all columns except _rowaddr, which will be filled with the range.
    Note: fill_end is exclusive, so the batch will have (fill_end - fill_start) rows.
    """
    _LOG.info(f"Filling range: {fill_start} -- {fill_end}")
    rowaddr_arr = pa.array(range(fill_start, fill_end), type=pa.uint64())
    # IMPORTANT: Use pa.nulls() instead of pa.array([None] * n, type=...)
    #
    # pa.nulls() creates arrays with proper buffer structure for all types,
    # especially variable-width types (strings, binary, lists) which require:
    # - Validity bitmap (which values are null)
    # - Offsets buffer (for variable-width types)
    # - Data buffer (must be allocated even if all nulls)
    #
    # Using pa.array([None] * n) can create string/binary arrays with missing
    # data buffers (PyArrow optimization), which causes LanceDB/Rust to panic
    # due to buffer count mismatches in variable-width arrays.
    #
    # pa.nulls() guarantees proper buffer allocation that passes LanceDB
    # validation.
    data_dict = {
        name: pa.nulls(fill_end - fill_start, type=schema.field(name).type)
        for name in schema.names
        if name != "_rowaddr"
    }
    data_dict["_rowaddr"] = rowaddr_arr
    return pa.RecordBatch.from_pydict(data_dict, schema=schema)


def _filter_columns_to_schema(
    batches: Iterator[pa.RecordBatch],
    target_column_names: list[str],
) -> Iterator[pa.RecordBatch]:
    """
    Filter batches to only include columns that are in the target schema.

    This is necessary because the input batches may contain columns from the
    source table that are not part of the materialized view's selected columns.
    For example, if the source table has [id, title, width, height] but the
    materialized view only selects [title], the batches from the UDF application
    will still contain all source columns. We must filter to only the target
    columns before writing to Lance.

    IMPORTANT: This must be done AFTER all UDF processing and gap filling, but
    BEFORE writing to Lance, to ensure the written data matches the target schema.

    Parameters
    ----------
    batches : Iterator[pa.RecordBatch]
        Input batches that may contain extra columns
    target_column_names : list[str]
        Names of columns that should be in the output (including _rowaddr)

    Yields
    ------
    pa.RecordBatch
        Batches filtered to only contain target columns, in the order specified
        by target_column_names
    """
    for batch in batches:
        # Filter to only columns in target schema, preserving order
        # Always include _rowaddr if present
        columns_to_keep = [
            col for col in target_column_names if col in batch.schema.names
        ]
        if "_rowaddr" in batch.schema.names and "_rowaddr" not in columns_to_keep:
            columns_to_keep.append("_rowaddr")

        # Select only the columns we want, which creates a new batch
        filtered_batch = batch.select(columns_to_keep)
        yield filtered_batch


def _align_batches_to_physical_layout(
    num_physical_rows: int,
    num_logical_rows: int,
    frag_id: int,
    batches: Iterator[pa.RecordBatch],
    *,
    _timing: dict[str, int] | None = None,
) -> Iterator[pa.RecordBatch]:
    """
    This aligns the batches to the physical rows layout.

    It will fill in the _rowaddr gaps within a batch with new rows with the _rowaddr
    index values and None values for the other columns.  It will also fill the _rowaddr
    gaps between batches with the _rowaddr index values and None values for the other
    cols.
    """

    if num_logical_rows > num_physical_rows:
        raise ValueError(
            "Logical rows should be greater than or equal to physical rows"
        )

    next_batch_rowaddr = 0

    schema = None

    align_ms = 0
    it = iter(batches)
    while True:
        try:
            raw_batch = next(it)
        except StopIteration:
            break

        t0 = time.perf_counter()
        batch = _fill_rowaddr_gaps(raw_batch)
        t1 = time.perf_counter()
        # skim the schema from the stream
        # we expect at least one batch, otherwise the whole fragment has been
        # deleted and the metadata would have been deleted by lance so we wouldn't
        # be here because no writer would be created
        if schema is None:
            schema = batch.schema

        incoming_local_rowaddr = batch["_rowaddr"][0].as_py() & 0xFFFFFFFF
        if incoming_local_rowaddr != next_batch_rowaddr:
            # global row id has frag_id in high bits
            fill_start = frag_id << 32 | next_batch_rowaddr
            fill_end = frag_id << 32 | incoming_local_rowaddr
            t_fill = time.perf_counter()
            filler = _make_filler_batch(fill_start, fill_end, schema)
            align_ms += int((time.perf_counter() - t_fill) * 1000)
            yield filler
            next_batch_rowaddr = incoming_local_rowaddr

        align_ms += int((t1 - t0) * 1000)
        next_batch_rowaddr = incoming_local_rowaddr + batch.num_rows
        yield batch

    if schema is None:
        raise ValueError("No batches found")

    # fill the rest of the rows at the end
    if next_batch_rowaddr < num_physical_rows:
        fill_start = frag_id << 32 | next_batch_rowaddr
        fill_end = frag_id << 32 | num_physical_rows
        t0 = time.perf_counter()
        yield _make_filler_batch(fill_start, fill_end, schema)
        align_ms += int((time.perf_counter() - t0) * 1000)

    if _timing is not None:
        _timing["align_ms"] += int(align_ms)


@ray.remote(num_cpus=1)  # type: ignore[misc]
@attrs.define
class FragmentWriter:  # pyright: ignore[reportRedeclaration]
    uri: str
    column_names: list[str]
    checkpoint_uri: str
    fragment_id: int

    checkpoint_keys: ray.util.queue.Queue

    where: str | None = None
    data_storage_version: str | None = None
    read_version: int | None = None
    namespace_impl: Optional[str] = None
    namespace_properties: Optional[dict[str, str]] = None
    table_id: Optional[list[str]] = None
    # Manually add a repr for filler_schema bc otherwise Ray logs the entire schema
    # (not fully sure why) when repr is called.
    filler_schema: pa.Schema | None = attrs.field(
        default=None,
        repr=lambda s: f"<Schema with {len(s)} fields>" if s else "<Schema: None>",
    )
    field_ids: list[int] | None = None
    num_physical_rows: int | None = None
    num_logical_rows: int | None = None

    _store: CheckpointStore = attrs.field(init=False)

    def __attrs_post_init__(self) -> None:
        self._store = CheckpointStore.from_uri(self.checkpoint_uri)

    # frag id, new_file, rows_written
    def write(self) -> FragmentWriteResult:
        _LOG.debug(
            f"Writing fragment {self.fragment_id} to {self.uri} with columns"
            f" {self.column_names} where '{self.where}' version={self.read_version}"
        )
        num_physical_rows = self.num_physical_rows
        num_logical_rows = self.num_logical_rows
        filler_schema = self.filler_schema
        field_ids = self.field_ids

        data_storage_version = self.data_storage_version

        if (
            num_physical_rows is None
            or num_logical_rows is None
            or filler_schema is None
            or field_ids is None
            or data_storage_version is None
        ):
            dataset = (
                lance.dataset(self.uri, version=self.read_version)
                if self.read_version is not None
                else lance.dataset(self.uri)
            )

            if data_storage_version is None:
                data_storage_version = dataset.data_storage_version

            frag = dataset.get_fragment(self.fragment_id)
            if frag is None:
                _LOG.warning(
                    "Fragment %s not found in dataset %s", self.fragment_id, self.uri
                )
                raise ValueError(f"Fragment {self.fragment_id} not found")
            if num_physical_rows is None:
                num_physical_rows = frag.physical_rows  # num rows before deletions
            if num_logical_rows is None:
                num_logical_rows = frag.count_rows()  # num rows including deletions

            if filler_schema is None:
                base_schema = dataset.schema
                fields: list[pa.Field] = [
                    base_schema.field(name)
                    for name in self.column_names
                    if name in base_schema.names
                ]
                fields.append(pa.field("_rowaddr", pa.uint64()))
                filler_schema = pa.schema(fields)

            if field_ids is None:
                field_ids = []
                for column_name in self.column_names:
                    field_ids.extend(
                        extract_field_ids(dataset.lance_schema, column_name)
                    )

        assert data_storage_version is not None
        assert num_physical_rows is not None
        assert num_logical_rows is not None
        assert filler_schema is not None
        assert field_ids is not None

        import more_itertools

        # we always write files that physically align with the fragment
        timings: dict[str, int] = {
            "align_ms": 0,
            "queue_wait_ms": 0,
            "checkpoint_read_ms": 0,
        }
        batch_stats: dict[str, int] = {
            "num_batches": 0,
            "total_rows": 0,
            "total_bytes": 0,
        }
        it = _buffer_and_sort_batches(
            num_logical_rows,
            self.fragment_id,
            filler_schema,
            self._store,
            self.checkpoint_keys,
            _timing=timings,
            _batch_stats=batch_stats,
        )

        it = _align_batches_to_physical_layout(
            num_physical_rows,
            num_logical_rows,
            self.fragment_id,
            it,
            _timing=timings,
        )

        # Filter batches to only include columns in the target schema.
        # This removes any source table columns that weren't selected in the
        # materialized view (e.g., if source has [id, title, width] but view
        # only selects [title], this removes id and width).
        it = _filter_columns_to_schema(it, self.column_names)

        it = more_itertools.peekable(it)

        file_id = str(uuid.uuid4())
        path = str(URL(self.uri) / "data" / f"{file_id}.lance")
        if not urllib.parse.urlparse(self.uri).scheme:
            path = f"file://{os.path.abspath(path)}"

        written = 0
        schema = it.peek().schema

        # Reconstruct storage_options_provider if we have namespace config
        # and the namespace provides storage_options (e.g., cloud credentials)
        storage_options_provider, storage_options = get_storage_options_provider(
            self.namespace_impl, self.namespace_properties, self.table_id
        )

        writer_write_ms = 0
        writer_cm = lance.file.LanceFileWriter(
            path,
            schema,
            storage_options=storage_options,
            storage_options_provider=storage_options_provider,
            version=data_storage_version,
        )
        t0 = time.perf_counter()
        writer = writer_cm.__enter__()
        writer_write_ms += int((time.perf_counter() - t0) * 1000)
        try:
            for batch in it:
                t1 = time.perf_counter()
                writer.write_batch(batch)
                writer_write_ms += int((time.perf_counter() - t1) * 1000)
                written += batch.num_rows
        finally:
            t1 = time.perf_counter()
            writer_cm.__exit__(None, None, None)
            writer_write_ms += int((time.perf_counter() - t1) * 1000)

        _LOG.debug(
            f"writing fragment file {file_id}.lance with cols:{self.column_names} "
        )
        dsv_major, dsv_minor = parse_data_storage_version(data_storage_version)
        new_datafile = lance.fragment.DataFile(
            f"{file_id}.lance",
            field_ids,
            list(range(len(field_ids))),
            dsv_major,
            dsv_minor,
        )

        align_ms = int(timings.get("align_ms", 0) or 0)
        queue_wait_ms = int(timings.get("queue_wait_ms", 0) or 0)
        checkpoint_read_ms = int(timings.get("checkpoint_read_ms", 0) or 0)
        num_batches = int(batch_stats.get("num_batches", 0) or 0)
        total_batch_rows = int(batch_stats.get("total_rows", 0) or 0)
        total_batch_bytes = int(batch_stats.get("total_bytes", 0) or 0)
        avg_batch_num_rows = (
            int(round(total_batch_rows / num_batches)) if num_batches else 0
        )
        avg_batch_size = (
            int(round(total_batch_bytes / num_batches)) if num_batches else 0
        )
        return FragmentWriteResult(
            frag_id=self.fragment_id,
            new_file=new_datafile,
            rows_written=written,
            align_ms=align_ms,
            write_ms=int(writer_write_ms),
            queue_wait_ms=queue_wait_ms,
            checkpoint_read_ms=checkpoint_read_ms,
            avg_batch_num_rows=avg_batch_num_rows,
            avg_batch_size=avg_batch_size,
        )


FragmentWriter: ray.actor.ActorClass = cast("ray.actor.ActorClass", FragmentWriter)
