from __future__ import annotations

from pathlib import Path
from typing import Protocol, runtime_checkable

from fatum.experiment.types import (
    ArtifactKey,
    ArtifactPath,
    ExperimentID,
    FilePath,
    MetricKey,
    ParameterKey,
    StorageKey,
)


@runtime_checkable
class StorageBackend(Protocol):
    def save(self, key: StorageKey, path: Path) -> None: ...
    def load(self, key: StorageKey, path: Path) -> None: ...
    def exists(self, key: StorageKey) -> bool: ...
    def delete(self, key: StorageKey) -> None: ...
    def list(self, prefix: StorageKey) -> list[StorageKey]: ...


@runtime_checkable
class ExperimentProtocol(Protocol):
    @property
    def id(self) -> ExperimentID: ...
    @property
    def metadata(self) -> ExperimentMetadata: ...
    @property
    def artifacts(self) -> dict[ArtifactKey, StorageKey]: ...
    @property
    def metrics(self) -> dict[MetricKey, float | int]: ...
    @property
    def parameters(self) -> dict[ParameterKey, object]: ...

    def save_artifact(self, path: FilePath, artifact_path: ArtifactPath | None = None) -> ArtifactKey: ...
    def load_artifact(self, key: ArtifactKey, path: FilePath) -> None: ...
    def record_metric(self, key: MetricKey, value: float) -> None: ...
    def add_parameter(self, key: ParameterKey, value: object) -> None: ...
    def complete(self) -> None: ...


@runtime_checkable
class ExperimentMetadata(Protocol):
    @property
    def experiment_id(self) -> ExperimentID: ...
    @property
    def created_at(self) -> float: ...
    @property
    def completed_at(self) -> float | None: ...
    @property
    def status(self) -> str: ...
    @property
    def tags(self) -> list[str]: ...
    @property
    def description(self) -> str: ...
