ai-images/Source/AIImages/Services/AdvancedPromptGenerator.cs

using System.Collections.Generic;
using System.Linq;
using System.Text;
using AIImages.Models;
using RimWorld;
using Verse;

namespace AIImages.Services
{
    /// <summary>
    /// Продвинутый генератор промптов для Stable Diffusion
    /// </summary>
    public class AdvancedPromptGenerator : IPromptGeneratorService
    {
        private static readonly Dictionary<string, string> TraitToMood = new Dictionary<
            string,
            string
        >
        {
            { "Kind", "warm smile, gentle expression, friendly eyes" },
            { "Bloodlust", "intense gaze, fierce expression, aggressive posture" },
            { "Psychopath", "cold eyes, emotionless face, calculating expression" },
            { "Pessimist", "sad eyes, worried expression, downcast gaze" },
            { "Optimist", "bright smile, hopeful expression, cheerful demeanor" },
            { "Nervous", "anxious expression, tense posture, worried eyes" },
            { "Careful", "focused expression, attentive gaze, composed posture" },
            { "Brave", "confident expression, determined gaze, strong posture" },
            { "Wimp", "fearful eyes, nervous expression, timid posture" },
            { "Greedy", "cunning expression, calculating eyes, sly smile" },
            { "Jealous", "envious gaze, bitter expression, tense face" },
            { "Ascetic", "serene expression, calm demeanor, peaceful eyes" },
            { "Beautiful", "stunning features, attractive appearance, graceful" },
            { "Ugly", "rough features, weathered appearance" },
            { "Pretty", "attractive features, pleasant appearance, charming" },
        };

        public string GeneratePositivePrompt(
            PawnAppearanceData appearanceData,
            StableDiffusionSettings settings
        )
        {
            if (appearanceData == null)
                return "portrait of a person";

            StringBuilder prompt = new StringBuilder();

            // 1. Базовый пользовательский промпт (если указан) - идет первым
            if (!string.IsNullOrEmpty(settings.PositivePrompt))
            {
                prompt.Append(settings.PositivePrompt.TrimEnd(',', ' '));
                prompt.Append(", ");
            }

            // 2. Художественный стиль
            var styleDef = DefDatabase<ArtStyleDef>.GetNamedSilentFail(settings.ArtStyleDefName);
            if (styleDef != null && !string.IsNullOrEmpty(styleDef.positivePrompt))
            {
                prompt.Append(styleDef.positivePrompt);
                prompt.Append(", ");
            }

            // 3. Тип кадра
            prompt.Append("portrait, ");

            // 4. Пол персонажа (в формате anime/SD теги)
            string genderTag = appearanceData.Gender == Gender.Female ? "1girl" : "1boy";
            prompt.Append(genderTag);
            prompt.Append(", ");

            // 5. Точный возраст
            prompt.Append($"{appearanceData.Age} y.o.");
            prompt.Append(", ");

            // 6. Тип тела
            string bodyType = GetBodyTypeDescription(appearanceData.BodyType);
            if (!string.IsNullOrEmpty(bodyType))
            {
                prompt.Append(bodyType);
                prompt.Append(", ");
            }

            // 7. Цвет кожи (сначала проверяем гены, затем используем цвет как fallback)
            string skinTone = GetSkinToneDescription(appearanceData);
            if (!string.IsNullOrEmpty(skinTone))
            {
                prompt.Append(skinTone);
                prompt.Append(", ");
            }

            // 8. Волосы
            string hairDescription = GetHairDescription(appearanceData);
            if (!string.IsNullOrEmpty(hairDescription))
            {
                prompt.Append(hairDescription);
                prompt.Append(", ");
            }

            // 9. Настроение и выражение на основе черт характера
            string moodDescription = GetMoodFromTraits(appearanceData.Traits);
            if (!string.IsNullOrEmpty(moodDescription))
            {
                prompt.Append(moodDescription);
                prompt.Append(", ");
            }

            // 10. Одежда
            string apparelDescription = GetApparelDescription(appearanceData.Apparel);
            if (!string.IsNullOrEmpty(apparelDescription))
            {
                prompt.Append(apparelDescription);
                prompt.Append(", ");
            }

            // 11. Качественные теги
            prompt.Append(GetQualityTags(settings.ArtStyleDefName));

            return prompt.ToString().Trim().TrimEnd(',');
        }

