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

import contextlib
import functools
import hashlib
import heapq
import logging
import random
import re
import uuid
from collections.abc import Iterator
from typing import cast

import attrs
import lance
import pyarrow as pa
import ray.actor
import ray.data
import ray.exceptions
import ray.util.queue

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

_LOG = logging.getLogger(__name__)

_SIMULATE_WRITE_FAILURE = False


@contextlib.contextmanager
def _simulate_write_failure(flag: bool) -> Iterator[None]:
    global _SIMULATE_WRITE_FAILURE
    current = _SIMULATE_WRITE_FAILURE
    try:
        _SIMULATE_WRITE_FAILURE = flag
        yield
    finally:
        _SIMULATE_WRITE_FAILURE = current


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


@ray.remote
@attrs.define
class ApplierActor:
    applier: LanceRecordBatchUDFApplier

    def run(self, task) -> tuple[ReadTask, dict[str, str]]:
        return task, self.applier.run(task)


ApplierActor: ray.actor.ActorClass = cast(ray.actor.ActorClass, ApplierActor)


@ray.remote(num_cpus=1)
@attrs.define
class FragmentWriter:
    uri: str
    column_name: str
    store: CheckpointStore

    fragment_id: int

    checkpoint_keys: ray.util.queue.Queue

    def write(self) -> tuple[int, lance.fragment.DataFile]:
        dataset = lance.dataset(self.uri)
        import more_itertools

        buffer: list[tuple[int, dict[str, str]]] = []
        num_rows = dataset.get_fragment(self.fragment_id).count_rows()
        written_rows = 0

        def _iter() -> Iterator[pa.RecordBatch]:
            nonlocal written_rows
            while written_rows < num_rows:
                while not buffer or buffer[0][0] != written_rows:
                    batch: tuple[int, dict[str, str]] = self.checkpoint_keys.get()
                    heapq.heappush(buffer, batch)

                data = heapq.heappop(buffer)[1]
                res = pa.RecordBatch.from_pydict(
                    {
                        key: _combine_chunks(self.store[value]["data"])
                        for key, value in data.items()
                    }
                )
                yield res
                written_rows += res.num_rows

        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,
            self.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: "{self.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 self.fragment_id, new_datafile


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


def _get_fragment_dedupe_key(
    uri: str, frag_id: int, output_column: str, transform: UDF
) -> str:
    key = f"{uri}:{frag_id}:{output_column}:{transform.checkpoint_key}"
    return hashlib.sha256(key.encode()).hexdigest()


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,
    concurrency: int = 8,
    applier_batch_size: int = 32,
    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

    assert len(transforms) == 1, "Only one column can be added at a time"

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

    pool = ray.util.ActorPool(
        [ApplierActor.remote(applier=applier) for _ in range(concurrency)]
    )

    plan = plan_read(
        uri,
        columns,
        batch_size=batch_size,
        read_version=read_version,
        **kwargs,
    )

    table_name = uri.split("/")[-1].split(".")[0]
    output_column = list(transforms.keys())[0]
    pbar_prefix = f"[{table_name} - {output_column}]"

    @functools.lru_cache
    def _fragment_is_done(frag_id: int) -> bool:
        dedupe_key = _get_fragment_dedupe_key(
            uri, frag_id, output_column, transforms[output_column]
        )
        return dedupe_key in checkpoint_store

    applier_iter = pool.map_unordered(
        lambda actor, value: actor.run.remote(value),
        # the API says list, but iterables are fine
        filter(lambda task: not _fragment_is_done(task.frag_id), plan),
    )

    writers = {}
    writer_queues: dict[int, ray.util.queue.Queue] = {}
    writer_futs_to_id = {}
    writer_task_cache = {}

    applier_pbar = tqdm(applier_iter, total=len(plan), position=0)

    applier_pbar.set_description(f"{pbar_prefix} Applying UDFs")

    ds = lance.dataset(uri, version=read_version)
    writer_pbar = tqdm(total=len(ds.get_fragments()), position=1)
    writer_pbar.set_description(f"{pbar_prefix} Writing Fragments")

    def _commit(frag_id, fut) -> None:
        nonlocal read_version
        _LOG.info("Committing fragment id: %d", frag_id)

        def _shutdown_writer() -> None:
            actor_handle: ray.actor.ActorHandle = writers[frag_id]
            ray.kill(actor_handle)
            del writers[frag_id]
            writer_queue: ray.util.queue.Queue = writer_queues[frag_id]
            writer_queue.shutdown()
            del writer_queues[frag_id]
            del writer_futs_to_id[fut]

        try:
            fut_frag_id, new_file = ray.get(fut)
            assert fut_frag_id == frag_id
        except (ray.exceptions.ActorDiedError, ray.exceptions.ActorUnavailableError):
            _LOG.exception("Failed to commit fragments, restarting writer")
            _shutdown_writer()
            # make new writer after the old one is shutdown
            queue = ray.util.queue.Queue()
            writer = FragmentWriter.remote(
                uri, output_column, checkpoint_store, frag_id, queue
            )
            writer_queues[frag_id] = queue
            writers[frag_id] = writer
            writer_futs_to_id[writer.write.remote()] = frag_id
            # re-add the tasks to the new writer
            for task in writer_task_cache[frag_id]:
                queue.put(task)
            _LOG.info("Restarted writer for fragment %d", frag_id)
            return

        operation = lance.LanceOperation.DataReplacement(
            replacements=[
                lance.LanceOperation.DataReplacementGroup(
                    fragment_id=frag_id,
                    new_file=new_file,
                )
            ]
        )

        lance.LanceDataset.commit(uri, operation, read_version=read_version)
        read_version += 1

        _shutdown_writer()
        del writer_task_cache[frag_id]

        writer_pbar.update(1)
        dedupe_key = _get_fragment_dedupe_key(
            uri, frag_id, output_column, transforms[output_column]
        )
        checkpoint_store[dedupe_key] = pa.RecordBatch.from_pydict({})

    for item in applier_pbar:
        task: ReadTask = item[0]
        result = item[1]

        frag_id = task.frag_id
        if frag_id not in writers:
            _LOG.info("Creating writer for fragment %d", frag_id)
            queue = ray.util.queue.Queue()
            writer = FragmentWriter.remote(
                uri, output_column, checkpoint_store, frag_id, queue
            )
            writer_queues[frag_id] = queue
            writers[frag_id] = writer
            writer_futs_to_id[writer.write.remote()] = frag_id
            writer_task_cache[frag_id] = []
        queue = writer_queues[frag_id]

        queue.put((task.offset, result))
        # save the task result before the fragment is committed
        # in the even the fragment writer fails we can restart it easily
        writer_task_cache[frag_id].append((task.offset, result))

        # FAULT INJECTION: simulate write failure for testing
        if _SIMULATE_WRITE_FAILURE and random.random() < 0.5:
            ray.kill(writers[frag_id])

        ready, _ = ray.wait(list(writer_futs_to_id.keys()), timeout=0)
        for fut in ready:
            frag_id = writer_futs_to_id[fut]
            _commit(frag_id, fut)

    # commit any remaining fragments
    # need to keep retry until all fragments are committed
    while writer_futs_to_id:
        for fut, frag_id in list(writer_futs_to_id.items()):
            _commit(frag_id, fut)
