from __future__ import annotations

from collections.abc import Iterable, Iterator
from dataclasses import dataclass

from metadata import TableName
from storage.object_store import ObjectStore

from .paths import SchemaPath, parse_schema_path
from .schema import TableDefinition, load_table_definition


@dataclass(frozen=True)
class StorageSchemaIndex:
    table_defs: dict[tuple[TableName, int], TableDefinition]

    @property
    def table_names(self) -> set[TableName]:
        return {table_name for table_name, _table_version in self.table_defs}

    def get(self, table_name: TableName, table_version: int) -> TableDefinition | None:
        return self.table_defs.get((table_name, table_version))

    def latest_table_definitions(self) -> dict[TableName, TableDefinition]:
        latest: dict[TableName, TableDefinition] = {}
        for (table_name, _table_version), table_def in self.table_defs.items():
            current = latest.get(table_name)
            if current is None or table_def.table_version > current.table_version:
                latest[table_name] = table_def
        return latest


def discover_schema_paths(
    store: ObjectStore,
    incremental_prefix: str,
    *,
    high_watermark: int | None = None,
) -> list[SchemaPath]:
    return sorted(
        _iter_schema_paths(store, incremental_prefix, high_watermark=high_watermark),
        key=lambda path: (path.table_name, path.table_version, path.key),
    )


def discover_latest_schema_paths(
    store: ObjectStore,
    incremental_prefix: str,
    *,
    high_watermark: int | None = None,
) -> list[SchemaPath]:
    latest: dict[TableName, SchemaPath] = {}
    for path in _iter_schema_paths(store, incremental_prefix, high_watermark=high_watermark):
        current = latest.get(path.table_name)
        if current is None or path.table_version > current.table_version:
            latest[path.table_name] = path
    return sorted(latest.values(), key=lambda path: (path.table_name, path.table_version, path.key))


def _iter_schema_paths(
    store: ObjectStore,
    incremental_prefix: str,
    *,
    high_watermark: int | None = None,
) -> Iterator[SchemaPath]:
    for key in store.iter_keys(incremental_prefix):
        path = parse_schema_path(key, incremental_prefix)
        if path is None:
            continue
        if high_watermark is not None and path.table_version > high_watermark:
            continue
        yield path


def _load_schema_index_from_paths(store: ObjectStore, paths: Iterable[SchemaPath]) -> StorageSchemaIndex:
    table_defs: dict[tuple[TableName, int], TableDefinition] = {}
    for path in paths:
        table_def = load_table_definition(store.read_json(path.key), path)
        table_defs[(table_def.name, table_def.table_version)] = table_def
    return StorageSchemaIndex(table_defs=table_defs)


def load_storage_schema_index(
    store: ObjectStore,
    incremental_prefix: str,
    *,
    high_watermark: int | None = None,
) -> StorageSchemaIndex:
    return _load_schema_index_from_paths(
        store,
        discover_schema_paths(store, incremental_prefix, high_watermark=high_watermark),
    )


def load_latest_storage_schema_index(
    store: ObjectStore,
    incremental_prefix: str,
    *,
    high_watermark: int | None = None,
) -> StorageSchemaIndex:
    return _load_schema_index_from_paths(
        store,
        discover_latest_schema_paths(store, incremental_prefix, high_watermark=high_watermark),
    )


__all__ = [
    "StorageSchemaIndex",
    "discover_latest_schema_paths",
    "discover_schema_paths",
    "load_latest_storage_schema_index",
    "load_storage_schema_index",
]