        public string GenerateNegativePrompt(StableDiffusionSettings settings)
        {
            StringBuilder negativePrompt = new StringBuilder();

            // 1. Пользовательский негативный промпт (если указан) - идет первым
            if (!string.IsNullOrEmpty(settings.NegativePrompt))
            {
                negativePrompt.Append(settings.NegativePrompt.TrimEnd(',', ' '));
            }

            // Получаем стиль из Def
            var styleDef = DefDatabase<ArtStyleDef>.GetNamedSilentFail(settings.ArtStyleDefName);
            if (styleDef == null || !styleDef.addBaseNegativePrompts)
            {
                // Для стилей без базовых негативов - используем только пользовательский промпт
                return negativePrompt.ToString().Trim();
            }

            // 2. Базовые негативные промпты
            if (negativePrompt.Length > 0)
            {
                negativePrompt.Append(", ");
            }

            negativePrompt.Append(
                "ugly, deformed, low quality, blurry, bad anatomy, worst quality, "
            );
            negativePrompt.Append(
                "mutated, disfigured, bad proportions, extra limbs, missing limbs"
            );

            // 3. Специфичные для стиля негативы из Def
            if (!string.IsNullOrEmpty(styleDef.negativePrompt))
            {
                negativePrompt.Append(", ");
                negativePrompt.Append(styleDef.negativePrompt);
            }

            return negativePrompt.ToString().Trim();
        }

        public string GetFullPromptDescription(
            PawnAppearanceData appearanceData,
            StableDiffusionSettings settings
        )
        {
            StringBuilder description = new StringBuilder();

            description.AppendLine("=== Positive Prompt ===");
            description.AppendLine(GeneratePositivePrompt(appearanceData, settings));
            description.AppendLine();

            description.AppendLine("=== Negative Prompt ===");
            description.AppendLine(GenerateNegativePrompt(settings));
            description.AppendLine();

            description.AppendLine("=== Technical Parameters ===");
            description.AppendLine($"Steps: {settings.Steps}");
            description.AppendLine($"CFG Scale: {settings.CfgScale}");
            description.AppendLine($"Size: {settings.Width}x{settings.Height}");
            description.AppendLine($"Sampler: {settings.Sampler}");
            description.AppendLine(
                $"Seed: {(settings.Seed == -1 ? "Random" : settings.Seed.ToString())}"
            );

            return description.ToString();
        }

        private string GetBodyTypeDescription(string bodyType)
        {
            if (string.IsNullOrEmpty(bodyType))
                return "";

            return bodyType.ToLower() switch
            {
                "thin" => "slender build, lean physique",
                "hulk" => "muscular build, strong physique, athletic body",
                "fat" => "heavyset build, stocky physique",
                "female" => "feminine build",
                "male" => "masculine build",
                _ => "average build",
            };
        }

        private string GetSkinToneDescription(PawnAppearanceData data)
        {
            // Этот метод будет использоваться в контексте генерации промптов
            // где у нас есть только данные внешности, но не сам Pawn
            // Поэтому используем fallback к определению по цвету
            return ColorDescriptionService.GetSkinToneDescription(data.SkinColor);
        }

