"""
Utility functions using slang for SystemVerilog analysis.

This is to avoid incompatiblity with pyslang (conflicts with some other
EDA tool installs) and it generally lagging behind slang releases.
"""

import json
import os
from pathlib import Path
import subprocess

from opencos import util
from opencos.commands import sim
from opencos.files import safe_shutil_which
from opencos.utils.str_helpers import sanitize_defines_for_sh


INSTALL_HELPER_STR = '''
Install "slang" with:
       git clone https://github.com/MikePopoloski/slang
       cd slang
       cmake -B build
       cmake --build build -j
       cd build
       sudo make install
'''

def warn_slang_not_available(extra_text: str = '') -> None:
    '''Warning if slang isn't available on this system'''
    if 'slang_not_installed' in util.AlreadyMessage.warn_names:
        return
    util.AlreadyMessage.warn_names.add('slang_not_installed')
    util.warning(f'External tool "slang" is required {extra_text}')
    for line in INSTALL_HELPER_STR.split('\n'):
        util.warning(line)


class SlangAvailable:
    '''
    Helper to report availability of 'slang', via PATH or caller supplied EXE

    Used by opencos.tools.slang.ToolSlang
    '''

    available = None
    version: str = ''
    exe: str = ''
    base_path: str = ''

    def init(self, exe: str = '', quiet: bool = False) -> None:
        '''Sets up class members, outside users can get at this via:
            > slang_helpers.available.available

        Does not run if it has already been run and exe arg is empty, which allows
        slang.ToolSlang.get_versions to re-run init() if using a specific exe path.
        '''
        if self.available is not None:
            if not exe:
                return # already set, on default exe
        if exe:
            assert os.path.split(exe)[-1] == 'slang', f'Bad {exe=} arg to init(..)'
        else:
            exe = 'slang'

        # members not set, or set but someone BYO'd exe
        self.exe = safe_shutil_which(exe)
        self.base_path, _ = os.path.split(self.exe)
        self.available = bool(self.exe)
        if self.available:
            self.version = self._get_version(self.exe, quiet=quiet)

    def _get_version(self, exe: str = '', quiet: bool = False) -> str:
        '''Gets the version from: slang --version, called by init()'''

        if not exe:
            exe = self.exe
        if not exe:
            return ''
        version_ret = subprocess.run([exe, '--version'], capture_output=True, check=False)
        stdout = version_ret.stdout.decode('utf-8', errors='replace')
        words = stdout.split() # slang version 8.0.6+b4a74b00
        if len(words) < 3:
            if not quiet:
                util.warn_once(f'{exe} --version: returned unexpected string {version_ret=}',
                               warn_name=f'slang_version__{exe}')
            return ''
        version = words[2]
        left, _ = version.split('+')
        ver_list = left.split('.')
        if len(ver_list) != 3:
            if not quiet:
                util.warn_once(f'{exe} --version: returned unexpected string',
                               f'{version_ret=} {version=}',
                               warn_name=f'slang_version__{exe}')
            return ''
        return left


# Have a global version holding available, exe, version:
slang_available = SlangAvailable()

def check_slang_available() -> bool:
    '''Returns True is slang is available for use'''
    if slang_available.available is None:
        slang_available.init()
    return slang_available.available


def slang_read_ast_json(filepath: str) -> dict:
    '''Get data from AST'''
    ret = {}
    with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
        ret = json.load(f)
    return ret


