import logging
import pickle
from pathlib import Path
from typing import Any

from hydra.types import TaskFunction
from omegaconf import DictConfig

from hydra.core.utils import JobReturn, JobStatus
from hydra.experimental.callback import Callback

def reorder_path(s: str) -> str:
    kvs = s.split(",")
    newkv = []
    for kv in kvs:
        newkv.append(kv.split("="))
    newkv = sorted(newkv, key=lambda x: (len(x[1]), x[0]))
    return ",".join("=".join(kv) for kv in newkv)

def replace_slash(s: str) -> str:
    return s.replace("/", ".-")

class RewriteJobDirCallback(Callback):
    """Rerwite the override_dirname"""

    def __init__(self) -> None:
        pass

    def on_run_start(self, config: DictConfig, **kwargs: Any) -> None:
        config.hydra.job.override_dirname = replace_slash(
            reorder_path(config.hydra.job.override_dirname)
        )
        return super().on_run_start(config, **kwargs)