import os
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# ComfyUI Settings
COMFY_HOST = os.getenv("COMFY_HOST", "127.0.0.1")
COMFY_PORT = os.getenv("COMFY_PORT", "8188")
COMFY_URL = f"http://{COMFY_HOST}:{COMFY_PORT}"
COMFY_WORKFLOW_FILE = "comfy/comfy_rpg_char_gen.json"

# Ollama Settings
OLLAMA_HOST = os.getenv("OLLAMA_HOST", "127.0.0.1")
OLLAMA_PORT = os.getenv("OLLAMA_PORT", "11434")
OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "llama3")
HF_TOKEN = os.getenv("HF_TOKEN")
HF_BASE_URL = "https://router.huggingface.co/v1"
HF_TEXT_MODEL = "Qwen/Qwen2.5-72B-Instruct"
HF_IMAGE_MODEL = "black-forest-labs/FLUX.1-dev"

HF_TEXT_MODELS = [
    "Qwen/Qwen2.5-72B-Instruct",
    "meta-llama/Llama-3.1-70B-Instruct",
    "mistralai/Mistral-7B-Instruct-v0.3",
    "microsoft/Phi-3-mini-4k-instruct"
]

HF_IMAGE_MODELS = [
    "black-forest-labs/FLUX.1-dev",
    "Tongyi-MAI/Z-Image-Turbo",
    "Qwen/Qwen-Image-2512"
]

# Gemini Settings
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
GEMINI_TEXT_MODEL = "gemini-3-pro-preview"
GEMINI_IMAGE_MODEL = "imagen-4.0-generate-001"

# Application Constants
EXAMPLES_DIR = "examples"
PROMPTS_FILE = "prompts.yaml"
FEATURES_FILE = "features.yaml"

# Define segments for prompt building
# (Category, Subcategory, Template Key)
FEATURE_SEQUENCE = [
    ('identity', 'race', 'race'), 
    ('identity', 'class', 'class'), 
    ('identity', 'gender', 'gender'), 
    ('identity', 'age', 'age'),
    ('expression_pose', 'expression', 'expression'), 
    ('expression_pose', 'pose', 'pose'),
    ('appearance', 'hair_color', 'hair_color'), 
    ('appearance', 'hair_style', 'hair_style'), 
    ('appearance', 'eye_color', 'eye_color'), 
    ('appearance', 'build', 'build'), 
    ('appearance', 'skin_tone', 'skin_tone'), 
    ('appearance', 'distinguishing_feature', 'distinguishing_feature'),
    ('equipment', 'armor', 'armor'), 
    ('equipment', 'weapon', 'weapon'), 
    ('equipment', 'accessory', 'accessory'), 
    ('equipment', 'accessory', 'accessory2'), 
    ('equipment', 'material', 'material'),
    ('environment', 'background', 'background'), 
    ('environment', 'lighting', 'lighting'), 
    ('environment', 'atmosphere', 'atmosphere'),
    ('vfx_style', 'vfx', 'vfx'), 
    ('vfx_style', 'style', 'style'), 
    ('vfx_style', 'mood', 'mood'), 
    ('vfx_style', 'camera', 'camera'),
    ('technical', 'aspect_ratio', 'aspect_ratio')
]

# Section names for extra info
SECTIONS = ['Identity', 'Appearance', 'Equipment', 'Environment', 'Style']