librelane

# Copyright 2025 LibreLane Contributors
#
# Adapted from OpenLane 2
#
# Copyright 2023 Efabless Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import sys
import json
from functools import partial, wraps
from concurrent.futures import ThreadPoolExecutor
from typing import Optional, Union, Tuple

from click import (
    Context,
    Parameter,
    echo,
)
from click.formatting import join_options
from cloup import (
    option,
    argument,
    option_group,
    Choice,
    Argument,
    Option as CloupOption,
    Path as CloupPath,
)
from cloup.constraints import (
    mutually_exclusive,
)
from cloup.typing import Decorator

from .flow import Flow
from ..common import set_tpe, cli, get_pdk_hash, _get_process_limit
from ..logging import set_log_level, verbose, err, options, LogLevels
from ..state import State, InvalidState


class Option(CloupOption):
    """
    A slight modification of cloup.Option that consumes the environment
    variable(s) in envvar upon use.
    """

    def resolve_envvar_value(self, ctx: Context) -> Optional[str]:
        if self.envvar is None:
            return None
        evs = self.envvar
        if isinstance(evs, str):
            evs = [evs]
        for envvar in self.envvar:
            rv = os.environ.pop(envvar, None)
            if rv:
                return rv
        return None


class Path(CloupPath):
    """
    A modification of cloup.Path that rejects paths starting with a tilde (~)
    as too ambiguous. This is because of user confusion.
    """

    def convert(
        self,
        value: Union[str, os.PathLike],
        param: Optional[Parameter],
        ctx: Optional[Context],
    ):
        value = str(value)
        assert param is not None
        assert ctx is not None

        is_cmd = False
        if sys.platform == "win32":
            import psutil

            is_cmd = psutil.Process(os.getppid()).name().lower().endswith("cmd.exe")
        if not is_cmd and value.startswith("~"):
            usage, _ = join_options(param.opts)
            if isinstance(param, Argument):
                usage = ""
            buffer = f"'{value}' starts with a tilde, which is ambiguous.\n"
            buffer += "  * If you meant your home directory, make sure your POSIX shell is able to expand the tilde:\n"
            buffer += f"    * GOOD: {ctx.command_path} {usage} {value}   …\n"
            if usage != "":
                buffer += f"    * BAD:  {ctx.command_path} {usage}={value}   …\n"
            buffer += f'    * BAD:  {ctx.command_path} {usage} "{value}" …\n'
            env_vars = []
            if param.envvar is not None:
                if isinstance(param.envvar, str):
                    env_vars = [param.envvar]
                else:
                    env_vars = list(param.envvar)
            for var in env_vars:
                buffer += f"    * GOOD: {var}={value}   {ctx.command_path} …\n"
                buffer += f'    * BAD:  {var}="{value}" {ctx.command_path} …\n'
            buffer += '  * If you want a relative file or directory that starts with a literal "~", use an absolute path.\n'
            self.fail(buffer, param, ctx)
        return super().convert(value, param, ctx)


def set_log_level_cb(
    ctx: Context,
    param: Parameter,
    value: Optional[str],
):
    if value is None:
        return

    level: Union[str, int] = value
    try:
        try:
            level = int(value)
        except ValueError:
            pass
        set_log_level(level)
    except ValueError as e:
        err(f"Invalid logging level {value}: {e}.")
        echo(ctx.get_help())
        ctx.exit(-1)


def set_worker_count_cb(
    ctx: Context,
    param: Parameter,
    value: Optional[int],
):
    if value is None:
        return None

    set_tpe(ThreadPoolExecutor(max_workers=value))


def initial_state_cb(
    ctx: Context,
    param: Parameter,
    value: Tuple[str],
):
    if len(value) == 0:
        return None

    raw = {}
    for state_json in value:
        try:
            with open(state_json, encoding="utf8") as f:
                state_dict = json.load(f)
                if not isinstance(state_dict, dict):
                    raise ValueError(f"JSON data {value} is not a dictionary")
                raw.update(state_dict)
        except json.JSONDecodeError as e:
            err(f"Invalid JSON file: {e}")
            ctx.exit(-1)
        except Exception as e:
            err(f"Failed to read initial state: {e}")
            ctx.exit(-1)
    try:
        initial_state = State.load(raw, validate_path=True)
    except InvalidState as e:
        err(e)
        ctx.exit(-1)

    return initial_state


