import os
from typing import Dict, Any, List

MASTER_SECRET = os.getenv("MASTER_SECRET", "nightglow-secret-key-2024")
REPO_ID = "MGFeng/Nightglow-AI"                # 训练数据仓库
BUCKET_ID = "MGFeng/Nightglow-AI-storage"      # 模型存储桶
HEARTBEAT_TIMEOUT = 90

DATA_SOURCES = [
    "nvidia/OpenCodeInstruct",
    "Fsoft-AIC/the-vault-function",
    "MGFeng/Nightglow-AI",
]

MODEL_PRESETS: Dict[str, Dict[str, Any]] = {
    "3B": {
        "hidden_size": 3200,
        "num_layers": 18,
        "num_heads": 20,
        "num_experts": 8,
        "num_kv_heads": 20,
        "intermediate_size": 12800,
        "vocab_size": 49152,
        "max_position_embeddings": 8192,
        "dropout": 0.1,
        "tasks": 1000,
    },
}

TRAINING_CONFIG = {
    "batch_size": 2,
    "max_seq_length": 512,
    "learning_rate": 1e-4,
    "gradient_clip": 1.0,
    "weight_decay": 0.01,
    "warmup_steps": 100,
    "lr_scheduler": "cosine",
    "use_fp16": True,
}

SORTED_SIZES = list(MODEL_PRESETS.keys())

class ConfigManager:
    _instance = None
    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._instance._initialized = False
        return cls._instance

    def __init__(self):
        if self._initialized:
            return
        self._config = {
            "target_size": "3B",
            "data_source": DATA_SOURCES[0],
            "language": "python",
            "total_steps": 1000,
            "steps_per_task": 100,
            **TRAINING_CONFIG,
        }
        self._initialized = True
        print("✅ ConfigManager initialized")

    def get(self, key: str, default=None):
        return self._config.get(key, default)

    def set(self, key: str, value: Any) -> bool:
        if key not in self._config:
            return False
        self._config[key] = value
        return True

    def get_all(self) -> Dict[str, Any]:
        return self._config.copy()

config = ConfigManager()