"""This module contains the CustomInferenceConfig class.
This class defines configuration parameters for a custom Inference object.
"""

import logging
from pathlib import Path
from typing import List, Optional

from pydantic import model_validator

from mac.config.inference.inference_config import InferenceConfig
from mac.types import SupportedFrameworks

logger = logging.getLogger(__name__)


class CustomInferenceConfig(InferenceConfig):
    """This class defines configuration parameters for a custom Inference object."""

    # Python modules to include
    matching_files: Optional[List[Path]] = None
    modules_to_include: List[Path]

    @model_validator(mode="before")
    @classmethod
    def raise_error_if_framework_invalid(cls, data):
        """Checks that the framework is supported."""
        if data["framework"] != SupportedFrameworks.CUSTOM:
            message = (
                "`framework` should be of type `SupportedFrameworks.CUSTOM`."
            )
            logger.error(message)
            raise ValueError(message)
        return data

    @model_validator(mode="before")
    @classmethod
    def raise_error_if_model_path_not_dir(cls, data):
        """Checks that the model_path is a directory.
        model_path should be a folder containing custom Python modules,
        model files and (optionally) pip requirements."""
        if not data["model_path"].is_dir():
            message = "`model_path` should be a directory."
            logger.error(message)
            raise ValueError(message)
        return data

    @model_validator(mode="after")  # type: ignore
    @classmethod
    def raise_error_if_modules_not_py_files(cls, model_instance):
        """Checks that all the module files are .py files."""
        model_instance.matching_files = [
            file
            for path in model_instance.modules_to_include
            for file in list(
                model_instance.model_path.glob(path.as_posix())
            )  # modules are always included in the model_path
        ]

        if not model_instance.matching_files:
            message = "No matching files found inside `model_path`."
            logger.error(message)
            raise FileNotFoundError(message)

        if any(file.suffix != ".py" for file in model_instance.matching_files):
            message = "`modules_to_include` must only match .py files."
            logger.error(message)
            raise ValueError(message)

        return model_instance
