from modules import script_callbacks

from lib.prompt_cleaning_hook.options import onUiSettings

from modules import scripts
from modules.processing import StableDiffusionProcessing

from src.prompt_cleaning_hook.cleanerModules import cleanerModules
from lib.prompt_cleaning_hook.options import getOption
from src.prompt_cleaning_hook.extension import extensionId, extensionTitle
from lib.prompt_cleaning_hook.logger import logger

def process(text: str, commaOnLinebreak: bool) -> dict[str, str | list[str]]:
  changingModules: list[str] = []
  for module in cleanerModules:
    newText = module['process'](text)
    if newText != text:
      changingModules.append(module['id'])
      text = newText
  return {'text': text, 'changingModules': changingModules}

def logChange(result, original: str, isPositive: bool = True):
  cleaned: str = result['text']
  originalPromptSingleLine = original.replace('\n', '¶')
  cleanedSingleLine = cleaned.replace('\n', '¶')
  promptName = 'prompt' if isPositive else 'negative prompt'
  modulesString = ' → '.join(result['changingModules'])
  logMessage = f'Cleaned {promptName} from “{originalPromptSingleLine}” to “{cleanedSingleLine}” using modules {modulesString}'
  logChanges = getOption('log_changes')
  if logChanges:
    logger.info(logMessage)
  else:
    logger.debug(logMessage)

class PromptCleaningHook(scripts.Script):
  def __init__(self):
    logger.debug('PromptCleaningHook initialized')
    super().__init__()

  def title(self):
    return extensionTitle

  def show(self, isImg2img):
    return scripts.AlwaysVisible

  def process(self, processing: StableDiffusionProcessing):
    if not getOption('enabled'):
      return

    commaOnLinebreak = getOption('comma_on_linebreak')
    saveMetadata = getOption('save_metadata')

    changedCount = 0
    originalPrompt = processing.prompt
    originalNegativePrompt = processing.negative_prompt

    for i in range(len(processing.all_prompts)):
      cleanResult = process(processing.all_prompts[i], commaOnLinebreak)
      if cleanResult['changingModules'].count == 0:
        continue
      logChange(cleanResult, processing.all_prompts[i])
      changedCount += 1
      processing.all_prompts[i] = cleanResult['text']

    for i in range(len(processing.all_negative_prompts)):
      cleanResult = process(processing.all_negative_prompts[i], commaOnLinebreak)
      if cleanResult['changingModules'].count == 0:
        continue
      logChange(cleanResult, processing.all_negative_prompts[i], False)
      changedCount += 1
      processing.all_negative_prompts[i] = cleanResult['text']

    if changedCount == 0:
      return

    logger.debug(f'Cleaned {changedCount}/{len(processing.all_prompts) + len(processing.all_negative_prompts)} prompts')
    if saveMetadata:
      processing.extra_generation_params.setdefault(f'{extensionId}_original_prompt', originalPrompt)
      processing.extra_generation_params.setdefault(f'{extensionId}_original_negative_prompt', originalNegativePrompt)

script_callbacks.on_ui_settings(onUiSettings)