# Copyright lowRISC contributors (OpenTitan project).
# Licensed under the Apache License, Version 2.0, see LICENSE for details.
# SPDX-License-Identifier: Apache-2.0

"""Flow config base class."""

import json
import os
import pprint
import sys
from abc import ABC, abstractmethod
from collections.abc import Mapping, Sequence
from datetime import datetime, timezone
from pathlib import Path
from typing import TYPE_CHECKING, ClassVar

import hjson

from dvsim.flow.hjson import set_target_attribute
from dvsim.job.data import CompletedJobStatus
from dvsim.launcher.factory import get_launcher_cls
from dvsim.logging import log
from dvsim.report.data import FlowResults, IPMeta, ResultsSummary
from dvsim.report.generate import gen_block_report, gen_reports
from dvsim.scheduler import Scheduler
from dvsim.utils import (
    find_and_substitute_wildcards,
    rm_path,
    subst_wildcards,
)
from dvsim.utils.git import git_commit_hash

if TYPE_CHECKING:
    from dvsim.job.deploy import Deploy

__all__ = ("FlowCfg",)


# Interface class for extensions.
class FlowCfg(ABC):
    """Base class for the different flows supported by dvsim.py.

    The constructor expects some parsed hjson data. Create these objects with
    the factory function in CfgFactory.py, which loads the hjson data and picks
    a subclass of FlowCfg based on its contents.
    """

    # Set in subclasses. This is the key that must be used in an hjson file to
    # tell dvsim.py which subclass to use.
    flow = None

    # Can be overridden in subclasses to configure which wildcards to ignore
    # when expanding hjson.
    ignored_wildcards: ClassVar = []

    def __str__(self) -> str:
        """Get string representation of the flow config."""
        return pprint.pformat(self.__dict__)

    def __init__(self, flow_cfg_file, hjson_data, args, mk_config) -> None:
        # Options set from command line
        # Uniquify input items, while preserving the order.
        self.items = list(dict.fromkeys(args.items))
        self.list_items = args.list
        self.select_cfgs = args.select_cfgs
        self.flow_cfg_file = flow_cfg_file
        self.args = args
        self.scratch_root = args.scratch_root
        self.branch = args.branch
        self.job_prefix = args.job_prefix
        self.gui = args.gui
        self.gui_debug = args.gui_debug
        if self.gui_debug:
            self.gui = 1

        self.interactive = args.interactive

        # Options set from hjson cfg.
        self.project = ""
        self.scratch_path = ""
        self.scratch_base_path = ""

        # Add exports using 'exports' keyword - these are exported to the child
        # process' environment.
        self.exports = []

        # Add overrides using the overrides keyword - existing attributes
        # are overridden with the override values.
        self.overrides = []

        # List of cfgs if the parsed cfg is a primary cfg list
        self.cfgs: Sequence[FlowCfg] = []

        # Add a notion of "primary" cfg - this is indicated using
        # a special key 'use_cfgs' within the hjson cfg.
        self.is_primary_cfg = False

        # For a primary cfg, it is the aggregated list of all deploy objects
        # under self.cfgs. For a non-primary cfg, it is the list of items
        # slated for dispatch.
        self.deploy: Sequence[Deploy] = []

        # Timestamp
        self.timestamp_long = args.timestamp_long
        self.timestamp = args.timestamp

        # Results
        self.errors_seen: bool = False
        self.rel_path = ""
        self.results_title = ""
        self.revision = ""
        self.css_file = Path(__file__).resolve().parent / "style.css"
        # `self.results_*` below will be updated after `self.rel_path` and
        # `self.scratch_base_root` variables are updated.
        self.results_dir = ""
        self.results_page = ""
        self.results_html_name = ""

        # Full results in md text
        self.results_md = ""
        # Selectively sanitized md results to be published
        self.publish_results_md = ""
        self.sanitize_publish_results = False
        # Summary results, generated by over-arching cfg
        self.results_summary_md = ""

        # Merge in the values from the loaded hjson file. If subclasses want to
        # add other default parameters that depend on the parameters above,
        # they can override _merge_hjson and add their parameters at the start
        # of that.
        self._merge_hjson(hjson_data)

        # Is this a primary config? If so, we need to load up all the child
        # configurations at this point. If not, we place ourselves into
        # self.cfgs and consider ourselves a sort of "degenerate primary
        # configuration".
        self.is_primary_cfg = "use_cfgs" in hjson_data

        if not self.is_primary_cfg:
            self.cfgs.append(self)
        else:
            for entry in self.use_cfgs:
                self._load_child_cfg(entry, mk_config)

        if self.rel_path == "":
            self.rel_path = str(
                Path(self.flow_cfg_file).parent.relative_to(self.proj_root),
            )

        # Process overrides before substituting wildcards
        self._process_overrides()

        # Expand wildcards. If subclasses need to mess around with parameters
        # after merging the hjson but before expansion, they can override
        # _expand and add the code at the start.
        self._expand()

        # Construct the path variables after variable expansion.
        self.results_dir = Path(self.scratch_base_path) / "reports" / self.rel_path
        self.results_page = self.results_dir / self.results_html_name

        # Run any final checks
        self._post_init()

    def _merge_hjson(self, hjson_data: Mapping) -> None:
        """Take hjson data and merge it into self.__dict__.

        Subclasses that need to do something just before the merge should
        override this method and call super()._merge_hjson(..) at the end.

        """
        for key, value in hjson_data.items():
            set_target_attribute(self.flow_cfg_file, self.__dict__, key, value)

    def _expand(self) -> None:
        """Expand wildcards after merging hjson.

        Subclasses can override this to do something just before expansion.

        """
        # If this is a primary configuration, it doesn't matter if we don't
        # manage to expand everything.
        partial = self.is_primary_cfg

        # If custom dump script is exist, replace with run_script attribute.
        if self.args.dump_script is not None:
            self.run_script = "{proj_root}/" + self.args.dump_script

        self.__dict__ = find_and_substitute_wildcards(
            self.__dict__,
            self.__dict__,
            self.ignored_wildcards,
            ignore_error=partial,
        )

    def _post_init(self) -> None:
        # Run some post init checks
        if not self.is_primary_cfg:  # noqa: SIM102
            # Check if self.cfgs is a list of exactly 1 item (self)
            if not (len(self.cfgs) == 1 and self.cfgs[0].name == self.name):
                log.error("Parse error!\n%s", self.cfgs)
                sys.exit(1)

    def create_instance(self, mk_config, flow_cfg_file):
        """Create a new instance of this class for the given config file.

        mk_config is a factory method (passed explicitly to avoid a circular
        dependency between this file and CfgFactory.py).

        """
        new_instance = mk_config(flow_cfg_file)

        # Sanity check to make sure the new object is the same class as us: we
        # don't yet support heterogeneous primary configurations.
        if type(self) is not type(new_instance):
            log.error(
                f"{self.flow_cfg_file}: Loading child configuration at {flow_cfg_file!r}, but the "
                f"resulting flow types don't match: ({type(self).__name__} vs. {type(new_instance).__name__}).",
            )
            sys.exit(1)

        return new_instance

    def _load_child_cfg(self, entry, mk_config) -> None:
        """Load a child configuration for a primary cfg."""
        if type(entry) is str:
            # Treat this as a file entry. Substitute wildcards in cfg_file
            # files since we need to process them right away.
            cfg_file = subst_wildcards(entry, self.__dict__, ignore_error=True)
            self.cfgs.append(self.create_instance(mk_config, cfg_file))

        elif type(entry) is dict:
            # Treat this as a cfg expanded in-line
            temp_cfg_file = self._conv_inline_cfg_to_hjson(entry)
            if not temp_cfg_file:
                return
            self.cfgs.append(self.create_instance(mk_config, temp_cfg_file))

            # Delete the temp_cfg_file once the instance is created
            log.verbose("Deleting temp cfg file:\n%s", temp_cfg_file)
            rm_path(temp_cfg_file, ignore_error=True)

        else:
            log.error(
                'Type of entry "%s" in the "use_cfgs" key is invalid: %s',
                entry,
                str(type(entry)),
            )
            sys.exit(1)

    def _conv_inline_cfg_to_hjson(self, idict: Mapping) -> str | None:
        """Dump a temp hjson file in the scratch space from input dict.

        This method is to be called only by a primary cfg.
        """
        if not self.is_primary_cfg:
            log.fatal("This method can only be called by a primary cfg")
            sys.exit(1)

        name = idict.get("name", None)
        if not name:
            log.error(
                "In-line entry in use_cfgs list does not contain a "
                '"name" key (will be skipped!):\n%s',
                idict,
            )
            return None

        # Check if temp cfg file already exists
        temp_cfg_file = self.scratch_root + "/." + self.branch + "__" + name + "_cfg.hjson"

        # Create the file and dump the dict as hjson
        log.verbose('Dumping inline cfg "%s" in hjson to:\n%s', name, temp_cfg_file)

        try:
            Path(temp_cfg_file).write_text(hjson.dumps(idict, for_json=True))

        except Exception as e:
            log.exception(
                'Failed to hjson-dump temp cfg file"%s" for "%s"(will be skipped!) due to:\n%s',
                temp_cfg_file,
                name,
                e,
            )
            return None

        # Return the temp cfg file created
        return temp_cfg_file

    def _process_overrides(self) -> None:
        # Look through the dict and find available overrides.
        # If override is available, check if the type of the value for existing
        # and the overridden keys are the same.
        overrides_dict = {}
        if hasattr(self, "overrides"):
            overrides = self.overrides
            if type(overrides) is not list:
                log.error(
                    'The type of key "overrides" is %s - it should be a list',
                    type(overrides),
                )
                sys.exit(1)

            # Process override one by one
            for item in overrides:
                if type(item) is dict and set(item.keys()) == {"name", "value"}:
                    ov_name = item["name"]
                    ov_value = item["value"]
                    if ov_name not in overrides_dict:
                        overrides_dict[ov_name] = ov_value
                        self._do_override(ov_name, ov_value)
                    else:
                        log.error(
                            'Override for key "%s" already exists!\nOld: %s\nNew: %s',
                            ov_name,
                            overrides_dict[ov_name],
                            ov_value,
                        )
                        sys.exit(1)
                else:
                    log.error(
                        '"overrides" is a list of dicts with '
                        '{"name": <name>, "value": <value>} pairs. '
                        "Found this instead:\n%s",
                        str(item),
                    )
                    sys.exit(1)

    def _do_override(self, ov_name: str, ov_value: object) -> None:
        # Go through self attributes and replace with overrides
        if hasattr(self, ov_name):
            orig_value = getattr(self, ov_name)
            if isinstance(ov_value, type(orig_value)):
                log.debug('Overriding "%s" value "%s" with "%s"', ov_name, orig_value, ov_value)
                setattr(self, ov_name, ov_value)
            else:
                log.error(
                    'The type of override value "%s" for "%s" '
                    'doesn\'t match the type of original value "%s"',
                    ov_value,
                    ov_name,
                    orig_value,
                )
                sys.exit(1)
        else:
            log.error('Override key "%s" not found in the cfg!', ov_name)
            sys.exit(1)

    @abstractmethod
    def _purge(self) -> None:
        """Purge the existing scratch areas in preparation for the new run."""

    def purge(self) -> None:
        """Public facing API for _purge()."""
        for item in self.cfgs:
            item._purge()

    @abstractmethod
    def _print_list(self) -> None:
        """Print the list of available items that can be kicked off."""

    def print_list(self) -> None:
        """Public facing API for _print_list()."""
        for item in self.cfgs:
            item._print_list()

    def prune_selected_cfgs(self) -> None:
        """Prune the list of configs for a primary config file."""
        # This should run after self.cfgs has been set
        assert self.cfgs

        # If the user didn't pass --select-cfgs, we don't do anything.
        if self.select_cfgs is None:
            return

        # If the user passed --select-cfgs, but this isn't a primary config
        # file, we should probably complain.
        if not self.is_primary_cfg:
            log.error(
                f"The configuration file at {self.flow_cfg_file!r} is not a primary "
                "config, but --select-cfgs was passed on the command "
                "line.",
            )
            sys.exit(1)

        # Filter configurations
        self.cfgs = [c for c in self.cfgs if c.name in self.select_cfgs]

    @abstractmethod
    def _create_deploy_objects(self) -> None:
        """Create deploy objects from items that were passed on for being run.

        The deploy objects for build and run are created from the objects that
        were created from the create_objects() method.
        """

    def create_deploy_objects(self) -> None:
        """Public facing API for _create_deploy_objects()."""
        self.prune_selected_cfgs()

        # GUI, GUI debug or Interactive mode is allowed only for one cfg.
        if (self.gui or self.gui_debug or self.interactive) and len(self.cfgs) > 1:
            log.fatal("In GUI mode, only one cfg can be run.")
            sys.exit(1)

        for item in self.cfgs:
            item._create_deploy_objects()

    def deploy_objects(self) -> Sequence[CompletedJobStatus]:
        """Public facing API for deploying all available objects.

        Runs each job and returns a map from item to status.
        """
        deploy = []
        for item in self.cfgs:
            deploy.extend(item.deploy)

        if not deploy:
            log.error("Nothing to run!")
            sys.exit(1)

        jobs = [d.get_job_spec() for d in deploy]

        if os.environ.get("DVSIM_DEPLOY_DUMP", "true"):
            filename = f"deploy_{self.branch}_{self.timestamp}.json"
            (Path(self.scratch_root) / filename).write_text(
                json.dumps(
                    # Sort on full name to ensure consistent ordering
                    sorted(
                        [
                            j.model_dump(
                                # callback functions can't be serialised
                                exclude={"pre_launch", "post_finish"},
                                mode="json",
                            )
                            for j in jobs
                        ],
                        key=lambda j: j["full_name"],
                    ),
                    indent=2,
                ),
            )

        return Scheduler(
            items=jobs,
            launcher_cls=get_launcher_cls(),
            interactive=self.interactive,
        ).run()

    def gen_results(self, results: Sequence[CompletedJobStatus]) -> None:
        """Generate flow results.

        Args:
            results: completed job status objects.

        """
        reports_dir = Path(self.scratch_base_path) / "reports"
        commit = git_commit_hash(path=Path(self.proj_root))
        url = f"https://github.com/lowrisc/opentitan/tree/{commit}"

        all_flow_results: Mapping[str, FlowResults] = {}

        for item in self.cfgs:
            item_results = [
                res
                for res in results
                if res.block.name == item.name and res.block.variant == item.variant
            ]

            flow_results: FlowResults = item._gen_json_results(
                run_results=item_results,
                commit=commit,
                url=url,
            )

            # Convert to lowercase to match filename
            block_result_index = (
                f"{item.name}_{item.variant}" if item.variant else item.name
            ).lower()

            all_flow_results[block_result_index] = flow_results

            # Generate the block's JSON/HTML reports to the report area.
            gen_block_report(
                results=flow_results,
                path=reports_dir,
            )

            self.errors_seen |= item.errors_seen

        if self.is_primary_cfg:
            # The timestamp for this run has been taken with `utcnow()` and is
            # stored in a custom format.  Store it in standard ISO format with
            # explicit timezone annotation.
            timestamp = (
                datetime.strptime(self.timestamp, "%Y%m%d_%H%M%S")
                .replace(tzinfo=timezone.utc)
                .isoformat()
            )

            results_summary = ResultsSummary(
                top=IPMeta(
                    name=self.name,
                    variant=self.variant,
                    commit=commit,
                    branch=self.branch,
                    url=url,
                ),
                timestamp=timestamp,
                flow_results=all_flow_results,
                report_path=reports_dir,
            )

            # Generate all the JSON/HTML reports to the report area.
            gen_reports(
                summary=results_summary,
                path=reports_dir,
            )

    def has_errors(self) -> bool:
        """Return error state."""
        return self.errors_seen
