import os

# ==========================================================================
# 🛠️ Environment Variables & Targets
# ==========================================================================
DEV_MODE = False 
DEV_TARGET_MODEL = ["dev_moss","dev_krea2","pony"]
PROD_TARGET_MODEL = ["moss", "krea2", "anima", "flux1_dev", "pony"]

# 🌟 核心新增：控制是否从独立的 GUI Space 动态同步前端文件
# True = 每次启动时从 GUI Space 拉取最新代码（适合线上开发阶段）
# False = 完全使用当前仓库内的静态文件（适合生产环境锁定或断网兜底）
SYNC_FROM_GUI_SPACE = True

# 🌟 新增：全局 TE 精度开关
# 可选值："w4a8","w6a8" or "int8"
MOSS_TE_PRECISION = "w4a8" 
KREA2_TE_PRECISION = "w4a8" 
FLUX1_DEV_T5_TE_PRECISION = "w4a8" 

# 根据精度开关，动态定义各模型的 TE 路径和文件名
TE_MAPPING = {
    "moss": {
        "w4a8": ("moss/text_encoders/qwen3.5_9b_w4a8_gptq.safetensors", "qwen3.5_9b_w4a8_gptq.safetensors"),
        "w6a8": ("moss/text_encoders/qwen3.5_9b_w6a8_convrot.safetensors", "qwen3.5_9b_w6a8_convrot.safetensors"),
        "int8": ("moss/text_encoders/qwen3.5_9b_int8_convrot.safetensors", "qwen3.5_9b_int8_convrot.safetensors")
    },
    "krea2": {
        "w4a8": ("kr2/text_encoders/qwen3vl_4b_w4a8_convrot.safetensors", "qwen3vl_4b_w4a8_convrot.safetensors"),
        "w6a8": ("kr2/text_encoders/qwen3vl_4b_w6a8_convrot.safetensors", "qwen3vl_4b_w6a8_convrot.safetensors"),
        "int8": ("kr2/text_encoders/qwen3vl_4b_int8_convrot.safetensors", "qwen3vl_4b_int8_convrot.safetensors")
    },
    "flux1_dev_t5": {
        "w4a8": ("f1d/text_encoders/t5xxl_w4a8_convrot.safetensors", "t5xxl_w4a8_convrot.safetensors"),
        "w6a8": ("f1d/text_encoders/t5xxl_w6a8_convrot.safetensors", "t5xxl_w6a8_convrot.safetensors"),
        "int8": ("f1d/text_encoders/t5xxl_int8_convrot.safetensors", "t5xxl_int8_convrot.safetensors")
    }
}

# 获取当前选中的路径
moss_te_path, moss_te_file = TE_MAPPING["moss"][MOSS_TE_PRECISION]
krea2_te_path, krea2_te_file = TE_MAPPING["krea2"][KREA2_TE_PRECISION]
flux_t5_path, flux_t5_file = TE_MAPPING["flux1_dev_t5"][FLUX1_DEV_T5_TE_PRECISION]

