import paramiko
import schedule
import time
import os
import sys
from flask import Flask, jsonify, render_template_string
from threading import Thread
import logging

app = Flask(__name__)

vps_status = {}

# 设置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout),
        logging.StreamHandler(sys.stderr)
    ]
)
logger = logging.getLogger()

def get_vps_configs():
    configs = []
    index = 1
    while True:
        hostname = os.environ.get(f'HOSTNAME_{index}')
        if not hostname:
            break
        
        config = {
            'index': index,
            'hostname': hostname,
            'username': os.environ.get(f'USERNAME_{index}'),
            'password': os.environ.get(f'PASSWORD_{index}'),
            'script_path': os.environ.get(f'SCRIPT_PATH_{index}')
        }
        configs.append(config)
        
        # 添加调试输出
        logger.info(f"Config {index}: {config}")
        print(f"Config {index}: {config}")
        sys.stdout.flush()
        
        index += 1
    return configs

def check_and_run_script(config):
    logger.info(f"Checking VPS {config['index']}: {config['hostname']}")
    print(f"Checking VPS {config['index']}: {config['hostname']}")
    sys.stdout.flush()
    client = None
    try:
        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        
        logger.info(f"Connecting to {config['hostname']}")
        client.connect(
            hostname=config['hostname'],
            username=config['username'],
            password=config['password'],
            port=22
        )
        
        script_path = config['script_path']
        script_name = os.path.basename(script_path)
        
        # 检查脚本是否正在运行 (使用 ps 命令，适用于FreeBSD)
        check_command = f"ps aux | grep {script_name} | grep -v grep"
        stdin, stdout, stderr = client.exec_command(check_command)
        if stdout.read():
            status = "Running"
            logger.info(f"Script is running on {config['hostname']}")
        else:
            logger.info(f"Script not running on {config['hostname']}. Executing restart script.")
            restart_command = f"/bin/sh {script_path}"  # 使用完整路径的sh
            stdin, stdout, stderr = client.exec_command(restart_command)
            exit_status = stdout.channel.recv_exit_status()
            
            if exit_status == 0:
                status = "Restarted"
                logger.info(f"Restart script executed successfully on {config['hostname']}")
            else:
                error_output = stderr.read().decode('utf-8')
                status = f"Error: {error_output}"
                logger.error(f"Error executing restart script on {config['hostname']}: {error_output}")
        
        vps_status[config['hostname']] = {
            'index': config['index'],
            'status': status,
            'last_check': time.strftime('%Y-%m-%d %H:%M:%S'),
            'username': config['username']
        }
        
    except Exception as e:
        logger.error(f"Error occurred while checking VPS {config['index']} - {config['hostname']}: {str(e)}")
        vps_status[config['hostname']] = {
            'index': config['index'],
            'status': f"Error: {str(e)}",
            'last_check': time.strftime('%Y-%m-%d %H:%M:%S'),
            'username': config['username']
        }
    finally:
        if client:
            client.close()
            logger.info(f"SSH connection closed for VPS {config['index']}: {config['hostname']}")
        logger.info(f"Finished checking VPS {config['index']}: {config['hostname']}")
        print(f"Finished checking VPS {config['index']}: {config['hostname']}")
        sys.stdout.flush()

def check_all_vps():
    logger.info("Starting VPS check")
    print("Starting VPS check")
    sys.stdout.flush()
    vps_configs = get_vps_configs()
    for config in vps_configs:
        check_and_run_script(config)
    logger.info("Finished VPS check")
    for hostname, status in vps_status.items():
        print(f"VPS {status['index']} - {hostname}: Status: {status['status']}, Username: {status['username']}")
    sys.stdout.flush()

@app.route('/')
def index():
    html = '''
    <h1>VPS Status Overview</h1>
    <table border="1">
        <tr>
            <th>Index</th>
            <th>Hostname</th>
            <th>Status</th>
            <th>Last Check</th>
            <th>Username</th>
        </tr>
        {% for hostname, data in vps_status.items() %}
        <tr>
            <td>{{ data.index }}</td>
            <td><a href="/status/{{ hostname }}">{{ hostname }}</a></td>
            <td>{{ data.status }}</td>
            <td>{{ data.last_check }}</td>
            <td>{{ data.username }}</td>
        </tr>
        {% endfor %}
    </table>
    '''
    return render_template_string(html, vps_status=vps_status)

@app.route('/status/<hostname>')
def vps_status_detail(hostname):
    if hostname in vps_status:
        return jsonify(vps_status[hostname])
    else:
        return jsonify({"error": "VPS not found"}), 404

@app.route('/health')
def health_check():
    return jsonify({"status": "healthy", "uptime": time.time() - start_time}), 200

def run_flask():
    app.run(host='0.0.0.0', port=8080)

def main():
    global start_time
    start_time = time.time()
    
    print("===== VPS monitoring script is starting =====")
    sys.stdout.flush()
    logger.info("===== VPS monitoring script started =====")
    
    flask_thread = Thread(target=run_flask)
    flask_thread.start()
    logger.info("Flask server started in background")
    print("Flask server started in background")
    sys.stdout.flush()

    vps_configs = get_vps_configs()
    logger.info(f"Found {len(vps_configs)} VPS configurations")
    print(f"Found {len(vps_configs)} VPS configurations")
    sys.stdout.flush()
    for config in vps_configs:
        logger.info(f"VPS configured: {config['hostname']}")
        print(f"VPS configured: {config['hostname']}")
        sys.stdout.flush()

    logger.info("Running initial VPS check")
    print("Running initial VPS check")
    sys.stdout.flush()
    check_all_vps()

    schedule.every(15).minutes.do(check_all_vps)
    logger.info("Scheduled VPS check every 15 minutes")
    print("Scheduled VPS check every 15 minutes")
    sys.stdout.flush()
    
    print("===== VPS monitoring script is running =====")
    sys.stdout.flush()
    
    heartbeat_count = 0
    while True:
        schedule.run_pending()
        time.sleep(60)
        heartbeat_count += 1
        if heartbeat_count % 5 == 0:  # 每5分钟输出一次心跳信息
            print(f"Heartbeat: Script is still running. Uptime: {heartbeat_count} minutes")
            sys.stdout.flush()

if __name__ == "__main__":
    main()