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

from __future__ import annotations

import logging
from typing import TYPE_CHECKING

import attrs
from tenacity import (
    Retrying,
    retry_if_exception,
    stop_after_attempt,
    wait_exponential_jitter,
)

from geneva.apply.task import CopyTask, ReadTask, ScanTask
from geneva.utils import object_store_retry

if TYPE_CHECKING:
    from geneva.table import Table, TableReference

_LOG = logging.getLogger(__name__)

TableCacheKey = tuple[
    tuple[str, ...],
    str | None,
    str | None,
    tuple[tuple[str, str], ...] | None,
    tuple[str, ...] | None,
    int | None,
]


def _table_cache_key(table_ref: TableReference) -> TableCacheKey:
    namespace_properties = None
    if table_ref.namespace_properties:
        namespace_properties = tuple(sorted(table_ref.namespace_properties.items()))

    system_namespace = None
    if table_ref.system_namespace is not None:
        system_namespace = tuple(table_ref.system_namespace)

    return (
        tuple(table_ref.table_id),
        table_ref.db_uri,
        table_ref.namespace_impl,
        namespace_properties,
        system_namespace,
        table_ref.version,
    )


@attrs.define
class TableCache:
    _tables: dict[TableCacheKey, Table] = attrs.field(factory=dict)

    def get_or_open(self, table_ref: TableReference) -> Table:
        key = _table_cache_key(table_ref)
        table = self._tables.get(key)
        if table is None:
            max_retries, retry_base_backoff, retry_max_backoff = (
                object_store_retry.get_applier_retry_settings()
            )

            def _before_sleep(retry_state) -> None:  # noqa: ANN001
                exc = (
                    retry_state.outcome.exception()
                    if retry_state.outcome is not None
                    else None
                )
                backoff = (
                    retry_state.next_action.sleep
                    if retry_state.next_action is not None
                    else 0.0
                )
                _LOG.warning(
                    (
                        "Transient object store error opening table %s "
                        "(retry %d/%d). Sleeping %.2fs before retry. Error: %s"
                    ),
                    table_ref.table_id,
                    retry_state.attempt_number,
                    max_retries,
                    backoff,
                    exc,
                )

            retrier = Retrying(
                retry=retry_if_exception(
                    object_store_retry.is_retryable_object_store_error
                ),
                stop=stop_after_attempt(max_retries + 1),
                wait=wait_exponential_jitter(
                    initial=retry_base_backoff,
                    max=retry_max_backoff,
                ),
                reraise=True,
                before_sleep=_before_sleep,
            )

            for attempt in retrier:
                with attempt:
                    table = table_ref.open()
            assert table is not None
            self._tables[key] = table
        assert table is not None
        return table


def bind_tables_for_task(task: ReadTask, table_cache: TableCache) -> None:
    if isinstance(task, ScanTask):
        table = table_cache.get_or_open(task.table_ref_for_read())
        task.bind_table(table)
        return

    if isinstance(task, CopyTask):
        src_table = table_cache.get_or_open(task.src)
        dst_table = table_cache.get_or_open(task.dst)
        task.bind_tables(src=src_table, dst=dst_table)
        return

    _LOG.warning("bind_tables_for_task: unhandled task type %s", type(task).__name__)


def clear_bound_tables(task: ReadTask) -> None:
    if isinstance(task, ScanTask):
        task.clear_table()
        return

    if isinstance(task, CopyTask):
        task.clear_tables()
        return

    _LOG.warning("clear_bound_tables: unhandled task type %s", type(task).__name__)
