import yaml
import os
from typing import Dict, Any, List
from src.ontology.models import PromptIR

class PromptBuilder:
    def __init__(self, profile_name: str = "generic", profiles_dir: str = "profiles"):
        self.profile_name = profile_name.lower()
        self.profiles_dir = profiles_dir
        self.profile_data = self._load_profile(self.profile_name)

    def _load_profile(self, name: str) -> Dict[str, Any]:
        path = os.path.join(self.profiles_dir, f"{name}.yaml")
        if not os.path.exists(path):
            # Fallback to generic if not found
            path = os.path.join(self.profiles_dir, "generic.yaml")
        
        with open(path, "r", encoding="utf-8") as f:
            return yaml.safe_load(f)

    def build(self, ir: PromptIR) -> Dict[str, str]:
        """
        Builds positive and negative prompts based on the IR and selected profile.
        """
        template = self.profile_data.get("positive_template", {})
        ordering = template.get("ordering", [])
        quality_tags = template.get("quality_tags", [])

        positive_tags = []
        
        # Mapping IR sections to ordering keys
        # We need a way to collect tags by category from the structured IR
        section_map = {
            "quality": quality_tags,
            "style": ir.style,
            "characters": [c.name for c in ir.characters if c.name != "Subject"],
            "appearance": [attr for c in ir.characters for attr in c.appearance],
            "clothing": [clo for c in ir.characters for clo in c.clothing],
            "accessories": [acc for c in ir.characters for acc in c.accessories],
            "pose": [p for c in ir.characters for p in c.pose],
            "expression": [e for c in ir.characters for e in c.expression],
            "scene": ir.scene.locations,
            "lighting": ir.scene.lighting,
            "atmosphere": ir.scene.atmosphere,
            "effects": ir.effects,
            "technical_details": ir.technical_details
        }

        for section in ordering:
            tags = section_map.get(section, [])
            for tag in tags:
                if tag and tag not in positive_tags:
                    positive_tags.append(tag)

        return {
            "positive": ", ".join(positive_tags),
            "negative": self.profile_data.get("negative_prompt", "")
        }