# /ComfyUI-Prompt-Formatter/categorized_random_prompt_formatter.py

import yaml
import re
import random

# Local utility imports
from .prompt_formatter_utils import (
    find_yaml_file,
    clean_output_string,
    resolve_category_tags,
    resolve_inline_choices,
    INCLUDE_DIRECTIVE,
    TAGS_KEY
)

class CategorizedRandomPromptFormatter:
    """
    A ComfyUI node to generate random prompts from categories defined in a YAML file.
    """
    NODE_NAME = "Categorized Random Prompt Formatter"

          
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "category_definition_file": ("STRING", {"default": "prompt_categories.yaml"}),
                "output_template": ("STRING", {
                    "multiline": True,
                    "default": "<|quality:1|>, <|character:1|>, <|setting:1|>"
                }),
                "seed": ("INT", {"default": -1, "min": -1, "max": 0xffffffffffffffff}),
            },
            "optional": { "output_delimiter": ("STRING", {"default": ", "}) }
        }

    RETURN_TYPES = ("STRING", "INT")
    RETURN_NAMES = ("random_prompt", "used_seed")
    FUNCTION = "generate_prompt"
    CATEGORY = "text/generation" 

    def generate_prompt(self, category_definition_file, output_template, seed, output_delimiter=", "):
        # --- 1. Handle Seed ---
        used_seed = random.randint(0, 0xffffffffffffffff) if seed == -1 else seed
        rng = random.Random(used_seed)

        # --- 2. Load & Resolve Categories ---
        resolved_categories = {}
        yaml_path = find_yaml_file(category_definition_file, self.NODE_NAME)
        raw_yaml_data = None

        if yaml_path:
            try:
                with open(yaml_path, 'r', encoding='utf-8') as f:
                    raw_yaml_data = yaml.safe_load(f)
                if not isinstance(raw_yaml_data, dict):
                     print(f"Warning [{self.NODE_NAME}]: YAML file '{yaml_path}' is not a dictionary.")
                     raw_yaml_data = None
            except Exception as e:
                 print(f"Error [{self.NODE_NAME}]: Loading YAML file {yaml_path}: {e}")

        if raw_yaml_data:
            cache = {}
            for cat_name in list(raw_yaml_data.keys()):
                if str(cat_name).strip() not in[INCLUDE_DIRECTIVE, TAGS_KEY]:
                    resolved_categories[str(cat_name).strip()] = resolve_category_tags(cat_name, raw_yaml_data, cache, self.NODE_NAME)
        else:
            return ("", used_seed)

        # --- 3. Process Template & Generate Prompt ---
        # First resolve inline choices in template to avoid breaking placeholder syntax if any
        template_resolved = resolve_inline_choices(output_template, rng)
        
        result_parts =[]
        last_end = 0
        
        for match in re.finditer(r"<\|([^:]+?)(?::(-?\d+))?(?::([0-9.]+))?\|>", template_resolved):
            cat_name, count_str, prob_str, start, end = match.group(1).strip(), match.group(2), match.group(3), *match.span()
            result_parts.append(template_resolved[last_end:start])
            
            probability = float(prob_str) if prob_str else 1.0
            num_to_pick = int(count_str) if count_str and int(count_str) >= 0 else 1
            
            # Probability test block
            if num_to_pick > 0 and rng.random() <= probability:
                available_tags = list(resolved_categories.get(cat_name,[]))
                if available_tags:
                    sample_count = min(num_to_pick, len(available_tags))
                    tags_to_join = rng.sample(sorted(available_tags), sample_count)
                    result_parts.append(output_delimiter.join(tags_to_join))
                else:
                    print(f"Warning [{self.NODE_NAME}]: Category '{cat_name}' not found or empty.")
            
            last_end = end
        
        result_parts.append(template_resolved[last_end:])
        
        # Second resolution step handles inline choices nested from inside categorized tags 
        joined_prompt = "".join(result_parts)
        final_prompt = resolve_inline_choices(joined_prompt, rng)
        final_prompt = clean_output_string(final_prompt, output_delimiter)
        
        return (final_prompt, used_seed)