import os
import sys
from dotenv import load_dotenv

load_dotenv()


def discover_models():
    """从环境变量中发现所有模型"""
    models = {}
    for key, value in os.environ.items():
        if key.endswith('_MODEL_ID'):
            prefix = key[:-9]
            service_name = os.getenv(f'{prefix}_SERVICE_NAME', f'{prefix.lower()}_classifier')
            service_folder = service_name.replace('_classifier', '')

            models[prefix.lower()] = {
                'prefix': prefix,
                'model_id': value,
                'port': os.getenv(f'{prefix}_PORT', '5002'),
                'service_name': service_name,
                'model_dir': os.getenv(f'{prefix}_MODEL_DIR', f'/app/services/{service_folder}/model'),
                'service_folder': service_folder,
            }
    return models


def generate_supervisor_conf(models):
    """生成supervisor配置文件"""
    conf = []

    conf.append("""[supervisord]
nodaemon=true
logfile=/var/log/supervisor/supervisord.log
pidfile=/tmp/supervisord.pid
childlogdir=/var/log/supervisor

""")

    programs = []
    for name, config in models.items():
        program_name = config['service_name']
        service_folder = config['service_folder']
        programs.append(program_name)

        conf.append(f"""[program:{program_name}]
command=gunicorn --bind 0.0.0.0:{config['port']} --workers %(ENV_MAX_WORKERS)s --threads 2 --timeout 120 services.{service_folder}.app:app
directory=/app
autostart=true
autorestart=true
startretries=3
stdout_logfile=/app/logs/{program_name}.log
stdout_logfile_maxbytes=50MB
stderr_logfile=/app/logs/{program_name}_error.log
stderr_logfile_maxbytes=50MB
environment=SERVICE_NAME="{program_name}",SERVICE_PORT="{config['port']}",MODEL_ID="{config['model_id']}",MODEL_DIR="{config['model_dir']}"

""")  # 注意：移除了 user=appuser

    conf.append(f"""[group:bert_services]
programs={','.join(programs)}
""")

    return ''.join(conf)


def main():
    supervisor_conf_path = './supervisor.conf'

    if os.path.exists(supervisor_conf_path):
        print("=" * 50)
        print("supervisor.conf already exists, skipping generation...")
        sys.exit(0)

    print("=" * 50)
    print("Generating supervisor.conf dynamically...")
    print("=" * 50)

    models = discover_models()

    if not models:
        print("⚠️  No models found")
        sys.exit(1)

    print(f"Found {len(models)} model(s):")
    for name, config in models.items():
        print(f"  - {name}: {config['model_id']} (port {config['port']})")

    conf_content = generate_supervisor_conf(models)

    with open(supervisor_conf_path, 'w') as f:
        f.write(conf_content)

    print(f"\n✓ supervisor.conf generated successfully!")
    print("=" * 50)


if __name__ == "__main__":
    main()