from __future__ import annotations

from typing import Tuple

from pydantic import BaseModel
from typing_extensions import Literal

from ..models import GraphId, Sampling, VersionedGraphMetadata, validate_sampling
from ..store.graph_metadata_store import GraphMetadataStore


class Params(BaseModel):
    graph_id: GraphId
    sampling: Literal["head", "random", "all"]
    max_rows: int
    ref_version_token: str


def update_sampling(store: GraphMetadataStore, params: Params) -> Tuple[bool, VersionedGraphMetadata]:
    """
    Raises:
        ModelValidationError
        GraphDoesNotExistError
        ConcurrentUpdateCollisionError
        GraphMetadataStoreError
    """
    graph_id = params.graph_id
    graph_metadata = store.get(graph_id)

    updated = (
        graph_metadata["sampling"]["max_rows"] != params.max_rows
        or graph_metadata["sampling"]["sampling"] != params.sampling
    )

    new_graph_metadata = graph_metadata

    if updated:
        sampling = Sampling(sampling=params.sampling, max_rows=params.max_rows)
        validate_sampling(sampling)

        graph_metadata["sampling"] = sampling

        new_graph_metadata = store.update(graph_metadata, params.ref_version_token)

    return (updated, new_graph_metadata)
