"""This modules uses the JSON files created by the 
https://github.com/butaixianran/Stable-Diffusion-Webui-Civitai-Helper extension"""

import sys
import os
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "scripts"))
from scripts.Logger import Logger
from scripts.Utils import Utils

import modules.shared as shared
import modules.scripts as scripts
import json


class CivitaihelperPrompts():
    """
    This class is used to get the prompts from the JSON files created by the
    Stable-Diffusion-Webui-Civitai-Helper extension
    """

    def __init__(self) -> None:
        self.model_path = self.get_custom_model_folder()
        self.utils = Utils()
        self.logger = Logger()
        self.logger.debug = False

    def get_custom_model_folder(self) -> str:
        """Get the custom model folder from the command line options or the default folder

        Returns:
            str: the model folder path
        """
        if shared.cmd_opts.ckpt_dir and os.path.isdir(shared.cmd_opts.ckpt_dir):
            return shared.cmd_opts.ckpt_dir
        else:
            return os.path.join(scripts.basedir(), "models", "Stable-diffusion")


    def get_civitAi_prompt_from_model(self, path: str) -> str:
        """Get the prompt from the civitai.info file

        Args:
            path (str): the path to the model

        Returns:
            str: the prompt
        """
        path = path.replace(".ckpt", ".civitai.info").replace(
            ".safetensors", ".civitai.info")
        path = self.utils.get_clean_checkpoint_path(path)
        path = os.path.join(self.model_path, path)
        self.logger.debug_log(f"{path} -> is file {os.path.isfile(path)}")
        if not os.path.exists(os.path.realpath(path)):
            return "{prompt};"
        model_info = None
        with open(os.path.realpath(path), 'r') as f:
            try:
                model_info = json.load(f)
            except Exception as e:
                self.logger.debug_log(f"Error loading civitai info: {e}", False)
                return "{prompt};"
        try:
            self.logger.debug_log(f"len: {len(model_info['images'])}")
            for i in range(0, len(model_info['images'])):
                try:
                    info = model_info['images'][i]['meta']['prompt']
                    self.logger.debug_log(f"Prompt: {info}")
                    if info:
                        return f"{info};"
                except:
                    pass
            return "{prompt};"
        except:
            return "{prompt};"

    def createCivitaiPromptString(self, checkpoints: str) -> str:
        """Create a string with all the prompts from the checkpoints
        Indices are added to the prompts

        Args:
            checkpoints (str): the checkpoints

        Returns:
            str: the prompt string with indices
        """
        checkpoints_list = self.utils.getCheckpointListFromInput(checkpoints)
        prompts = ""
        prompts_with_info = ""
        for i, checkpoint in enumerate(checkpoints_list):
            prompts += self.get_civitAi_prompt_from_model(checkpoint)

        prompts_with_info += self.utils.add_index_to_string(
            prompts, is_checkpoint=False)

        self.logger.log_info("loaded all prompts")
        return prompts_with_info