# SPDX-License-Identifier: PROPRIETARY
# SPDX-FileCopyrightText: Copyright The Geneva Authors
import json
import logging
import time
from typing import Any

import attr
import attrs
import ray
from lancedb import AsyncConnection, AsyncTable
from lancedb.namespace import AsyncLanceNamespaceDBConnection
from lancedb.util import (
    value_to_sql,
)

from geneva.jobs.jobs import GENEVA_JOBS_TABLE_NAME
from geneva.table import TableReference
from geneva.utils import dt_now_utc, escape_sql_string

_LOG = logging.getLogger(__name__)

# Performance metrics (cumulative across all tasks/fragments unless noted).
# Time metrics are in milliseconds and may overlap; they are not additive.
#
# UDF/read path (ApplierActor + CheckpointingApplier):
# Performance metrics (cumulative across all tasks/fragments unless noted).
# Time metrics are in milliseconds and may overlap; they are not additive.
#
# UDF/read path (ApplierActor + CheckpointingApplier):
METRIC_UDF_PROCESSING_TIME = "udf_processing_time"
METRIC_BATCH_CHECKPOINTING_TIME = "batch_checkpointing_time"
METRIC_READ_IO_TIME = "read_io_time_ms"
METRIC_CHECKPOINT_LOAD_TIME = "checkpoint_load_time_ms"
# Time spent checking task-level checkpoint existence (task_key in store).
METRIC_CHECKPOINT_EXISTS_TIME = "checkpoint_exists_time_ms"
# Time spent listing per-fragment checkpoints (plan/checkpoint scan).
METRIC_CHECKPOINT_LIST_TIME = "checkpoint_list_time_ms"
# Wall time per read task (includes read IO + UDF + checkpointing); overlaps
# with other metrics, so do not sum with them.
METRIC_READ_TASK_TOTAL_TIME = "read_task_total_time_ms"

# Planning (per job):
METRIC_PLAN_READ_TIME = "plan_read_time_ms"

# Writer path (FragmentWriter + FragmentWriterManager):
METRIC_FRAGMENT_CHECKPOINTING_TIME = "fragment_checkpointing_time"
METRIC_WRITER_ALIGN_TIME = "writer_align_time"
METRIC_WRITER_WRITE_TIME = "writer_write_time"
METRIC_WRITER_QUEUE_WAIT_TIME = "writer_queue_wait_time_ms"
METRIC_WRITER_CHECKPOINT_READ_TIME = "writer_checkpoint_read_time_ms"
# Deprecated compatibility metric. Internal timing has been removed and this is
# expected to stay at 0.
METRIC_WRITER_BUFFER_SORT_TIME = "writer_buffer_sort_time"

# Commit path (dataset commit attempts and retries):
METRIC_COMMIT_TIME = "commit_time_ms"
METRIC_COMMIT_BACKOFF_TIME = "commit_backoff_time_ms"
METRIC_COMMIT_CONFLICT_RETRIES = "commit_conflict_retries"
METRIC_COMMIT_CONCURRENT_WRITER_RETRIES = "commit_concurrent_writer_retries"

# Batch shape (writer-reconciled averages):
METRIC_AVG_BATCH_NUM_ROWS = "avg_batch_num_rows"
METRIC_AVG_BATCH_SIZE = "avg_batch_size"


ROW_PROGRESS_METRICS = frozenset(
    {"rows_checkpointed", "rows_ready_for_commit", "rows_committed"}
)


