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

from __future__ import annotations

import logging
from typing import TYPE_CHECKING

import attrs

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

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:
            table = table_ref.open()
            self._tables[key] = table
        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__)
