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)