@ray.remote
@attrs.define
class JobTracker:
    """
    Centralized progress ledger for a job.
    - Metrics are arbitrary strings -> {n, total, done, desc}
    - All ops are atomic via the actor's mailbox.
    - DB saves can be disabled entirely with enable_saves=False
    """

    job_id: str

    table_ref: TableReference = attrs.field()

    metrics: dict[str, dict] = attrs.field(factory=dict)

    enable_saves: bool = attrs.field(default=True)

    min_time_between_updates_secs: float = attrs.field(default=5.0)

    # use async lancedb connection to avoid blocking Ray Tasks
    _db: AsyncConnection | AsyncLanceNamespaceDBConnection | None = attrs.field(
        init=False, default=None
    )

    _jobs_table: AsyncTable | None = attrs.field(init=False, default=None)

    _last_updated: float = attrs.field(init=False, default=-float("inf"))
    _rows_finalized: bool = attrs.field(init=False, default=False)

    def _should_skip_row_update(self, name: str) -> bool:
        return self._rows_finalized and name in ROW_PROGRESS_METRICS

    _job_done: bool = attrs.field(init=False, default=False)

    async def _upsert(
        self, name: str, *, total: int | None = None, desc: str | None = None
    ) -> None:
        """
        Given a metric 'name', ensure that the metric entry is present with specified
        values or initialized with 0s.
        This will flush the metrics to the underlying lance table if
        min_time_between_updates_secs is exceeded
        """
        m = self.metrics.get(name)
        if m is None:
            m = {"n": 0, "total": 0, "done": False, "desc": name}
            self.metrics[name] = m
        if total is not None:
            m["total"] = int(total)
        if desc is not None:
            m["desc"] = desc
        # auto-done if we're already at/over total
        if m["total"] and m["n"] >= m["total"]:
            m["done"] = True

        await self._save_with_throttle()

    # Generic metric API
    async def set_total(self, name: str, total: int) -> None:
        if self._should_skip_row_update(name):
            return
        await self._upsert(name, total=total)

    async def set_desc(self, name: str, desc: str) -> None:
        await self._upsert(name, desc=desc)

    async def set(self, name: str, n: int) -> None:
        if self._should_skip_row_update(name):
            return
        await self._upsert(name)
        m = self.metrics[name]
        m["n"] = int(n)
        if m["total"] and m["n"] >= m["total"]:
            m["done"] = True
            await self._save_with_throttle(force=True)

    async def increment(self, name: str, delta: int = 1) -> None:
        if self._should_skip_row_update(name):
            return
        await self._upsert(name)
        m = self.metrics[name]
        m["n"] += int(delta)
        if m["total"] and m["n"] >= m["total"]:
            m["done"] = True
            await self._save_with_throttle(force=True)

    async def batch_increment(self, deltas: dict[str, int]) -> None:
        """Increment multiple metrics in a single actor call/save cycle."""
        if not deltas:
            return

        if self._rows_finalized:
            deltas = {
                name: delta
                for name, delta in deltas.items()
                if name not in ROW_PROGRESS_METRICS
            }
            if not deltas:
                return

        force_save = False
        updated = False
        for name, delta in deltas.items():
            inc = int(delta)
            if inc == 0:
                continue
            m = self.metrics.get(name)
            if m is None:
                m = {"n": 0, "total": 0, "done": False, "desc": name}
                self.metrics[name] = m
            m["n"] += inc
            updated = True
            if m["total"] and m["n"] >= m["total"]:
                m["done"] = True
                force_save = True

        if not updated:
            return
        await self._save_with_throttle(force=force_save)

    async def mark_done(self, name: str) -> None:
        if self._should_skip_row_update(name):
            return
        await self._upsert(name)
        self.metrics[name]["done"] = True
        await self._save_with_throttle(force=True)

    # Lifecycle
    def mark_job_done(self) -> None:
        """Mark this job as finished (success or failure)."""
        self._job_done = True

    def is_job_done(self) -> bool:
        return self._job_done

    async def finalize_rows(
        self,
        rows_checkpointed: int,
        rows_ready_for_commit: int,
        rows_committed: int,
    ) -> None:
        """Atomically finalize core row metrics for a completed job."""
        self._rows_finalized = True
        for name, value in (
            ("rows_checkpointed", rows_checkpointed),
            ("rows_ready_for_commit", rows_ready_for_commit),
            ("rows_committed", rows_committed),
        ):
            m = self.metrics.get(name)
            if m is None:
                m = {"n": 0, "total": 0, "done": False, "desc": name}
                self.metrics[name] = m
            m["n"] = int(value)
            m["total"] = int(value)
            m["done"] = True
        await self._save_with_throttle(force=True)

    # Read APIs
    async def get_progress(self, name: str) -> dict:
        await self._upsert(name)
        m = self.metrics[name]
        return {"n": m["n"], "total": m["total"], "done": m["done"], "desc": m["desc"]}

    def get_all(self) -> dict[str, dict]:
        # Shallow copy for safety
        return {k: dict(v) for k, v in self.metrics.items()}

    async def _save_with_throttle(self, force: bool = False) -> None:
        """Save metrics to DB, optionally bypassing throttle."""
        current_time = time.time()
        if (
            not force
            and self._last_updated + self.min_time_between_updates_secs > current_time
        ):
            # don't update metrics too frequently to avoid excessive DB writes
            return

        # bump last update and save metrics
        self._last_updated = current_time
        await self._save_metrics(attr.asdict(self)["metrics"])

    async def _save_metrics(self, _metrics: dict[str, dict]) -> None:
        if not self.enable_saves:
            return
        if not _metrics:
            return

        try:
            # lazy load db and table
            if not self._db:
                if not self.table_ref:
                    _LOG.error(
                        "JobTracker missing table_ref, metrics will not be saved"
                    )
                    return
                self._db = await self.table_ref.open_system_db_async()
                namespace = self.table_ref.system_table_namespace
                self._jobs_table = await self._db.open_table(
                    GENEVA_JOBS_TABLE_NAME, namespace=namespace
                )

            updates = {
                "metrics": self._convert_metrics(_metrics),
                "updated_at": dt_now_utc(),
            }
            updates_sql = {k: value_to_sql(v) for k, v in updates.items()}
            _LOG.debug(f"Saving metrics for job {self.job_id}: {updates_sql}")
            start = time.time()

            await self._jobs_table.update(
                where=f"job_id = '{escape_sql_string(self.job_id)}'",
                updates_sql=updates_sql,
            )
            _LOG.info(
                f"Saved metrics for job "
                f"{self.job_id} in {(time.time() - start) * 1000:.0f}ms"
            )

        except AttributeError as e:
            _LOG.warning(
                f"unable to save metrics. The execution manifest "
                f"may be using an older Geneva version than the client"
                f": {e}"
            )
        except Exception as e:
            _LOG.error(f"error saving metrics {e}")
            self._db = None  # force retry on next save

    def _convert_metrics(self, _metrics: dict[str, dict]) -> list[str]:
        """Convert metrics to list of json strings for storage"""
        m = []
        for name, v in _metrics.items():
            vm = dict(v)
            vm["name"] = name
            m.append(json.dumps(vm))
        return m

    def __str__(self) -> str:
        return f"JobTracker(job_id={self.job_id})"

    def __repr__(self) -> str:
        return str(self)


def job_tracker_options(**kwargs: Any) -> Any:
    """Create JobTracker actor options with serialized method execution."""
    return JobTracker.options(max_concurrency=1, **kwargs)
