"""Allocate an experiment to a workflow run."""

import tempfile
import warnings
from collections.abc import Callable, Sequence
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime
from pathlib import Path
from types import NoneType
from typing import TYPE_CHECKING, Any, Generic, Literal, overload

import pandas as pd
import yaml
from hatchet_sdk import TaskRunRef, TriggerWorkflowOptions, WorkflowRunRef
from hatchet_sdk.runnables.workflow import Workflow
from pydantic import BaseModel, Field, FilePath, field_serializer, field_validator
from tqdm import tqdm

from scythe.registry import ExperimentRegistry, Standalone, TInput, TOutput
from scythe.scatter_gather import (
    RecursionMap,
    ScatterGatherInput,
    ScatterGatherResult,
    scatter_gather,
)
from scythe.settings import ScytheStorageSettings
from scythe.utils.filesys import S3Url
from scythe.utils.results import save_and_upload_parquets
from scythe.utils.s3 import s3_client as s3

if TYPE_CHECKING:
    from mypy_boto3_s3.client import S3Client
    from mypy_boto3_s3.type_defs import CommonPrefixTypeDef, ObjectTypeDef
else:
    S3Client = object
    CommonPrefixTypeDef = object
    ObjectTypeDef = object


class ExperimentSpecsMismatchError(Exception):
    """An error raised when the specs for an experiment do not match the expected type."""

    def __init__(self, expected_type: type, actual_type: type):
        """Initialize the error."""
        self.expected_type = expected_type
        self.actual_type = actual_type
        super().__init__(
            f"Expected type {expected_type.__name__}, got {actual_type.__name__}."
        )


class DuplicateInputArtifactsError(Exception):
    """An error raised when a file is duplicated in the input artifacts."""

    def __init__(self, file_name: str, field_name: str):
        """Initialize the error."""
        self.file_name = file_name
        self.field_name = field_name
        super().__init__(
            f"File with name {file_name} for field {field_name} is duplicated in the "
            f"input artifacts (i.e. multiple files with the same name but in different "
            "directories.) "
        )


class BaseNotFoundError(Exception):
    """An error raised when a resource is not found."""


class ExperimentNotFoundError(BaseNotFoundError):
    """An error raised when an experiment is not found."""

    def __init__(self, experiment_name: str):
        """Initialize the error."""
        self.experiment_name = experiment_name
        super().__init__(f"Experiment {experiment_name} not found.")


class ExperimentVersionNotFoundError(BaseNotFoundError):
    """An error raised when an experiment version is not found."""

    def __init__(self, experiment_name: str, version: "SemVer"):
        """Initialize the error."""
        self.experiment_name = experiment_name
        self.version = version
        super().__init__(f"Experiment {experiment_name}/{version} not found.")


class ExperimentRunNotFoundError(BaseNotFoundError):
    """An error raised when an experiment run is not found."""

    def __init__(self, experiment_name: str, version: "SemVer", run_name: str):
        """Initialize the error."""
        self.experiment_name = experiment_name
        self.version = version
        self.run_name = run_name
        super().__init__(
            f"Experiment {experiment_name}/{version}/{run_name} not found."
        )


class InputArtifactLocations(BaseModel):
    """The locations of the input artifacts for an experiment."""

    files: dict[str, set[S3Url]]


class ExperimentRunManifest(BaseModel):
    """The manifest for an experiment run."""

    workflow_run_id: str
    experiment_id: str
    experiment_name: str
    specs_uri: S3Url
    io_spec: S3Url
    input_artifacts: S3Url | None = None


