import yaml
import os
from collections import OrderedDict

CHECKPOINT_DIR = "models/checkpoints"
LORA_DIR = "models/loras"
EMBEDDING_DIR = "models/embeddings" 
CONTROLNET_DIR = "models/controlnet"
MODEL_PATCHES_DIR = "models/model_patches"
DIFFUSION_MODELS_DIR = "models/diffusion_models"
VAE_DIR = "models/vae"
TEXT_ENCODERS_DIR = "models/text_encoders"
STYLE_MODELS_DIR = "models/style_models"
CLIP_VISION_DIR = "models/clip_vision"
IPADAPTER_DIR = "models/ipadapter"
IPADAPTER_FLUX_DIR = "models/ipadapter-flux"
INPUT_DIR = "input"
OUTPUT_DIR = "output"

CATEGORY_TO_DIR_MAP = {
    "diffusion_models": DIFFUSION_MODELS_DIR,
    "text_encoders": TEXT_ENCODERS_DIR,
    "vae": VAE_DIR,
    "checkpoints": CHECKPOINT_DIR,
    "loras": LORA_DIR,
    "controlnet": CONTROLNET_DIR,
    "model_patches": MODEL_PATCHES_DIR,
    "embeddings": EMBEDDING_DIR,
    "style_models": STYLE_MODELS_DIR,
    "clip_vision": CLIP_VISION_DIR,
    "ipadapter": IPADAPTER_DIR,
    "ipadapter-flux": IPADAPTER_FLUX_DIR
}

_PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
_FILE_LIST_PATH = os.path.join(_PROJECT_ROOT, 'yaml', 'file_list.yaml')
_CONSTANTS_PATH = os.path.join(_PROJECT_ROOT, 'yaml', 'constants.yaml')

def load_constants_from_yaml(filepath=_CONSTANTS_PATH):
    if not os.path.exists(filepath):
        print(f"Warning: Constants file not found at {filepath}. Using fallback values.")
        return {}
    with open(filepath, 'r', encoding='utf-8') as f:
        return yaml.safe_load(f)

def load_file_download_map(filepath=_FILE_LIST_PATH):
    if not os.path.exists(filepath):
        raise FileNotFoundError(f"The file list (for downloads) was not found at: {filepath}")
    
    with open(filepath, 'r', encoding='utf-8') as f:
        file_list_data = yaml.safe_load(f)
    
    download_info_map = {}
    for category, files in file_list_data.get('file', {}).items():
        if isinstance(files, list):
            for file_info in files:
                if 'filename' in file_info:
                    file_info['category'] = category
                    download_info_map[file_info['filename']] = file_info
    return download_info_map

try:
    ALL_FILE_DOWNLOAD_MAP = load_file_download_map()
    
    category_to_model_type = {
        "diffusion_models": "UNET",
        "text_encoders": "TEXT_ENCODER",
        "vae": "VAE",
        "checkpoints": "SDXL",
        "loras": "LORA",
        "controlnet": "CONTROLNET",
        "model_patches": "MODEL_PATCH",
        "style_models": "STYLE",
        "clip_vision": "CLIP_VISION",
        "ipadapter": "IPADAPTER",
        "ipadapter-flux": "IPADAPTER_FLUX"
    }

    MODEL_MAP_CHECKPOINT = OrderedDict()
    ALL_MODEL_MAP = OrderedDict()

    for filename, file_info in ALL_FILE_DOWNLOAD_MAP.items():
        category = file_info.get('category')
        repo_id = file_info.get('repo_id', '')
        model_type = category_to_model_type.get(category, 'UNKNOWN')
        
        model_tuple = (repo_id, filename, model_type, "latent", category)
        ALL_MODEL_MAP[filename] = model_tuple
        
        if category == 'checkpoints':
            MODEL_MAP_CHECKPOINT[filename] = model_tuple

    MODEL_TYPE_MAP = {k: v[2] for k, v in ALL_MODEL_MAP.items()}

    ARCH_CATEGORIES_MAP = {}
    for display_name, info in MODEL_MAP_CHECKPOINT.items():
        arch = info[2]
        cat = info[4] if len(info) > 4 else None
        if arch not in ARCH_CATEGORIES_MAP:
            ARCH_CATEGORIES_MAP[arch] = []
        if cat and cat not in ARCH_CATEGORIES_MAP[arch]:
            ARCH_CATEGORIES_MAP[arch].append(cat)

except Exception as e:
    print(f"FATAL: Could not load file download map from YAML. Error: {e}")
    ALL_FILE_DOWNLOAD_MAP = {}
    MODEL_MAP_CHECKPOINT, ALL_MODEL_MAP = OrderedDict(), OrderedDict()
    MODEL_TYPE_MAP = {}
    ARCH_CATEGORIES_MAP = {}

try:
    _constants = load_constants_from_yaml()
    MAX_LORAS = _constants.get('MAX_LORAS', 10)
    MAX_EMBEDDINGS = _constants.get('MAX_EMBEDDINGS', 5)
    MAX_CONDITIONINGS = _constants.get('MAX_CONDITIONINGS', 10)
    MAX_CONTROLNETS = _constants.get('MAX_CONTROLNETS', 5)
    MAX_H3_GUIDES = _constants.get('MAX_H3_GUIDES', 5)
    MAX_IPADAPTERS = _constants.get('MAX_IPADAPTERS', 5)
    LORA_SOURCE_CHOICES = _constants.get('LORA_SOURCE_CHOICES', ["Hugging Face", "Civitai", "File"])
    RESOLUTION_MAP = _constants.get('RESOLUTION_MAP', {})
    MULTIPLIERS_MAP = _constants.get('MULTIPLIERS_MAP', {})
    ARCHITECTURES_CONFIG = {"architectures": {}, "architecture_order": []}
    FEATURES_CONFIG = {"default": {"enabled_chains": ["lora"]}}
    MODEL_DEFAULTS_CONFIG = {"Default": {}}
except Exception as e:
    print(f"FATAL: Could not load constants from YAML. Error: {e}")
    MAX_LORAS, MAX_EMBEDDINGS, MAX_CONDITIONINGS, MAX_CONTROLNETS, MAX_IPADAPTERS = 10, 5, 10, 5, 5
    MAX_H3_GUIDES = 5
    LORA_SOURCE_CHOICES = ["Hugging Face", "Civitai", "File"]
    RESOLUTION_MAP = {}
    MULTIPLIERS_MAP = {}
    ARCHITECTURES_CONFIG = {"architectures": {}, "architecture_order": []}
    FEATURES_CONFIG = {"default": {"enabled_chains": ["lora"]}}
    MODEL_DEFAULTS_CONFIG = {"Default": {}}