# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The Geneva Authors

import logging
import re
import uuid
from collections.abc import Callable, Iterator

import lance
import pandas as pd
import pyarrow as pa
import ray.data

from geneva.apply import LanceRecordBatchUDFApplier, plan_read
from geneva.checkpoint import CheckpointStore
from geneva.debug.logger import CheckpointStoreErrorLogger
from geneva.transformer import UDF

_LOG = logging.getLogger(__name__)


def map_with_value(  # noqa: ANN201
    fn: Callable,
):
    def _wrapped(x: dict):  # noqa: ANN202
        result = {
            "task": x["item"],
            "checkpoint_key": fn(x["item"]),
            "frag_id": x["item"].frag_id,
        }
        return result

    return _wrapped


def sort_and_collect_batch(df: pd.DataFrame):  # noqa: ANN201
    data = df.to_dict(orient="records")
    frag_id = data[0]["task"].frag_id
    _LOG.info("Sorting %d keys for frag: %s", len(data), frag_id)
    data.sort(key=lambda x: x["task"].offset)

    checkpoint_keys = [x["checkpoint_key"] for x in data]
    return {"item": [{"frag_id": frag_id, "checkpoint_keys": checkpoint_keys}]}


def _combine_chunks(arr: pa.Array) -> pa.Array:
    if isinstance(arr, pa.ChunkedArray):
        assert len(arr.chunks) == 1
        return arr.chunks[0]
    return arr


def write_fragment(  # noqa: ANN201
    uri: str,
    column_name: str,
    store: CheckpointStore,
):
    dataset = lance.dataset(uri)

    def _write_fragment(x):  # noqa: ANN202
        frag_id = x["item"]["frag_id"]
        checkpoint_keys = x["item"]["checkpoint_keys"]
        batches: list[dict[str, str]] = list(checkpoint_keys)
        import more_itertools

        def _iter() -> Iterator[pa.RecordBatch]:
            for batch in batches:
                res = {
                    key: _combine_chunks(store[value]["data"])
                    for key, value in batch.items()
                }
                yield pa.RecordBatch.from_pydict(res)

        it = more_itertools.peekable(_iter())
        rbr = pa.RecordBatchReader.from_batches(it.peek().schema, it)

        # TODO: this doesn't support struct or complex schema yet
        new_data = lance.fragment.write_fragments(
            rbr,
            uri,
            max_rows_per_file=1 << 31,
            max_bytes_per_file=1 << 40,
        )

        assert len(new_data) == 1
        new_data = new_data[0]
        assert len(new_data.files) == 1
        new_datafile = new_data.files[0]

        # MASSIVE HACK: open up an API to get the field id from the column name
        field_id = re.compile(
            rf'name: "{column_name}", id: (?P<field_id>[\d]*),'
        ).findall(str(dataset.lance_schema))
        assert len(field_id) == 1
        field_id = int(field_id[0])

        new_datafile.fields = [field_id]
        new_datafile.column_indices = [0]

        return {
            "frag_id": frag_id,
            "new_datafile": new_datafile,
        }

    return _write_fragment


def commit_fragments(  # noqa: ANN201
    uri: str,
    version: int,
):
    def _commit_fragments(x) -> None:
        _LOG.info("Committing %d fragments", len(x))
        operation = lance.LanceOperation.DataReplacement(
            replacements=[
                lance.LanceOperation.DataReplacementGroup(
                    fragment_id=item["frag_id"],
                    new_file=item["new_datafile"],
                )
                for item in x
            ]
        )
        lance.LanceDataset.commit(uri, operation, read_version=version)

    return _commit_fragments


def run_ray_add_column(
    uri: str,
    columns: list[str],
    transforms: dict[str, UDF],
    checkpoint_store: CheckpointStore,
    /,
    job_id: str | None = None,
    batch_size: int = 8,
    read_version: int | None = None,
    test_run: bool = True,
    **kwargs,
) -> None:
    if read_version is None:
        read_version = lance.dataset(uri).version

    job_id = job_id or uuid.uuid4().hex

    applier = LanceRecordBatchUDFApplier(
        udfs=transforms,
        checkpoint_store=checkpoint_store,
        error_logger=CheckpointStoreErrorLogger(job_id, checkpoint_store),
    )

    tasks = list(
        plan_read(uri, columns, batch_size=batch_size, read_version=read_version)
    )

    if not len(tasks):
        return None

    output_column = list(transforms.keys())[0]

    ds = (
        ray.data.from_items(tasks)
        .random_shuffle()
        .map(map_with_value(applier.run))
        .groupby("frag_id")
        .map_groups(sort_and_collect_batch)
        .map(write_fragment(uri, output_column, checkpoint_store))
    )
    # trigger collection

    frags = list(ds.iterator().iter_rows())
    _LOG.info("Committing %d fragments", len(frags))
    operation = lance.LanceOperation.DataReplacement(
        replacements=[
            lance.LanceOperation.DataReplacementGroup(
                fragment_id=item["frag_id"],
                new_file=item["new_datafile"],
            )
            for item in frags
        ]
    )
    ds = lance.dataset(uri)
    lance.LanceDataset.commit(uri, operation, read_version=ds.version)