class SlangAst:
    '''Methods for getting data from an existing AST slang.json file, or generating

    a new one via subprocess'''

    hier_sep: str = '/'
    uvm_home: str = os.environ.get('UVM_HOME', '') # best effort path to location of uvm_pkg.sv
    needs_uvm: bool = False # default
    handler_obj = None
    threads: int = 1 # we try to get from the handler_obj if possible

    slang_command_args = [
        '--timescale=1ns/1ns', '--single-unit', '--allow-dup-initial-drivers',
        '--ast-json-source-info', '--ast-json-detailed-types',
        '-Wno-explicit-static', '-Wno-missing-top',
    ]
    slang_json_fname = 'slang.json'
    slang_sh_fname = 'run_slang_for_ast.sh'
    work_dir = ''
    timeout: int = 180 # seconds, 3 min for slang to run upper-bound? Usually its 1sec or less

    def __init__(self, handler_obj: object,
                 filepath: str, hier_sep: str = '', uvm_home: str = ''):
        self.handler_obj = handler_obj

        if hier_sep:
            self.hier_sep = hier_sep

        self.ast_data = None
        self.root_node = None
        self.all_items = [] # clear these

        # deal with UVM:
        self.set_uvm_and_uvm_home(needs_uvm=True, uvm_home=uvm_home)

        # load AST if provided:
        self.init_ast_filepath(filepath)

        # inherit the max threads from the handler obj:
        self.threads = getattr(handler_obj, 'threads',
                               getattr(handler_obj, 'args', {}).get('threads', 1))


    def init_ast_filepath(self, filepath: str) -> None:
        '''Call if you set your constructor with filepath='' or None,

        sets self.ast_data, self.root_node; clears self.all_items.
        '''
        if not filepath:
            return

        self.ast_data = slang_read_ast_json(filepath)
        self.root_node = self.ast_data.get('design', {})
        self.all_items = []


    def set_uvm_and_uvm_home(self, needs_uvm: bool, uvm_home: str) -> None:
        '''Sets self.needs_uvm, overwrites self.uvm_home'''
        if not needs_uvm or not uvm_home:
            # do not update, but clear self.needs_uvm
            self.needs_uvm = False
            return


        # deal with Windows:
        p = Path(uvm_home)
        if p.is_dir() and (p / 'uvm_pkg.sv').is_file():
            self.needs_uvm = True
            self.uvm_home = str(p.resolve())
            return

        self.needs_uvm = False
        self.uvm_home = ''
        util.error(f'Attempt to set SlangAst uvm_home={uvm_home} failed, either not',
                   'a directory, or file uvm_pkg.sv not present')


    def run_slang_for_ast_auto(self) -> None:
        '''Wrapper for run_slang_for_ast(...)

        Will set self.handler_obj.has_pre_compile_slang_sh = self.slang_sh_fname,
        if that attr exists in self.handler_obj
        '''
        assert self.handler_obj, 'Need a handler_obj set to run this method'
        self.run_slang_for_ast(
            defines=self.handler_obj.defines,
            incdirs=self.handler_obj.incdirs,
            parameters=self.handler_obj.parameters,
            files=(self.handler_obj.files_sv + self.handler_obj.files_v),
            work_dir=self.handler_obj.args.get('work-dir', '')
        )


    def run_slang_for_ast( # pylint: disable=too-many-branches
            self,
            defines: dict,
            incdirs: list,
            parameters: dict,
            files: list,
            work_dir: str = ''
    ) -> None:
        '''Attemps to run slang to get an AST slang.json

        Will set self.handler_obj.has_pre_compile_slang_sh = self.slang_sh_fname,
        if that attr exists in self.handler_obj
        '''

        slang_available.init()
        if not slang_available.available or not slang_available.exe:
            if t := getattr(self.handler_obj, '_TOOL', ''):
                tool_str = f'by --tool={t}'
            else:
                tool_str = ''
            warn_slang_not_available(f'for detailed coverage {tool_str}')
            return


        command_list = [
            slang_available.exe, # from pymodule global
        ] + self.slang_command_args + [
            '--ast-json', self.slang_json_fname
        ]

        if not any((a.startswith('-j') or a.startswith('--threads')) for a in command_list):
            command_list.extend( ['-j', str(self.threads) ] )

        # incdirs:
        if self.needs_uvm and self.uvm_home:
            command_list.extend(['--include-directory', self.uvm_home])
        for value in incdirs:
            command_list.extend(['--include-directory', value])

        # defines (caller may choose to send no defines, and instead use
        # CommandSim.create_ext_defines_sv(obj)
        for k,v in defines.items():
            command_list.append( '--define-macro' )
            if v is None:
                command_list.append(k)
            else:
                # Generally we should only support int and str python types passed as
                # --define-macro {k}={v}
                command_list.append(f'{k}={sanitize_defines_for_sh(v)}')

        # parameters:
        command_list.extend(
            sim.parameters_dict_get_command_list(params=parameters, arg_prefix='-G')
        )

        # Note - As of slang <= 10.0.30, do not be tempted to add uvm_pkg.sv or uvm.sv,
        # even if you have an know their whereabouts (self.uvm_home and self.needs_uvm).
        # Slang can parse them, but will not build you an AST slang.json, which is what
        # we ultimately need.

        # files:
        command_list.extend(files)

        # run from work-dir
        assert hasattr(self.handler_obj, 'exec'), f'{self.handler_obj=} does not have exec method'

        # This command needs to make it into all.sh in the work-dir, otherwise
        # the only way to re-run is through "eda COMMAND ..."
        if self.slang_sh_fname:
            util.write_shell_command_file(
                dirpath=work_dir, filename=self.slang_sh_fname,
                command_lists=[command_list], line_breaks=False
            )
            if os.path.exists(os.path.join(work_dir, self.slang_sh_fname)) and \
               hasattr(self.handler_obj, 'has_pre_compile_slang_sh'):
                self.handler_obj.has_pre_compile_slang_sh = self.slang_sh_fname

        # Using self.handler_obj.exec timeouts, we may want to use a default value
        # if the handler_obj doesn't have a subprocess-timeout arg set.
        _, stdout, rc = self.handler_obj.exec(
            work_dir=work_dir,
            command_list=command_list,
            background=False,
            stop_on_error=False,
            quiet=False,
            tee_fpath='slang.ast.log',
            shell=False,
            ignore_error=True,
            timeout=self.timeout,
            timeout_str=f'SlangAst internal timeout ({self.timeout} seconds)'
        )
        slang_json_fpath = os.path.join(work_dir, self.slang_json_fname)
        if os.path.isfile(slang_json_fpath):
            util.info(f'slang_helpers.SlangAst.run_slang_for_ast({work_dir=}, ...):',
                      f'wrote slang AST to: {slang_json_fpath}')
            self.init_ast_filepath(filepath=slang_json_fpath)
            return

        util.warning(f'slang_helpers.SlangAst.run_slang_for_ast({work_dir=}, ...)',
                     f'unable to generate {self.slang_json_fname}')

        if rc != 0:
            util.warning(f'slang_helpers.SlangAst.run_slang_for_ast({work_dir=}, ...)',
                         f'returned non-zero return code: {rc}')
        elif 'Build failed:' in stdout:
            util.warning(f'slang_helpers.SlangAst.run_slang_for_ast({work_dir=}, ...)',
                         'build was unsuccessful')


    @staticmethod
    def static_covergroups_extract_details(name: str, member, only_covergroups: list):
        """Specialized helper to dive into CovergroupType / CovergroupBody."""
        if not name:
            cvg_name = member.get("name", "unnamed_cg")
        else:
            cvg_name = name
        points = {}

        collect_enabled = only_covergroups is not None

        if not collect_enabled:
            return {}

        collect_all = collect_enabled and not only_covergroups # empty list


        # Covergroups have a 'CovergroupBody' in their members
        for m in member.get("members", []):
            if m.get("kind") == "CovergroupBody":
                # The Body contains the actual Coverpoints
                for body_m in m.get("members", []):
                    if body_m.get("kind") == "Coverpoint":
                        cp_name = body_m.get("name", "unnamed_cp")

                        collect = bool(
                            collect_all or
                            cvg_name in only_covergroups or
                            cp_name in only_covergroups
                        )

                        # Extract Bins if you need them for explicit coverage hits
                        bins = [b.get("name") for b in body_m.get("members", [])
                                if b.get("kind") == "CoverageBin"]

                        if collect:
                            points[cp_name] = bins

        if points or \
           not only_covergroups or \
           cvg_name in only_covergroups:
            return {"name": cvg_name, "points": points}

        return {}

    def get_items_all(self) -> list:
        '''Gets all items from get_items(..) starting at root'''
        self.all_items = self.get_items(node=self.root_node)
        return self.all_items


    # pylint: disable=unknown-option-value
    # pylint: disable=dangerous-default-value
    # pylint: disable=too-many-locals,too-many-branches,too-many-statements
    # pylint: disable=too-many-arguments
    # pylint: disable=too-many-positional-arguments
    def get_items(
            self,
            node: dict = {},  # current node in dict AST.
            path: str = "",  # current path in the hier, "" is start at root.
            traversed_modules: list = [], # list of all modules up to this point (incl itself)
            only_top_module: str = '',    # if set, value must be in traversed_modules
            only_module_types: list = [], # if empty, use all
            only_files: list = [],        # leaf file names, if empty use all.
            only_paths: list = [],        # paths if current path starts with them
            only_ports: list = [],        # if empty, use all: 'In', 'Out', 'Inout']
            only_signals: list = [],      # if empty, use all
            only_covergroups: list = [],
            max_depth: int = 99
    ) -> list:
        """
        Recursively walks the 'design' hierarchy in a slang JSON file.
        """

        collect_ports_enabled = only_ports is not None
        collect_signals_enabled = only_signals is not None
        collect_covergroups_enabled = only_covergroups is not None


        if not node:
            node = self.root_node
        if not node:
            return []

        # 1. Identify if this node is an Instance or GenerateBlock
        #    but collect some details first, if they exist
        node_kind = node.get("kind", "")
        members = node.get("members", [])
        body = node.get("body", {})
        if body and isinstance(body, dict):
            if members_alt := body.get("members", []):
                members = members_alt
        else:
            # TODO(drew): shortcut here, but if body is a str,
            # it is a space separatated "(addr) (module-type)"
            # you would have to go find the addr previously in the
            # ast... skipping this for now.
            body = {}

        inst_name = node.get("name", "unnamed")
        if path:
            current_path = f"{path}{self.hier_sep}{inst_name}"
        else:
            current_path = inst_name

        # Usually will traverse $root, '', all packages, then TOP/, ..
        util.debug(f'SlangAst.get_items -- start {current_path=} {max_depth=}')

        if node_kind in ("Instance", "GenerateBlock", "Package"):

            # v10.0 uses 'source_file' directly
            source_file = node.get("source_file", "unknown")

            # 2. Extract Ports and find Child Instances from the body
            if body.get('kind', '') in ('InstanceBody', 'PackageBody'):
                module_type = body.get('name', 'unknown')
                source_file = body.get("source_file", "unknown")
            else:
                module_type = ''

            ports = []
            signals = []
            covergroups = []
            child_results = []

            collect = True
            if all((only_module_types, module_type, module_type not in only_module_types)):
                collect = False
                util.debug(f'SlangAst.get_items -- not collecting at {current_path},',
                           f'{module_type=} not in {only_module_types=}')
            elif only_files and os.path.split(source_file)[-1] not in only_files:
                collect = False
                util.debug(f'SlangAst.get_items -- not collecting at {current_path},',
                           f'file {os.path.split(source_file)[-1]} not in {only_files=})')
            elif only_paths and not any(current_path.startswith(p) for p in only_paths):
                # TODO(drew): it might be worth skipping the only_top_module path is that
                # is set, that way we can match on u_dut/ without needing testbench_name/u_dut/
                collect = False
                util.debug(f'SlangAst.get_items -- not collecting at {current_path},',
                           f'not in {only_paths=})')
            elif only_top_module and (only_top_module not in traversed_modules + [module_type]):
                collect = False
                util.debug(f'SlangAst.get_items -- not collecting at {current_path},',
                           f'{only_top_module=} not in traversed modules')

            for m in members:
                kind = m.get("kind", "")
                name = m.get("name", "")


                # Collect ports
                if collect and collect_ports_enabled and kind == "Port":
                    direction = m.get("direction", "")
                    real_name = m.get("internalSymbol", "").split()[-1]
                    if real_name and (not only_ports or direction in only_ports):
                        ports.append({
                            real_name: {
                                "direction": m.get("direction", "")
                            }
                        })

                # Collect signals
                elif collect and collect_signals_enabled and kind in ("Net", "Variable"):
                    if name and (not only_signals or name in only_signals):
                        signals.append(name)

                # Collect Coverage (Covergroups/Points)
                elif collect and collect_covergroups_enabled and kind == "CovergroupType":
                    # Covergroups often have nested Coverpoints in their own body, call helper
                    covergroups.append(SlangAst.static_covergroups_extract_details(
                            name=name, member=m, only_covergroups=only_covergroups
                    ))

                # 3. RECURSE: Find children inside the body
                elif any((kind.startswith(x) for x in ("Instance", "GenerateBlock",
                                                       "Package"))):
                    if module_type and kind == 'Package':
                        new_traversed_modules = [module_type]
                    elif module_type:
                        new_traversed_modules = traversed_modules + [module_type]
                    else:
                        new_traversed_modules = traversed_modules

                    if max_depth > 0:
                        child_results.extend(self.get_items(
                            node=m, path=current_path,
                            traversed_modules=new_traversed_modules,
                            max_depth=(max_depth-1),
                            only_top_module=only_top_module,
                            only_paths=only_paths,
                            only_ports=only_ports, only_module_types=only_module_types,
                            only_covergroups=only_covergroups, only_files=only_files
                        ))


            if collect and node_kind == "Instance":
                if module_type:
                    new_traversed_modules = traversed_modules + [module_type]
                else:
                    new_traversed_modules = traversed_modules
                self_data = [{
                    "path": current_path,
                    "module": module_type,
                    "traversed_modules": new_traversed_modules,
                    "inst_name": inst_name,
                    "file": str(Path(source_file).absolute()),
                    "ports": ports,
                    "signals": signals,
                    "covergroups": covergroups
                }]
            else:
                self_data = []

            return self_data + child_results



        # 4. Handle the top-level design container
        results = []
        # If the current node isn't an instance but has members (like the root), check them
        for m in node.get("members", []):
            results.extend(self.get_items(
                node=m, path=path,
                traversed_modules=traversed_modules,
                max_depth=max_depth,
                only_top_module=only_top_module,
                only_paths=only_paths,
                only_ports=only_ports, only_module_types=only_module_types,
                only_covergroups=only_covergroups, only_files=only_files
            ))

        if results:
            return results

        if m := node.get('design', {}):
            results.extend(self.get_items(
                node=m, path=path,
                traversed_modules=[],
                max_depth=max_depth,
                only_top_module=only_top_module,
                only_paths=only_paths,
                only_ports=only_ports, only_module_types=only_module_types,
                only_covergroups=only_covergroups, only_files=only_files
            ))

        return results