class SemVer(BaseModel):
    """A semantic version."""

    major: int
    minor: int = 0
    patch: int = 0

    @classmethod
    def FromString(cls, version: str) -> "SemVer":
        """Parse a semantic version from a string."""
        if not version.startswith("v"):
            msg = "Version must start with 'v'"
            raise ValueError(msg)
        version = version[1:]
        delimiter = "." if "." in version else "-"
        parts = version.split(delimiter)
        if len(parts) == 0:
            msg = "Version must have at least one part"
            raise ValueError(msg)
        major = int(parts[0])
        minor = int(parts[1]) if len(parts) > 1 else 0
        patch = int(parts[2]) if len(parts) > 2 else 0
        return cls(major=major, minor=minor, patch=patch)

    def __lt__(self, other: "SemVer") -> bool:
        """Compare two semantic versions."""
        if self.major != other.major:
            return self.major < other.major
        if self.minor != other.minor:
            return self.minor < other.minor
        return self.patch < other.patch

    def __le__(self, other: "SemVer") -> bool:
        """Compare two semantic versions."""
        if self.major != other.major:
            return self.major <= other.major
        if self.minor != other.minor:
            return self.minor <= other.minor
        return self.patch <= other.patch

    def __gt__(self, other: "SemVer") -> bool:
        """Compare two semantic versions."""
        if self.major != other.major:
            return self.major > other.major
        if self.minor != other.minor:
            return self.minor > other.minor
        return self.patch > other.patch

    def __ge__(self, other: "SemVer") -> bool:
        """Compare two semantic versions."""
        if self.major != other.major:
            return self.major >= other.major
        if self.minor != other.minor:
            return self.minor >= other.minor
        return self.patch >= other.patch

    def next_major_version(self) -> "SemVer":
        """Get the next major version."""
        return SemVer(major=self.major + 1, minor=0, patch=0)

    def next_minor_version(self) -> "SemVer":
        """Get the next minor version."""
        return SemVer(major=self.major, minor=self.minor + 1, patch=0)

    def next_patch_version(self) -> "SemVer":
        """Get the next patch version."""
        return SemVer(major=self.major, minor=self.minor, patch=self.patch + 1)

    def __str__(self) -> str:
        """Get the string representation of the semantic version."""
        return f"v{self.major}.{self.minor}.{self.patch}"


VersioningStrategy = Literal["bumpmajor", "bumpminor", "bumppatch", "keep"]
DatetimeFormat = "%Y-%m-%d_%H-%M-%S"


class SerializableRunnable(
    BaseModel, Generic[TInput, TOutput], arbitrary_types_allowed=True
):
    """A serializable runnable."""

    runnable: Standalone[TInput, TOutput] | Workflow[TInput]

    def __deepcopy__(self, memo: dict[int, Any] | None = None):
        """Deep copy the runnable which is not copyable."""
        warnings.warn(
            "Deep copying a SerializableRunnable will not copy the runnable itself.",
            stacklevel=3,
        )
        runnable = self.runnable
        self.runnable = None  # pyright: ignore [reportAttributeAccessIssue]
        new_self = super().__deepcopy__(memo)
        new_self.runnable = runnable
        self.runnable = runnable
        return new_self

    @field_validator("runnable", mode="before")
    def get_runnable_from_str(cls, v):
        """Get the runnable from a string."""
        if isinstance(v, str):
            runnable = ExperimentRegistry.get_runnable(v)
            return runnable
        return v

    @field_serializer("runnable")
    def serialize_runnable(self, v, b) -> str:
        """Serialize the runnable to a string."""
        return v.name


