from datetime import datetime
import pathlib
from piel.config import piel_path_types
from piel.file_system import return_path
from typing import Literal

__all__ = [
    "extract_datetime_from_path",
    "find_all_design_runs",
    "find_latest_design_run",
    "get_gds_path_from_design_run",
    "get_design_run_version",
    "sort_design_runs",
]


def extract_datetime_from_path(run_path: pathlib.Path) -> str:
    """
    Extracts the datetime from a given `run_path` and returns it as a string.
    """
    run_str = run_path.name.replace("RUN_", "")
    if "." in run_str.split("_")[0]:
        date_format = "%Y.%m.%d_%H.%M.%S"
    else:
        date_format = "%Y-%m-%d_%H-%M-%S"
    return datetime.strptime(run_str, date_format)


def find_all_design_runs(
    design_directory: piel_path_types,
    run_name: str | None = None,
) -> list[pathlib.Path]:
    """
    For a given `design_directory`, the `openlane` output can be found in the `runs` subdirectory. This function sorts the runs according to the default notations between both `openlane` and `openlane2` run formats.

    If a `run_name` is specified, then the function will return the exact run if it exists. Otherwise, it will return the latest run

    Args:
        design_directory (piel_path_types): The path to the design directory
        run_name (str, optional): The name of the run to return. Defaults to None.
        version (Literal["v1", "v2"], optional): The version of OpenLane to use. Defaults to None.

    Raises:
        ValueError: If the run_name is specified but not found in the design_directory

    Returns:
        list[pathlib.Path]: A list of pathlib.Path objects corresponding to the runs
    """
    design_directory = return_path(design_directory, as_piel_module=True)
    runs_design_directory = design_directory / "runs"
    # Convert to path so that it can be found and compared within design_directory
    all_runs_list = list(runs_design_directory.iterdir())
    if run_name is not None:
        run_path = runs_design_directory / run_name
        # Return this exact run
        if run_path in all_runs_list:
            # Check that the run exists
            pass
        else:
            raise ValueError(
                "Run: " + str(run_path) + "not found in " + str(all_runs_list)
            )
    else:
        # Take the latest design run if it exists
        if len(all_runs_list) > 0:
            sorted_runs_per_version = sort_design_runs(all_runs_list)
        else:
            # If there are no runs
            raise ValueError(
                "No OpenLane design runs were found in: " + str(runs_design_directory)
            )
    return sorted_runs_per_version


def find_latest_design_run(
    design_directory: piel_path_types,
    run_name: str | None = None,
    version: Literal["v1", "v2"] | None = None,
) -> (pathlib.Path, str):
    """
    For a given `design_directory`, the `openlane` output can be found in the `runs` subdirectory. This function sorts the runs according to the default notations between both `openlane` and `openlane2` run formats.

    If a `run_name` is specified, then the function will return the exact run if it exists. Otherwise, it will return the latest run.

    Args:
        design_directory (piel_path_types): The path to the design directory
        run_name (str, optional): The name of the run to return. Defaults to None.
        version (Literal["v1", "v2"], optional): The version of the run to return. Defaults to None.

    Raises:
        ValueError: If the run_name is specified but not found in the design_directory

    Returns:
        (pathlib.Path, str): A tuple of the latest run path and the version
    """
    latest_path = None
    latest_datetime = None
    latest_version = None

    sorted_runs_per_version = find_all_design_runs(
        design_directory=design_directory, run_name=run_name
    )

    # If a specific version is specified
    if version:
        filtered_paths = sorted_runs_per_version.get(version, [])
        filtered_versions = [version] * len(filtered_paths)
    else:
        filtered_paths = []
        filtered_versions = []
        for v, path_list in sorted_runs_per_version.items():
            filtered_paths.extend(path_list)
            filtered_versions.extend([v] * len(path_list))

    for path, v in zip(filtered_paths, filtered_versions, strict=True):
        path_datetime = extract_datetime_from_path(path)
        if not latest_datetime or path_datetime > latest_datetime:
            latest_datetime = path_datetime
            latest_path = path
            latest_version = v

    return latest_path, latest_version


def get_gds_path_from_design_run(
    design_directory: piel_path_types,
    run_directory: piel_path_types | None = None,
) -> pathlib.Path:
    """
    Returns the path to the final GDS generated by OpenLane.

    Args:
        design_directory (piel_path_types): The path to the design directory
        run_directory (piel_path_types, optional): The path to the run directory. Defaults to None. Otherwise gets the latest run.

    Returns:
        pathlib.Path: The path to the final GDS
    """
    design_directory = return_path(design_directory)
    if run_directory is None:
        run_directory, run_version = find_latest_design_run(design_directory)
    else:
        run_directory = return_path(run_directory)
        run_version = get_design_run_version(run_directory)

    if run_version == "v1":
        final_gds_run = run_directory / "results" / "final" / "gds"
    elif run_version == "v2":
        final_gds_run = run_directory / "final" / "gds"

    final_gds = list(final_gds_run.glob("*.gds"))
    if len(final_gds) == 0:
        raise ValueError(
            "No final GDS was found in: "
            + str(final_gds_run)
            + " for "
            + str(run_directory)
        )
    elif len(final_gds) > 1:
        raise ValueError(
            "Multiple final GDS were found in: "
            + str(final_gds_run)
            + " for "
            + str(run_directory)
        )
    else:
        return final_gds[0]


def get_design_run_version(
    run_directory: piel_path_types,
) -> Literal["v1", "v2"]:
    """
    Returns the version of the design run.
    """
    run_str = run_directory.name.replace("RUN_", "")
    version = "v1" if "." in run_str.split("_")[0] else "v2"
    return version


def sort_design_runs(
    path_list: list[pathlib.Path],
) -> dict[str, list[pathlib.Path]]:
    """
    For a given `design_directory`, the `openlane` output can be found in the `runs` subdirectory. This function sorts the runs according to the default notations between both `openlane` and `openlane2` run formats.

    Args:
        path_list (list[pathlib.Path]): A list of pathlib.Path objects corresponding to the runs

    Returns:
        dict[str, list[pathlib.Path]]: A dictionary of sorted runs
    """
    # Initialize a dictionary to hold sorted PosixPaths categorized as v1 and v2
    sorted_paths = {"v2": list(), "v1": list()}

    for path in path_list:
        version = get_design_run_version(path)
        # Parsing datetime to make it sortable
        datetime_obj = extract_datetime_from_path(path)

        sorted_paths[version].append((datetime_obj, path))

    # Sort the paths by datetime
    for version, paths in sorted_paths.items():
        sorted_paths[version] = [x[1] for x in sorted(paths)]

    # Return the sorted paths
    return sorted_paths