def only_cb(
    ctx: Context,
    param: Parameter,
    value: Optional[str],
):
    if value is not None:
        ctx.obj = ctx.obj or {}
        ctx.obj["only"] = value
    return value


def from_to_cb(
    ctx: Context,
    param: Parameter,
    value: Optional[str],
):
    if isinstance(ctx.obj, dict) and ctx.obj.get("only"):
        return ctx.obj.get("only")
    return value


def condensed_cb(ctx: Context, param: Parameter, value: bool):
    if value:
        options.set_condensed_mode(True)
        options.set_show_progress_bar(False)


def progressbar_cb(ctx: Context, param: Parameter, value: Optional[bool]):
    if value is not None:
        options.set_show_progress_bar(value)


def cloup_flow_opts(
    *,
    config_options: bool = True,
    run_options: bool = True,
    sequential_flow_controls: bool = True,
    sequential_flow_reproducible: bool = False,
    pdk_options: bool = True,
    log_level: bool = True,
    jobs: bool = True,
    accept_config_files: bool = True,
    volare_by_default: bool = True,
    volare_pdk_override: Optional[str] = None,
    _enable_debug_flags: bool = False,
    enable_overwrite_flag: bool = False,
    enable_initial_state_element: bool = False,
) -> Decorator:
    """
    Creates a wrapper that appends a number of LibreLane flow-related flags to a
    function decorated with @cloup.command (https://cloup.readthedocs.io/en/stable/autoapi/cloup/index.html#cloup.command).

    The following keyword arguments will be passed to the decorated function.

    * Those postfixed ‡ are compatible with the constructor for :class:`Flow`.
    * Those postfixed § are compatible with the :meth:`Flow.start`.

    ---

    * Flow configuration options (if parameter ``config_options`` is ``True``):
        * ``flow_name``: ``Optional[str]``: A valid flow ID to be used with :meth:`Flow.factory.get`
        * ``config_override_strings`` ‡: ``Optional[Iterable[str]]``
    * Sequential flow controls (if parameter ``sequential_flow_controls`` is ``True``)
        * ``frm`` §: ``Optional[str]``: Start from a step with this ID. Supported by sequential flows.
        * ``to`` §: ``Optional[str]``: Stop at a step with this id. Supported by sequential flows.
        * ``skip`` §: ``Iterable[str]``: Skip these steps. Supported by sequential flows.
    * Sequential flow reproducible (if parameter ``sequential_flow_reproducible`` is ``True``)
        * ``reproducible`` §: ``str``: Create a reproducible for a step with is ID, aborting the flow afterwards. Supported by sequential flows.
    * Flow run options (if parameter ``run_options`` is ``True``):
        * ``tag`` §: ``Optional[str]``
        * ``last_run`` §: ``bool``: If ``True``, ``tag`` is guaranteed to be None.
        * ``with_initial_state`` §: ``Optional[State]``
    * PDK options
        * ``use_volare`` : ``bool``
        * ``pdk_root`` ‡: ``Optional[str]``
        * ``pdk`` ‡: ``str``
        * ``scl`` ‡: ``Optional[str]``
        * ``pad`` ‡: ``Optional[str]``
    * ``config_files``: ``Iterable[str]``: Paths to configuration files (if
      parameter  ``accept_config_files`` is ``True``)

    :param config_options: Enables flow configuration and starting CLI flags
    :param sequential_flow_controls: Enables flow control CLI flags
    :param flow_run_options: Enables tag CLI flags
    :param pdk_options: Enables PDK CLI flags
    :param log_level: Enables ``--log-level`` CLI flag
    :param jobs: Enables ``-j/--jobs`` CLI flag
    :param accept_config_files: Accepts configuration file paths as CLI arguments
    :param volare_by_default: If ``pdk_options`` is ``True``, this changes whether
        Ciel is used by default for this CLI or not.
    :returns: The wrapper
    """
    o = partial(option, cls=CloupOption, show_default=True)

    def decorate(f):
        if config_options:
            f = option_group(
                "Flow configuration options",
                o(
                    "-f",
                    "--flow",
                    "flow_name",
                    type=Choice(Flow.factory.list(), case_sensitive=False),
                    default=None,
                    help="The built-in LibreLane flow to use for this run",
                ),
                o(
                    "-c",
                    "--override-config",
                    "config_override_strings",
                    type=str,
                    multiple=True,
                    help="For this run only- override a configuration variable with a certain value. In the format KEY=VALUE. Can be specified multiple times. Values must be valid JSON values, and keys must not use their deprecated names.",
                ),
            )(f)
        if run_options:
            f = o(
                "-i",
                "--with-initial-state",
                type=Path(
                    exists=True,
                    file_okay=True,
                    dir_okay=False,
                ),
                multiple=True,
                callback=initial_state_cb,
                help="Use these JSON files as an initial state. If multiple are provided, they are merged with keys in later files overriding keys in prior files. If none are specified, the latest `state_out.json` of the run directory will be used. If none exist, an empty initial state is created.",
            )(f)
            f = o(
                "--design-dir",
                "design_dir",
                type=Path(
                    exists=True,
                    file_okay=False,
                    dir_okay=True,
                ),
                default=None,
                help="The top-level directory for your design that configuration objects may resolve paths relative to.",
            )(f)
            if enable_overwrite_flag:
                f = o(
                    "--overwrite",
                    is_flag=True,
                    default=False,
                    help="Overwrite run, if exists.",
                )(f)
            if _enable_debug_flags:
                f = option_group(
                    "Debug flags",
                    o(
                        "--force-run-dir",
                        "_force_run_dir",
                        type=Path(
                            exists=True,
                            file_okay=False,
                            dir_okay=True,
                        ),
                        hidden=True,
                        default=None,
                    ),
                )(f)
            f = option_group(
                "Run options",
                o(
                    "--run-tag",
                    "tag",
                    default=None,
                    type=str,
                    help="An optional name to use for this particular run of an LibreLane-based flow. Used to create the run directory.",
                ),
                o(
                    "--last-run",
                    is_flag=True,
                    default=False,
                    help="Use the last run as the run tag.",
                ),
                constraint=mutually_exclusive,
            )(f)
        if sequential_flow_controls:
            f = option_group(
                "Sequential flow controls",
                o(
                    "-F",
                    "--from",
                    "frm",
                    type=str,
                    default=None,
                    callback=from_to_cb,
                    help="Start from a step with this id. Supported by sequential flows.",
                ),
                o(
                    "-T",
                    "--to",
                    type=str,
                    default=None,
                    callback=from_to_cb,
                    help="Stop at a step with this id. Supported by sequential flows.",
                ),
                o(
                    "--only",
                    type=str,
                    default=None,
                    expose_value=False,
                    is_eager=True,
                    callback=only_cb,
                    help="Shorthand to set both --from and --to to the same value. Overrides the values from both.",
                ),
                o(
                    "-S",
                    "--skip",
                    type=str,
                    multiple=True,
                    help="Skip these steps. Supported by sequential flows.",
                ),
            )(f)
        if sequential_flow_reproducible:
            f = o(
                "--reproducible",
                type=str,
                help="Create a reproducible for the step matching this ID, then abort the flow. Supported by sequential flows.",
            )(f)
        if log_level:
            f = o(
                "--log-level",
                type=cli.IntEnumChoice(LogLevels),
                default=None,
                help="A logging level. Set to VERBOSE or higher to silence subprocess logs. [default: unchanged from SUBPROCESS]",
                callback=set_log_level_cb,
                expose_value=False,
                show_default=False,
            )(f)
            f = o(
                "--show-progress-bar/--hide-progress-bar",
                type=bool,
                help="Whether to show the progress bar when running Flows. [default: show]",
                default=None,
                callback=progressbar_cb,
                expose_value=False,
            )(f)
            f = o(
                "--condensed/--full",
                type=bool,
                help="In condensed mode, subprocess logs are suppressed regardless of step, --hide-progress-bar is the default, and the log messages themselves are a bit more terse. Useful for debugging.",
                default=False,
                is_eager=True,
                callback=condensed_cb,
                expose_value=False,
            )(f)
        if pdk_options:
            f = option_group(
                "PDK options",
                o(
                    "--volare-pdk/--manual-pdk",
                    "--ciel-pdk/--manual-pdk",
                    "use_ciel",
                    is_eager=True,
                    default=volare_by_default,
                    help="Automatically use Ciel for PDK version installation and enablement. Set --manual if you want to use a custom PDK version.",
                ),
                o(
                    "--pdk-root",
                    type=Path(
                        file_okay=False,
                        dir_okay=True,
                    ),
                    is_eager=True,
                    envvar=["PDK_ROOT"],
                    help="Override Ciel PDK root folder. Required if Ciel is not installed, but a default value can also be set via the environment variable PDK_ROOT.",
                ),
                o(
                    "-p",
                    "--pdk",
                    type=str,
                    envvar=["PDK"],
                    default="sky130A",
                    help="The process design kit to use.",
                ),
                o(
                    "-s",
                    "--scl",
                    type=str,
                    envvar=["STD_CELL_LIBRARY"],
                    # no default, default is obtained dynamically from PDK
                    help="The standard cell library to use. If None, the PDK's default standard cell library is used.",
                ),
                o(
                    "--pad",
                    type=str,
                    envvar=["PAD_CELL_LIBRARY"],
                    # no default, default is obtained dynamically from PDK
                    help="The standard pad library to use. If None, the PDK's default standard cell library is used (if it exists).",
                ),
            )(f)
        if jobs:
            f = o(
                "-j",
                "--jobs",
                type=int,
                default=_get_process_limit(),
                help="The maximum number of threads or processes that can be used by LibreLane.",
                callback=set_worker_count_cb,
                expose_value=False,
            )(f)
        if enable_initial_state_element:
            f = o(
                "-e",
                "--initial-state-element-override",
                type=str,
                multiple=True,
                default=(),
                help="Elements to override in the used initial state in the format DESIGN_FORMAT_ID=PATH",
            )(f)
        if accept_config_files:
            f = argument(
                "config_files",
                nargs=-1,
                type=Path(
                    exists=True,
                    file_okay=True,
                    dir_okay=True,
                ),
            )(f)
        if pdk_options:

            @wraps(f)
            def pdk_resolve_wrapper(
                *args,
                pdk_root: Optional[str],
                pdk: str,
                scl: Optional[str],
                pad: Optional[str],
                use_ciel: bool,
                **kwargs,
            ) -> str:
                if not use_ciel:
                    if pdk_root is None:
                        err("Argument --pdk-root must be present with --manual-pdk.")
                        exit(1)
                else:
                    import ciel
                    from ciel.source import StaticWebDataSource

                    opdks_rev = volare_pdk_override or get_pdk_hash(pdk)
                    ciel_home = ciel.get_ciel_home(pdk_root)

                    include_libraries = ["default"]
                    if scl is not None:
                        include_libraries.append(scl)

                    if pad is not None:
                        include_libraries.append(pad)

                    pdk_family = None
                    if family := ciel.Family.by_name.get(pdk):
                        pdk = family.default_variant
                        pdk_family = family.name
                        verbose(f"Resolved PDK variant {family.default_variant}.")
                    else:
                        for family in ciel.Family.by_name.values():
                            if pdk in family.variants:
                                pdk_family = family.name
                                break

                    if pdk_family is None:
                        err(f"Could not resolve the PDK '{pdk}'.")
                        exit(1)

                    try:
                        version = ciel.fetch(
                            ciel_home,
                            pdk_family,
                            opdks_rev,
                            data_source=StaticWebDataSource(
                                "https://fossi-foundation.github.io/ciel-releases"
                            ),
                            include_libraries=include_libraries,
                        )
                        pdk_root = version.get_dir(ciel_home)
                    except ValueError as e:
                        err(f"Failed to download PDK: {e}")
                        exit(1)

                return f(*args, pdk_root=pdk_root, pdk=pdk, scl=scl, pad=pad, **kwargs)

            return pdk_resolve_wrapper
        else:
            return f

    return decorate