class BaseExperiment(
    SerializableRunnable[TInput, TOutput], arbitrary_types_allowed=True
):
    """A base experiment."""

    run_name: str | None = None
    storage_settings: ScytheStorageSettings = Field(
        default_factory=lambda: ScytheStorageSettings()
    )

    _latest_version_cache: "VersionedExperiment[TInput, TOutput] | None" = None

    @property
    def base_id(self) -> str:
        """The base experiment id."""
        return self.run_name or self.runnable.name

    @property
    def prefix(self) -> str:
        """The prefix for the experiment."""
        return f"{self.storage_settings.BUCKET_PREFIX}/{self.base_id}/"

    def list_versions(
        self, s3_client: S3Client | None = None
    ) -> list["VersionedExperiment[TInput, TOutput]"]:
        """Get all of the versions of the experiment."""
        s3_client = s3_client or s3

        paginator = s3_client.get_paginator("list_objects_v2")

        common_prefixes: list[CommonPrefixTypeDef] = []
        for page in paginator.paginate(
            Bucket=self.storage_settings.BUCKET,
            Prefix=self.prefix,
            Delimiter="/",
            PaginationConfig={"PageSize": 1000},
        ):
            common_prefixes.extend(page.get("CommonPrefixes", []))

        prefixes = [d.get("Prefix") for d in common_prefixes]
        version_strings = [p.split("/")[-2] for p in prefixes if p]

        versions = [SemVer.FromString(vs) for vs in version_strings]
        return [VersionedExperiment(base_experiment=self, version=v) for v in versions]

    def latest_version(
        self, s3_client: S3Client | None = None, from_cache: bool = True
    ) -> "VersionedExperiment[TInput, TOutput] | None":
        """Get the latest version of the experiment."""
        if not from_cache or "_latest_version_cache" not in self.model_fields_set:
            versions = self.list_versions(s3_client)
            latest = max(versions, key=lambda v: v.version) if versions else None
            self._latest_version_cache = latest
        return self._latest_version_cache

    @property
    def latest_run(self) -> "ExperimentRun":
        """The latest run for the experiment."""
        latest_version = self.latest_version()
        if latest_version is None:
            raise ExperimentNotFoundError(self.base_id)
        return latest_version.latest_run

    def latest_results_for_version(self, version: SemVer) -> dict[str, str]:
        """The latest results for a given version."""
        versioned_experiment = VersionedExperiment(
            base_experiment=self, version=version
        )
        return versioned_experiment.latest_run_results

    @property
    def latest_results(self) -> dict[str, str]:
        """The latest results for the experiment."""
        latest_version = self.latest_version()
        if latest_version is None:
            raise ExperimentNotFoundError(self.runnable.name)
        return latest_version.latest_run_results

    def resolve_next_version(
        self, version: SemVer | VersioningStrategy, s3_client: S3Client | None = None
    ) -> SemVer:
        """Resolve the next version of the experiment."""
        s3_client = s3_client or s3
        if not isinstance(version, SemVer):
            latest_version = self.latest_version(s3_client)
            if latest_version is None:
                version = SemVer(major=1, minor=0, patch=0)
            elif version == "bumpmajor":
                version = latest_version.version.next_major_version()
            elif version == "bumpminor":
                version = latest_version.version.next_minor_version()
            elif version == "bumppatch":
                version = latest_version.version.next_patch_version()
            else:
                version = latest_version.version
        return version

    def check_spec_type(self, spec: TInput) -> None:
        """Check that the type of the spec matches the expected type."""
        if not isinstance(spec, self.runnable.input_validator_type):
            raise ExperimentSpecsMismatchError(
                expected_type=self.runnable.input_validator_type,
                actual_type=type(spec),
            )

    def check_spec_types(self, specs: Sequence[Any]) -> None:
        """Check that the types of the specs match the expected type."""
        mismatching_types = [
            type(spec)
            for spec in specs
            if not isinstance(spec, self.runnable.input_validator_type)
        ]
        if mismatching_types:
            raise ExperimentSpecsMismatchError(
                expected_type=self.runnable.input_validator_type,
                actual_type=mismatching_types[0],
            )

    @overload
    def allocate(
        self,
        specs: Sequence[TInput],
        version: SemVer | VersioningStrategy = "bumpmajor",
        overwrite_sort_index: bool = True,
        overwrite_experiment_id: bool = True,
        recursion_map: RecursionMap | None = None,
        s3_client: S3Client | None = None,
    ) -> tuple[
        "ExperimentRun", TaskRunRef[ScatterGatherInput, ScatterGatherResult]
    ]: ...

    @overload
    def allocate(
        self,
        specs: TInput,
        version: SemVer | VersioningStrategy = "bumpmajor",
        overwrite_sort_index: bool = True,
        overwrite_experiment_id: bool = True,
        recursion_map: RecursionMap | None = None,
        s3_client: S3Client | None = None,
    ) -> tuple["ExperimentRun", WorkflowRunRef]: ...

    def allocate(  # noqa: C901
        self,
        specs: Sequence[TInput] | TInput,
        version: SemVer | VersioningStrategy = "bumpmajor",
        overwrite_sort_index: bool = True,
        overwrite_experiment_id: bool = True,
        recursion_map: RecursionMap | None = None,
        s3_client: S3Client | None = None,
    ) -> tuple[
        "ExperimentRun",
        TaskRunRef[ScatterGatherInput, ScatterGatherResult] | WorkflowRunRef,
    ]:
        """Allocate an experiment to a workflow run."""
        s3_client = s3_client or s3
        version = self.resolve_next_version(version, s3_client)
        self.check_spec_types(specs) if isinstance(
            specs, Sequence
        ) else self.check_spec_type(specs)

        cur_time = datetime.now()
        experiment_run = ExperimentRun(
            versioned_experiment=VersionedExperiment(
                base_experiment=self, version=version
            ),
            timestamp=cur_time,
        )

        spec_list = specs if isinstance(specs, Sequence) else [specs]
        experiment_run.overwrite_spec_meta(
            spec_list, overwrite_experiment_id, overwrite_sort_index
        )

        # handle local file transfer to s3
        # TODO: this will cause a race condition if multiple files have the same name but
        # live in different directories.
        # hence the checks below
        def construct_artifact_key(pth: FilePath, field_name: str):
            # The race condition happens here because we are only using pth.name
            # which loses any namespacing from ancestral directories.
            return experiment_run.construct_artifact_key(field_name, pth.name)

        # first, we need to get a record of the `field: FilePath` pairs
        # for each set of specs
        local_input_artifact_paths = [
            spec._local_artifact_file_paths for spec in spec_list
        ]
        # then, we transpose the list of dicts into a dict of sets
        # this is so that we can deduplicate shared input artifacts
        # grouped by the field they are used for.
        input_artifacts: dict[str, set[FilePath]] = {}
        at_least_one_input_artifact = False
        for spec_paths in local_input_artifact_paths:
            for field_name, fpath in spec_paths.items():
                input_artifacts.setdefault(field_name, set()).add(fpath)
                at_least_one_input_artifact = True

        # next, we want to detect collisions where input artifacts
        # are used for the same field and have the same filename but live in different
        # directories, meaning they could theoretically be different files, and would
        # result in a race condition since currently they would be uploaded with the same
        # key.
        all_input_artifact_names: dict[str, set[str]] = {}
        for field_name, fpaths in input_artifacts.items():
            for fpath in fpaths:
                file_key = construct_artifact_key(fpath, field_name)
                if file_key in all_input_artifact_names.get(field_name, set()):
                    raise DuplicateInputArtifactsError(file_key, field_name)
                all_input_artifact_names.setdefault(field_name, set()).add(file_key)

        # finally, we upload the input artifacts to s3
        input_artifacts_s3_urls = None
        if at_least_one_input_artifact:
            input_artifacts_s3_urls, input_artifacts_s3_url_maps = (
                upload_input_artifacts(
                    input_artifacts,
                    s3_client,
                    self.storage_settings.BUCKET,
                    construct_artifact_key,
                )
            )
            # and then we update the specs with the new s3 urls
            for spec in spec_list:
                for field_name, fpath in spec._local_artifact_file_paths.items():
                    uri_map = input_artifacts_s3_url_maps[field_name]
                    uri = uri_map[fpath]
                    setattr(spec, field_name, uri)

        # next, we save the specs to a parquet file and upload it to s3
        df = pd.DataFrame([s.model_dump(mode="json") for s in spec_list])
        df_name = experiment_run.specs_filename

        uris = save_and_upload_parquets(
            collected_dfs={df_name: df},
            s3=s3_client,
            bucket=self.storage_settings.BUCKET,
            output_key_constructor=experiment_run.construct_specs_filekey,
        )
        specs_uri = uris[df_name]

        if isinstance(self.runnable, Workflow):
            if isinstance(specs, Sequence):
                msg = "Workflows do not yet support batching."
                raise TypeError(msg)
            run_ref = self.runnable.run_no_wait(
                specs,
                options=TriggerWorkflowOptions(
                    additional_metadata={
                        "experiment_id": experiment_run.experiment_id,
                        "experiment_name": self.runnable.name,
                    },
                ),
            )
        else:
            # Now, we can finally allocate the experiment to a workflow run
            # of the scatter/gather task
            scatter_gather_input = ScatterGatherInput(
                experiment_id=experiment_run.experiment_id,
                task_name=self.runnable.name,
                specs_uri=specs_uri,
                recursion_map=recursion_map
                or RecursionMap(path=None, factor=10, max_depth=0),
                storage_settings=self.storage_settings,
            )

            run_ref = scatter_gather.run_no_wait(
                scatter_gather_input,
                options=TriggerWorkflowOptions(
                    additional_metadata={
                        "experiment_id": experiment_run.experiment_id,
                        "experiment_name": self.runnable.name,  # TOD: should this be experiment run name?
                        "level": 0,
                    },
                ),
            )

        workflow_run_id = run_ref.workflow_run_id

        # Now, we can upload some various metadata
        # files to the run dir
        input_validator = self.runnable.input_validator_type
        if isinstance(self.runnable, Workflow):
            output_validator = NoneType
        else:
            output_validator = self.runnable.output_validator_type

        # if output_validator is None:
        #     msg = "Output validator is not set for experiment"
        #     raise ValueError(msg)

        class ExperimentIO(BaseModel):
            """The input and output schema for the experiment."""

            input: input_validator = Field(  # pyright: ignore [reportInvalidTypeForm]
                default=..., description="The input for the experiment."
            )
            output: output_validator = Field(  # pyright: ignore [reportInvalidTypeForm]
                default=..., description="The output for the experiment."
            )

        schema = ExperimentIO.model_json_schema()

        with tempfile.TemporaryDirectory() as temp_dir:
            tdir = Path(temp_dir)
            manifest_path = tdir / "manifest.yml"
            io_path = tdir / "experiment_io_spec.yml"
            input_artifacts_path = tdir / "input_artifacts.yml"
            workflow_spec_path = tdir / "workflow_spec.yml"

            # dump and upload the schema
            with open(io_path, "w") as f:
                yaml.dump(schema, f, indent=2, sort_keys=False)
            io_file_key = experiment_run.io_spec_filekey
            s3_client.upload_file(
                Bucket=self.storage_settings.BUCKET,
                Key=io_file_key,
                Filename=io_path.as_posix(),
            )

            # dump and upload the source files record
            input_artifacts_key = experiment_run.input_artifacts_filekey
            if input_artifacts_s3_urls:
                with open(input_artifacts_path, "w") as f:
                    yaml.dump(
                        input_artifacts_s3_urls.model_dump(mode="json"),
                        f,
                        indent=2,
                        sort_keys=False,
                    )
                s3_client.upload_file(
                    Bucket=self.storage_settings.BUCKET,
                    Key=input_artifacts_key,
                    Filename=input_artifacts_path.as_posix(),
                )

            manifest = experiment_run.construct_manifest(workflow_run_id)
            with open(manifest_path, "w") as f:
                yaml.dump(
                    manifest.model_dump(mode="json"), f, indent=2, sort_keys=False
                )
            manifest_file_key = experiment_run.manifest_filekey
            s3_client.upload_file(
                Bucket=self.storage_settings.BUCKET,
                Key=manifest_file_key,
                Filename=manifest_path.as_posix(),
            )
            if not isinstance(specs, Sequence):
                with open(workflow_spec_path, "w") as f:
                    yaml.dump(
                        specs.model_dump(mode="json"), f, indent=2, sort_keys=False
                    )
                s3_client.upload_file(
                    Bucket=self.storage_settings.BUCKET,
                    Key=experiment_run.workflow_spec_filekey,
                    Filename=workflow_spec_path.as_posix(),
                )

        return experiment_run, run_ref


