# ruff: noqa: PERF203

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

# super simple applier

import logging
import time
from collections.abc import Iterator

import attrs
import pyarrow as pa

from geneva.apply.applier import BatchApplier
from geneva.apply.error_handling import BatchStrategy, ErrorHandlingContext
from geneva.apply.task import MapTask, ReadTask
from geneva.apply.utils import _iter_with_next_duration
from geneva.debug.logger import ErrorLogger

_LOG = logging.getLogger(__name__)


@attrs.define
class SimpleApplier(BatchApplier):
    """
    A simple applier that applies a function to each element in the batch.
    """

    job_id: str = attrs.field(default="unknown")
    udf_processing_time_ms: int = attrs.field(default=0, init=False)
    read_io_time_ms: int = attrs.field(default=0, init=False)

    def reset_timing_metrics(self) -> None:
        self.udf_processing_time_ms = 0
        self.read_io_time_ms = 0

    def run(
        self,
        read_task: ReadTask,
        map_task: MapTask,
        error_logger: ErrorLogger,
    ) -> Iterator[pa.RecordBatch]:
        self.reset_timing_metrics()
        batch_iter = read_task.to_batches(batch_size=map_task.batch_size())
        for seq, (read_ms, batch) in enumerate(
            _iter_with_next_duration(iter(batch_iter))
        ):
            self.read_io_time_ms += read_ms
            # Create error context and strategy for this batch
            ctx = ErrorHandlingContext.create(map_task, read_task, self.job_id, seq)
            strategy = BatchStrategy.from_context(ctx, map_task, error_logger)

            # Apply strategy - no branching needed!
            start = time.perf_counter()
            result_batch, error_records = strategy.apply(batch)
            self.udf_processing_time_ms += int((time.perf_counter() - start) * 1000)

            # Log errors if any were collected
            if error_records:
                error_logger.log_errors(error_records)

            yield result_batch
