#
#   Imandra Inc.
#
#   codelogician/modeling/mlist.py
#

from functools import cmp_to_key
from typing import cast

from pydantic import BaseModel

from .model import Model


class ModelList(BaseModel):
    models: list[Model]


def calc_upstream_affected(models: dict[str, Model], rel_path: str) -> int:
    """
    Return the number of models affected by the model with specified 'rel_path'
    """

    if rel_path not in models:
        raise Exception(f'Specified model [{rel_path}] is not present!')

    def helper(current_path: str, visited: list[str]):
        """Our helper function to help us traverse the models"""

        for path in models:
            if path in visited:
                continue

            found = False

            for d in models[path].dependencies:
                if current_path == cast(Model, d).rel_path:
                    found = True
                    break

            if found:
                visited.append(path)
                visited = helper(path, visited)

        return visited

    return len(helper(rel_path, []))


def gen_listing(models: dict[str, Model], category: str = 'frm_status') -> ModelList:
    """
    Return a sorted list of models by a specified category.

    Possible values are:
    - Formalization Status [value='frm_status'] - formalization status
    - Upstream dependencies [value='upstream'] - number of upstream dependencies
    - Number of opaque functions [value='opaques'] - number of opaque functions
    - Number of failed VGs [value='failed_vgs'] - number of failed verification goals
    """

    supported = ['frm_status', 'upstream', 'opaques', 'failed_vgs']

    # Let's make sure the category is supported first
    if category.lower() not in supported:
        raise Exception(
            f'Supplied category is not supported: {category}. Possible values are: {supported}'
        )

    stats = {}

    # This should get us the number of models that are ultimatey affected by this model
    for rel_path in models:
        upstream = calc_upstream_affected(models, rel_path)
        stats[rel_path] = {'path': rel_path, 'upstream': upstream}

    for rel_path in models:
        modelStats = models[rel_path].gen_stats()

        # Let's now append the statistics here
        for key in modelStats:
            stats[rel_path][key] = modelStats[key]

    def compare(s1, s2):
        try:
            if category.lower() == 'frm_status':
                return s1['frm_status'] >= s2['frm_status']
            elif category.lower() == 'upstream':
                return s1['upstream'] >= s2['upstream']
            elif category.lower() == 'opaques':
                return s1['num_opaques'] >= s2['num_opaques']
            else:
                return s1['failed_vgs'] >= s2['failed_vgs']
        except Exception as e:
            raise Exception(f'Failed to compare model statistcs: {str(e)}')

    return ModelList(models=list(sorted(stats.values(), key=cmp_to_key(compare))))