class VersionedExperiment(BaseModel, Generic[TInput, TOutput]):
    """A versioned experiment."""

    base_experiment: BaseExperiment[TInput, TOutput]
    version: SemVer

    @property
    def base_id(self) -> str:
        """The base id for the versioned experiment."""
        return f"{self.base_experiment.base_id}/{self.version}"

    @property
    def prefix(self) -> str:
        """The prefix for the versioned experiment."""
        return f"{self.base_experiment.prefix}{self.version}/"

    def list_runs(self, s3_client: S3Client | None = None) -> list["ExperimentRun"]:
        """List all of the runs for the versioned experiment."""
        s3_client = s3_client or s3
        paginator = s3_client.get_paginator("list_objects_v2")

        common_prefixes: list[CommonPrefixTypeDef] = []
        for page in paginator.paginate(
            Bucket=self.base_experiment.storage_settings.BUCKET,
            Prefix=self.prefix,
            Delimiter="/",
            PaginationConfig={"PageSize": 1000},
        ):
            common_prefixes.extend(page.get("CommonPrefixes", []))

        prefixes = [d.get("Prefix") for d in common_prefixes]
        prefixes = [p.split("/")[-2] for p in prefixes if p is not None]
        timestamps = [datetime.strptime(p, DatetimeFormat) for p in prefixes]
        return [
            ExperimentRun(versioned_experiment=self, timestamp=t) for t in timestamps
        ]

    @property
    def latest_run(self) -> "ExperimentRun":
        """The latest run for the versioned experiment."""
        runs = self.list_runs()
        if not runs:
            raise ExperimentNotFoundError(self.base_experiment.runnable.name)
        return max(runs, key=lambda r: r.timestamp)

    @property
    def latest_run_results(self) -> dict[str, str]:
        """The latest run results for the versioned experiment."""
        latest_run = self.latest_run
        return latest_run.list_results_files()