# ==========================================================================
# ⚙️ Model Configuration
# ==========================================================================
MODEL_CONFIG = {
    # ==========================================
    # 🌐 Frontend & Sync Config
    # ==========================================
    "ui_space_url": "https://huggingface.co/spaces/waveshaker/universal-lora-studio-gui/resolve/main",
    "frontend_cache_dir": "frontend_cache", 

    # ==========================================
    # ⚙️ Comfy & Model Config
    # ==========================================
    
    "config_url": "https://huggingface.co/waveshaker/comfy_models/resolve/main/studio_config.json",
    "char_lora_repo": "waveshaker/loras",
    
    "name": "Waveshaker-Universal-T2I-Studio-Core", 
    "exclude_folders": ["Heels"],

    "default_prompt": "20 years old woman, 1girl, RAW photo, PORTRAIT, (photorealistic:1.37, realistic), (straight from front:1.2), looking at viewer, ((wearing a simple white t-shirt:1.4)), (casual clothing:1.2), soft lighting, high quality, film grain, Fujifilm XT3, ((dynamic pose:1.2)), ((upper body:1.6)), (professional lighting:1.6), (blurred cafe or minimalist room background:1.2), clean environment, centered composition, dead center, medium shot, ample negative space around.",
    
    "node_types": {
        "unet_loader": "UNETLoader",
        "clip_loader": "CLIPLoader", 
        "vae_loader": "VAELoader",
        "lora_loader": "LoraLoaderModelOnly"
    },
    
    # === Base Model Config === 
    "models": {
        "moss": {
            "name": "Moss-Image-SFT",
            "companion_downloads": [
                ("waveshaker/comfy_models", moss_te_path),
                ("waveshaker/comfy_models", "moss/vae/flux2-vae_small.safetensors"),
            ],
            "local_te_filename": [moss_te_file], 
            "local_vae_filename": "flux2-vae_small.safetensors",
            "clip_type": "moss_image",
            "char_lora_prefix": "moss/",
            "defaults": { "steps": 8, "cfg": 1.0, "sampler": "euler", "scheduler": "simple", "char_lora_strength": 1.0 },
            
            "fallback_unets": {
                "Moss Image SFT Int8 Convrot (8Steps DMD2 Merged)": {
                    "repo": "waveshaker/comfy_models", 
                    "subfolder": "moss/diffusion_models", 
                    "filename": "moss_image_sft_8steps_dmd2_int8_convrot.safetensors"
                }
            },
            "fallback_base_loras": {}
        },
        
        "dev_moss": {
            "name": "Moss-Dev-Test",
            "companion_downloads": [
                ("waveshaker/comfy_models", moss_te_path),
                ("waveshaker/comfy_models", "moss/vae/flux2-vae_small.safetensors"),
            ],
            "local_te_filename": [moss_te_file], 
            "local_vae_filename": "flux2-vae_small.safetensors",
            "clip_type": "moss_image",
            "char_lora_prefix": "moss/",
            "defaults": { "steps": 8, "cfg": 1.0, "sampler": "euler", "scheduler": "simple", "char_lora_strength": 1.0 },
            
            "fallback_unets": {
                "Moss Image SFT Int8 Convrot (8Steps DMD2 Merged)": {
                    "repo": "waveshaker/comfy_models", 
                    "subfolder": "moss/diffusion_models", 
                    "filename": "moss_image_sft_8steps_dmd2_int8_convrot.safetensors"
                }
            },
            "fallback_base_loras": {}
        },

        "krea2": {
            "name": "Krea-2-Turbo",
            "companion_downloads": [
                ("waveshaker/comfy_models", krea2_te_path),
                ("waveshaker/comfy_models", "kr2/vae/qwen_image_vae.safetensors"),
            ],
            "local_te_filename": [krea2_te_file], 
            "local_vae_filename": "qwen_image_vae.safetensors",
            "clip_type": "krea2",
            "char_lora_prefix": "krea2/",
            "defaults": { "steps": 8, "cfg": 1.0, "sampler": "euler_ancestral", "scheduler": "beta", "char_lora_strength": 1.0 },
            
            "fallback_unets": {
                "Krea 2 Turbo Int8 Convrot": {
                    "repo": "waveshaker/comfy_models", 
                    "subfolder": "kr2/diffusion_models", 
                    "filename": "krea2_turbo_int8_convrot.safetensors"
                }
            },
            "fallback_base_loras": {},
        },

        "dev_krea2": {
            "name": "Krea2-Dev-Test",
            "companion_downloads": [
                ("waveshaker/comfy_models", krea2_te_path),
                ("waveshaker/comfy_models", "kr2/vae/qwen_image_vae.safetensors"),
            ],
            "local_te_filename": [krea2_te_file], 
            "local_vae_filename": "qwen_image_vae.safetensors",
            "clip_type": "krea2",
            "char_lora_prefix": "krea2/",
            "defaults": { "steps": 8, "cfg": 1.0, "sampler": "euler_ancestral", "scheduler": "beta", "char_lora_strength": 1.0 },
            
            "fallback_unets": {
                "Krea 2 Turbo Int4 Convrot": {
                    "repo": "waveshaker/comfy_models", 
                    "subfolder": "kr2/diffusion_models", 
                    "filename": "Krea2-Turbo-int4-ConvRot-ComfyUI.safetensors"
                }
            },
            "fallback_base_loras": {},
        },
        "anima": {
            "name": "Anima",
            "companion_downloads": [
                ("waveshaker/comfy_models", "anima/text_encoders/Qwen3-0.6B_base-int8_convrot.safetensors"),
                ("waveshaker/comfy_models", "anima/vae/qwen_image_vae.safetensors"),
            ],
            "local_te_filename": ["Qwen3-0.6B_base-int8_convrot.safetensors"], 
            "local_vae_filename": "qwen_image_vae.safetensors",
            "clip_type": "stable_diffusion",
            "char_lora_prefix": "anima/",
            "defaults": { "steps": 8, "cfg": 1.0, "sampler": "euler", "scheduler": "simple", "char_lora_strength": 1.0 },
            
            "fallback_unets": {
                "Diving Anima V4 Int8 Convrot (Anima Turbo LoRA v0.2 merged)": {
                    "repo": "waveshaker/comfy_models",
                    "subfolder": "anima/diffusion_models",
                    "filename": "divingAnima_v40_turbo_int8_convrot.safetensors"
                }
            },
            "fallback_base_loras": {
                "Anima Highres/Aesthetic Boost": {
                    "repo": "waveshaker/comfy_models",
                    "subfolder": "anima/loras",
                    "filename": "anima-highres-aesthetic-boost.safetensors",
                }
            },
        },
        
        "dev_anima": {
            "name": "Anima-Dev-Test",
            "companion_downloads": [
                ("waveshaker/comfy_models", "anima/text_encoders/Qwen3-0.6B_base-int8_convrot.safetensors"),
                ("waveshaker/comfy_models", "anima/vae/qwen_image_vae.safetensors"),
            ],
            "local_te_filename": ["Qwen3-0.6B_base-int8_convrot.safetensors"], 
            "local_vae_filename": "qwen_image_vae.safetensors",
            "clip_type": "stable_diffusion",
            "char_lora_prefix": "anima/",
            "defaults": { "steps": 4, "cfg": 1.0, "sampler": "euler", "scheduler": "simple" },
            
            "fallback_unets": {
                "Anima Int4 Convrot MSE": {
                    "repo": "waveshaker/comfy_models",
                    "subfolder": "anima/diffusion_models",
                    "filename": "animaINT4Quants_int4ConvrotMSE.safetensors"
                }
            },
            "fallback_base_loras": {}
        },
        "flux1_dev": {
            "name": "Flux.1-dev",
            "clip_loader_override": "DualCLIPLoader",
            "companion_downloads": [
                ("waveshaker/comfy_models", "f1d/text_encoders/clip_l.safetensors"),
                ("waveshaker/comfy_models", flux_t5_path),
                ("waveshaker/comfy_models", "f1d/vae/ae.safetensors"),
            ],
            "local_te_filename": ["clip_l.safetensors",flux_t5_file], 
            "local_vae_filename": "ae.safetensors",
            "clip_type": "flux",
            "char_lora_prefix": "flux1_dev/",
            "defaults": { "steps": 8, "cfg": 1.0, "sampler": "euler_ancestral", "scheduler": "beta", "char_lora_strength": 1.0 },
            
            "fallback_unets": {
                "Flux.1-dev Int8 Convrot (Flux.1 Turbo Alpha merged)": {
                    "repo": "waveshaker/comfy_models",
                    "subfolder": "f1d/diffusion_models",
                    "filename": "flux1-dev_turbo_ctq_int8_convrot.safetensors"
                }
            },
            "fallback_base_loras": {}
        },

        "pony": {
            "name": "Pony-Diffusion-V6",
            "clip_loader_override": "DualCLIPLoader",
            "default_prompt_override": "score_9,score_8_up,score_7_up,score_6_up,20 years old asian woman, RAW photo, PORTRAIT,PORTRAIT ONLY, (photorealistic:1.37, realistic), 1girl, black hair, looking at viewer, (((straight from front))), soft lighting, high quality, film grain, Fujifilm XT3, ((dynamic pose:1.2)) ((upper body:1.6)) (professional lighting:1.6) , simple and clean environment Background, varying from a blurred cafe, a minimalist room, to a soft outdoor natural setting.Centered composition, positioned exactly in the dead center. Medium shot with ample negative space all around.",
            "companion_downloads": [
                ("waveshaker/comfy_models", "pony/text_encoders/waiANIHENTAIPONYXL_v60_clipL.safetensors"),
                ("waveshaker/comfy_models", "pony/text_encoders/waiANIHENTAIPONYXL_v60_clipG.safetensors"),
                ("waveshaker/comfy_models", "pony/vae/sdxl_vae.safetensors"),
            ],
            "local_te_filename": ["waiANIHENTAIPONYXL_v60_clipL.safetensors","waiANIHENTAIPONYXL_v60_clipG.safetensors"], 
            "local_vae_filename": "sdxl_vae.safetensors",
            "clip_type": "sdxl",
            "char_lora_prefix": "pony/",
            "defaults": { "steps": 24, "cfg": 4.0, "sampler": "euler_ancestral", "scheduler": "normal", "char_lora_strength": 1.0 },
            
            "fallback_unets": {
                "WAI-ANI-HENTAI-PONYXL V6": {
                    "repo": "waveshaker/comfy_models",
                    "subfolder": "pony/diffusion_models",
                    "filename": "waiANIHENTAIPONYXL_v60_Diffusion.safetensors"
                }
            },
            "fallback_base_loras": {
                "PAseer-SDXL/PONY-LCM and Turbo Accelerator": {
                        "repo": "waveshaker/comfy_models",
                        "subfolder": "pony/loras",
                        "filename": "LCMV2-PONYplus-PAseer.safetensors"
                }
            }
        }
    },
    
    "json_keys": {
        "unets": "unets",
        "base_loras": "base_loras",
        "root_mapping": {
            "moss": "moss",
            "krea2": "krea2",
            "anima": "anima",
            "flux1_dev": "flux1_dev",
            "pony": "pony",
            "dev_moss": "dev_moss",
            "dev_krea2": "dev_krea2",
            "dev_anima": "dev_anima"
        }
    }
}

# ==========================================================================
# 🌟 Environment Filtering
# ==========================================================================
if DEV_MODE:
    print(f"🛠️ [DEV_MODE = True] 开发环境：仅加载 {DEV_TARGET_MODEL}。")
    TARGET_MODELS = DEV_TARGET_MODEL
else:
    print(f"🚀 [DEV_MODE = False] 生产环境：加载 {PROD_TARGET_MODEL}。")
    TARGET_MODELS = PROD_TARGET_MODEL

MODEL_CONFIG["models"] = {
    k: v for k, v in MODEL_CONFIG["models"].items() 
    if k in TARGET_MODELS
}

MODEL_CONFIG["json_keys"]["root_mapping"] = {
    k: v for k, v in MODEL_CONFIG["json_keys"]["root_mapping"].items()
    if k in TARGET_MODELS
}