        private string GetHairDescription(PawnAppearanceData data)
        {
            if (string.IsNullOrEmpty(data.HairDefName))
                return "";

            StringBuilder hair = new StringBuilder();

            // Цвет волос
            string hairColor = ColorDescriptionService.GetHairColorDescription(data.HairColor);
            hair.Append(hairColor);
            hair.Append(" hair, ");

            // Стиль прически - используем DefName для английского названия
            string style = CleanDefName(data.HairDefName)
                .Replace("shaved", "very short")
                .Replace("mohawk", "mohawk hairstyle");

            hair.Append(style);
            hair.Append(" hair");

            return hair.ToString();
        }

        private string GetMoodFromTraits(List<Trait> traits)
        {
            if (traits == null || !traits.Any())
                return "neutral expression";

            // Ищем черты, которые влияют на внешний вид
            foreach (var trait in traits)
            {
                string traitDefName = trait.def.defName;
                if (TraitToMood.TryGetValue(traitDefName, out string mood))
                {
                    return mood;
                }
            }

            return "calm expression";
        }

        private string GetApparelDescription(List<ApparelData> apparel)
        {
            if (apparel == null || !apparel.Any())
                return "simple clothes";

            StringBuilder apparelDesc = new StringBuilder("wearing ");

            // Берем топ 5 наиболее заметных предметов одежды
            var visibleApparel = apparel.Take(5).ToList();

            List<string> items = new List<string>();
            foreach (var item in visibleApparel)
            {
                StringBuilder itemDesc = new StringBuilder();

                // Цвет (если не белый)
                if (item.Color != UnityEngine.Color.white)
                {
                    string colorDesc = ColorDescriptionService.GetApparelColorDescription(
                        item.Color
                    );
                    itemDesc.Append(colorDesc);
                    itemDesc.Append(" ");
                }

                // Материал - используем DefName для английского названия
                if (!string.IsNullOrEmpty(item.MaterialDefName))
                {
                    itemDesc.Append(CleanDefName(item.MaterialDefName));
                    itemDesc.Append(" ");
                }

                // Название предмета - используем DefName для английского названия
                itemDesc.Append(CleanDefName(item.DefName));

                items.Add(itemDesc.ToString());
            }

            apparelDesc.Append(string.Join(", ", items));

            return apparelDesc.ToString();
        }

        private string GetQualityTags(string styleDefName)
        {
            var styleDef = DefDatabase<ArtStyleDef>.GetNamedSilentFail(styleDefName);
            if (styleDef == null)
            {
                return "";
            }

            StringBuilder tags = new StringBuilder();

            // Базовые теги качества
            if (styleDef.addBaseQualityTags)
            {
                tags.Append("highly detailed, professional, masterpiece, best quality");
            }

            // Специфичные для стиля теги качества
            if (!string.IsNullOrEmpty(styleDef.qualityTags))
            {
                if (tags.Length > 0)
                {
                    tags.Append(", ");
                }
                tags.Append(styleDef.qualityTags);
            }

            return tags.ToString();
        }

        /// <summary>
        /// Преобразует defName в читаемый английский текст для промпта
        /// Пример: "Apparel_Pants" -> "pants", "Synthread" -> "synthread"
        /// </summary>
        private string CleanDefName(string defName)
        {
            if (string.IsNullOrEmpty(defName))
                return "";

            string cleaned = defName;

            // Убираем распространенные префиксы RimWorld
            cleaned = System.Text.RegularExpressions.Regex.Replace(
                cleaned,
                "^(Apparel_|Armor_|Weapon_|Thing_)",
                "",
                System.Text.RegularExpressions.RegexOptions.IgnoreCase
            );

            // Разделяем CamelCase на слова
            cleaned = System.Text.RegularExpressions.Regex.Replace(
                cleaned,
                "([a-z])([A-Z])",
                "$1 $2"
            );

            // Заменяем подчеркивания на пробелы
            cleaned = cleaned.Replace("_", " ");

            // Убираем множественные пробелы
            cleaned = System.Text.RegularExpressions.Regex.Replace(cleaned, @"\s+", " ");

            // Приводим к нижнему регистру и убираем лишние пробелы
            return cleaned.Trim().ToLower();
        }
    }
}