class ExperimentRun(BaseModel, Generic[TInput, TOutput]):
    """An experiment run."""

    versioned_experiment: VersionedExperiment[TInput, TOutput]
    timestamp: datetime

    @property
    def dt_str(self) -> str:
        """The timestamp as a string."""
        return self.timestamp.strftime("%Y-%m-%d_%H-%M-%S")

    @property
    def experiment_id(self) -> str:
        """The base id for the experiment run."""
        return f"{self.versioned_experiment.base_id}/{self.dt_str}"

    @property
    def prefix(self) -> str:
        """The prefix for the run."""
        return f"{self.versioned_experiment.prefix}{self.dt_str}/"

    @property
    def artifact_prefix(self) -> str:
        """The prefix for the artifacts for the run."""
        return f"{self.prefix}artifacts/"

    def construct_artifact_key(self, field_name: str, file_name: str) -> str:
        """Construct the key for an artifact."""
        return f"{self.artifact_prefix}{field_name}/{file_name}"

    def construct_specs_filekey(self, filename: str) -> str:
        """Construct the key for the specs file."""
        return f"{self.prefix}{filename}.pq"

    @property
    def specs_filename(self) -> str:
        """The filename for the specs file."""
        return "specs"

    @property
    def specs_filekey(self) -> str:
        """The key for the specs file."""
        return self.construct_specs_filekey(self.specs_filename)

    @property
    def manifest_filekey(self) -> str:
        """The key for the manifest file."""
        return f"{self.prefix}manifest.yml"

    @property
    def io_spec_filekey(self) -> str:
        """The key for the io spec file."""
        return f"{self.prefix}experiment_io_spec.yml"

    @property
    def input_artifacts_filekey(self) -> str:
        """The key for the input artifacts file."""
        return f"{self.prefix}input_artifacts.yml"

    @property
    def workflow_spec_filekey(self) -> str:
        """The key for the workflow spec file."""
        return f"{self.prefix}workflow_spec.yml"

    def as_uri(self, key: str) -> S3Url:
        """Convert a key to a uri."""
        return S3Url(
            f"s3://{self.versioned_experiment.base_experiment.storage_settings.BUCKET}/{key}"
        )

    def construct_manifest(self, workflow_run_id: str) -> ExperimentRunManifest:
        """The manifest for the experiment run."""
        manifest = ExperimentRunManifest(
            workflow_run_id=workflow_run_id,
            experiment_id=self.experiment_id,
            experiment_name=self.versioned_experiment.base_experiment.runnable.name,
            io_spec=self.as_uri(self.io_spec_filekey),
            input_artifacts=self.as_uri(self.input_artifacts_filekey),
            specs_uri=self.as_uri(self.specs_filekey),
        )
        return manifest

    def overwrite_spec_meta(
        self,
        specs: Sequence[TInput],
        overwrite_experiment_id: bool = True,
        overwrite_sort_index: bool = True,
    ):
        """Overwrite the metadata for the specs."""
        for i, spec in enumerate(specs):
            if overwrite_experiment_id:
                spec.experiment_id = self.experiment_id
            if overwrite_sort_index:
                spec.sort_index = i

    @property
    def final_results_dirkey(self) -> str:
        """The key for the final results directory."""
        return f"{self.prefix}final/"

    @property
    def scalars_filekey(self) -> str:
        """The key for the scalars file."""
        return f"{self.final_results_dirkey}scalars.pq"

    def list_results_files(self, s3_client: S3Client | None = None) -> dict[str, str]:
        """List the results files for the experiment run."""
        s3_client = s3_client or s3
        paginator = s3_client.get_paginator("list_objects_v2")

        contents: list[ObjectTypeDef] = []
        for page in paginator.paginate(
            Bucket=self.versioned_experiment.base_experiment.storage_settings.BUCKET,
            Prefix=self.final_results_dirkey,
            Delimiter="/",
            PaginationConfig={"PageSize": 1000},
        ):
            contents.extend(page.get("Contents", []))

        contents = [c for c in contents if c.get("Key") is not None]
        keys = [c.get("Key") for c in contents]
        keys_safe = [k for k in keys if k and not k.endswith("/")]
        stems = [Path(k).stem for k in keys_safe]
        return dict(zip(stems, keys_safe, strict=True))


def upload_input_artifacts(
    input_artifacts: dict[str, set[FilePath]],
    s3_client: S3Client,
    bucket: str,
    construct_artifact_key: Callable[[FilePath, str], str],
) -> tuple[InputArtifactLocations, dict[str, dict[FilePath, S3Url]]]:
    """Upload source files to S3."""

    def handle_path(pth: FilePath, field_name: str):
        filekey = construct_artifact_key(pth, field_name)
        uri = S3Url(f"s3://{bucket}/{filekey}")
        s3_client.upload_file(
            Bucket=bucket,
            Key=filekey,
            Filename=pth.as_posix(),
        )
        return field_name, uri

    args: list[tuple[FilePath, str]] = []
    for field_name, paths in input_artifacts.items():
        for pth in paths:
            args.append((pth, field_name))
    with ThreadPoolExecutor(max_workers=10) as executor:
        first_args = [a[0] for a in args]
        second_args = [a[1] for a in args]
        results = list(
            tqdm(
                executor.map(handle_path, first_args, second_args),
                total=len(args),
                desc="Uploading source files",
            )
        )
    uris: dict[str, set[S3Url]] = {}
    for field_name, uri in results:
        uris.setdefault(field_name, set()).add(uri)
    uri_maps: dict[str, dict[FilePath, S3Url]] = {}
    for (field_name, uri), (pth, _) in zip(results, args, strict=True):
        uri_maps.setdefault(field_name, {})[pth] = uri
    return InputArtifactLocations(files=uris), uri_maps
