from flask import Flask, request, jsonify, render_template, session, redirect, url_for
from flask_socketio import SocketIO, emit
import json
import os
import sqlite3
from datetime import datetime, timedelta
import logging
from functools import wraps

# 设置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

app = Flask(__name__)
app.secret_key = os.urandom(24)  # 用于加密 session
socketio = SocketIO(app, cors_allowed_origins="*")

DB_NAME = 'client_management.db'

# 设置访问密码（在实际应用中，你应该使用更安全的方式存储密码）
ACCESS_PASSWORD = "131417aa"

# 用于存储客户端的最后命令执行状态
client_command_status = {}

def init_db():
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    c.execute('''CREATE TABLE IF NOT EXISTS clients
                 (id TEXT PRIMARY KEY, group_name TEXT, last_seen TIMESTAMP, 
                  system_info TEXT, program_status TEXT)''')
    c.execute('''CREATE TABLE IF NOT EXISTS groups
                 (name TEXT PRIMARY KEY)''')
    c.execute('''CREATE TABLE IF NOT EXISTS configurations
                 (name TEXT PRIMARY KEY, data TEXT)''')
    c.execute('''CREATE TABLE IF NOT EXISTS command_results
                 (client_id TEXT, command TEXT, status TEXT, 
                  output TEXT, error TEXT, program_status TEXT, 
                  timestamp DATETIME DEFAULT CURRENT_TIMESTAMP)''')
    c.execute("INSERT OR IGNORE INTO groups (name) VALUES ('default'), ('离线')")
    
    # 检查并添加新列
    columns_to_check = ['system_info', 'program_status']
    for column in columns_to_check:
        c.execute(f"PRAGMA table_info(clients)")
        columns = [column_info[1] for column_info in c.fetchall()]
        if column not in columns:
            c.execute(f"ALTER TABLE clients ADD COLUMN {column} TEXT")
    
    conn.commit()
    conn.close()

# 登录检查装饰器
def login_required(f):
    @wraps(f)
    def decorated_function(*args, **kwargs):
        if not session.get('logged_in'):
            return redirect(url_for('login'))
        return f(*args, **kwargs)
    return decorated_function

@app.route('/login', methods=['GET', 'POST'])
def login():
    if request.method == 'POST':
        if request.form['password'] == ACCESS_PASSWORD:
            session['logged_in'] = True
            return redirect(url_for('index'))
        else:
            return render_template('login.html', error="密码错误")
    return render_template('login.html')

@app.route('/logout')
def logout():
    session.pop('logged_in', None)
    return redirect(url_for('login'))

@app.route('/')
@login_required
def index():
    return render_template('index.html')

@socketio.on('connect')
def on_connect():
    logger.info(f"Client connected: {request.sid}")

@socketio.on('disconnect')
def on_disconnect():
    logger.info(f"Client disconnected: {request.sid}")

@socketio.on('register')
def on_register(data):
    client_id = data['client_id']
    group = data.get('group', 'default')
    system_info = json.dumps(data.get('system_info', {}))
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    try:
        c.execute("INSERT OR REPLACE INTO clients (id, group_name, last_seen, system_info, program_status) VALUES (?, ?, ?, ?, ?)",
                  (client_id, group, datetime.now(), system_info, 'unknown'))
        conn.commit()
        logger.info(f"Client registered: {client_id}")
        # Join the client to a room with its client_id
        socketio.server.enter_room(request.sid, client_id)
    except Exception as e:
        logger.error(f"Error registering client {client_id}: {e}")
    finally:
        conn.close()

@socketio.on('heartbeat')
def on_heartbeat(data):
    client_id = data['client_id']
    system_info = json.dumps(data.get('system_info', {}))
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    try:
        c.execute("UPDATE clients SET last_seen = ?, system_info = ? WHERE id = ?", 
                  (datetime.now(), system_info, client_id))
        conn.commit()
    except Exception as e:
        logger.error(f"Error updating heartbeat for client {client_id}: {e}")
    finally:
        conn.close()

@socketio.on('status_update')
def on_status_update(data):
    client_id = data['client_id']
    system_info = json.dumps(data['system_info'])
    program_status = data.get('program_status', 'unknown')
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    try:
        c.execute("UPDATE clients SET system_info = ?, program_status = ?, last_seen = ? WHERE id = ?", 
                  (system_info, program_status, datetime.now(), client_id))
        conn.commit()
        logger.info(f"Status updated for client: {client_id}")
    except Exception as e:
        logger.error(f"Error updating status for client {client_id}: {e}")
    finally:
        conn.close()
        
@app.route('/api/command_results', methods=['GET'])
@login_required
def get_command_results():
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    try:
        c.execute("SELECT * FROM command_results ORDER BY timestamp DESC LIMIT 10")
        results = c.fetchall()
        return jsonify([{
            'client_id': r[0],
            'command': r[1],
            'status': r[2],
            'output': r[3],
            'error': r[4],
            'program_status': r[5],
            'timestamp': r[6]
        } for r in results])
    except sqlite3.OperationalError as e:
        logger.error(f"Database error in get_command_results: {e}")
        return jsonify({'error': 'Database error', 'message': str(e)}), 500
    except Exception as e:
        logger.error(f"Unexpected error in get_command_results: {e}")
        return jsonify({'error': 'Unexpected server error', 'message': str(e)}), 500
    finally:
        conn.close()

@socketio.on('command_result')
def on_command_result(data):
    client_id = data.get('client_id')
    command = data.get('command')
    result = data.get('result', {})
    logger.info(f"Received command result from {client_id}: {command} - {result}")
    
    # 更新内存中的命令执行状态
    client_command_status[client_id] = result
    
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    try:
        # 更新程序状态
        c.execute("UPDATE clients SET program_status = ? WHERE id = ?", 
                  (result.get('program_status', 'unknown'), client_id))
        
        # 插入命令结果到 command_results 表
        c.execute("INSERT INTO command_results (client_id, command, status, output, error, program_status) VALUES (?, ?, ?, ?, ?, ?)",
                  (client_id, command, result.get('status'), result.get('output'), result.get('error'), result.get('program_status')))
        
        conn.commit()
        socketio.emit('update_client_status', {'client_id': client_id, 'command_status': result, 'program_status': result.get('program_status', 'unknown')}, broadcast=True)
        logger.info(f"Broadcasted status update for client {client_id}")
    except Exception as e:
        logger.error(f"Error updating command result for client {client_id}: {e}")
    finally:
        conn.close()

def check_offline_clients():
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    try:
        offline_threshold = datetime.now() - timedelta(minutes=5)
        c.execute("SELECT id, group_name FROM clients WHERE last_seen < ?", (offline_threshold,))
        offline_clients = c.fetchall()
        for client_id, group_name in offline_clients:
            if group_name != '离线':
                c.execute("UPDATE clients SET group_name = '离线' WHERE id = ?", (client_id,))
        conn.commit()
        return [client_id for client_id, _ in offline_clients]
    except Exception as e:
        logger.error(f"Error checking offline clients: {e}")
        return []
    finally:
        conn.close()

@app.route('/api/clients', methods=['GET'])
@login_required
def get_clients():
    offline_clients = check_offline_clients()
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    try:
        c.execute("SELECT id, group_name, system_info, program_status FROM clients")
        clients = {}
        for row in c.fetchall():
            client_id, group_name, system_info, program_status = row
            try:
                system_info_json = json.loads(system_info) if system_info else {}
            except json.JSONDecodeError:
                system_info_json = {}
                logger.warning(f"Invalid system_info JSON for client {client_id}")
            
            # 从内存中获取命令执行状态
            command_status_json = client_command_status.get(client_id, {})
            
            clients[client_id] = {
                'group': group_name,
                'system_info': system_info_json,
                'command_status': command_status_json,
                'program_status': program_status or 'unknown'
            }
        
        c.execute("SELECT name FROM groups")
        groups = [row[0] for row in c.fetchall()]
        return jsonify({
            'clients': clients,
            'groups': groups,
            'offline_clients': offline_clients
        })
    except Exception as e:
        logger.error(f"Error getting clients: {e}")
        return jsonify({'error': 'Internal server error'}), 500
    finally:
        conn.close()

@app.route('/api/execute_command', methods=['POST'])
@login_required
def execute_command():
    try:
        data = request.json
        client_ids = data['client_ids']
        command = data['command']
        params = data.get('params', {})
        affected_clients = 0
        
        logger.info(f"Executing command: {command} for clients: {client_ids}")
        logger.info(f"Command params: {params}")
        
        for client_id in client_ids:
            try:
                socketio.emit('execute_command', {'command': command, 'params': params}, room=client_id)
                affected_clients += 1
                logger.info(f"Command '{command}' sent to client {client_id}")
            except Exception as e:
                logger.error(f"Error sending command to client {client_id}: {e}")
        
        return jsonify({"status": "Command sent", "affected_clients": affected_clients})
    except Exception as e:
        logger.error(f"Error in execute_command: {e}")
        return jsonify({"status": "error", "message": str(e)}), 500

@app.route('/api/update_settings', methods=['POST'])
@login_required
def update_settings():
    data = request.json
    socketio.emit('update_settings', data, broadcast=True)
    return jsonify({"status": "Settings updated"})

@app.route('/api/update_client_group', methods=['POST'])
@login_required
def update_client_group():
    data = request.json
    client_ids = data['client_ids']
    new_group = data['group']
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    try:
        for client_id in client_ids:
            c.execute("UPDATE clients SET group_name = ? WHERE id = ?", (new_group, client_id))
        c.execute("INSERT OR IGNORE INTO groups (name) VALUES (?)", (new_group,))
        conn.commit()
        return jsonify({"status": "Client group updated"})
    except Exception as e:
        logger.error(f"Error updating client group: {e}")
        return jsonify({'error': 'Internal server error'}), 500
    finally:
        conn.close()

@app.route('/api/add_group', methods=['POST'])
@login_required
def add_group():
    data = request.json
    new_group = data['group']
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    try:
        c.execute("INSERT OR IGNORE INTO groups (name) VALUES (?)", (new_group,))
        conn.commit()
        return jsonify({"status": "Group added"})
    except Exception as e:
        logger.error(f"Error adding group: {e}")
        return jsonify({'error': 'Internal server error'}), 500
    finally:
        conn.close()

@app.route('/api/delete_group', methods=['POST'])
@login_required
def delete_group():
    data = request.json
    group_to_delete = data['group']
    if group_to_delete in ['default', '离线']:
        return jsonify({"status": "Cannot delete default or offline group"}), 400
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    try:
        c.execute("UPDATE clients SET group_name = 'default' WHERE group_name = ?", (group_to_delete,))
        c.execute("DELETE FROM groups WHERE name = ?", (group_to_delete,))
        conn.commit()
        return jsonify({"status": "Group deleted"})
    except Exception as e:
        logger.error(f"Error deleting group: {e}")
        return jsonify({'error': 'Internal server error'}), 500
    finally:
        conn.close()

@app.route('/api/save_config', methods=['POST'])
@login_required
def save_config():
    data = request.json
    config_name = data.get('name')
    config_data = json.dumps(data.get('config'))
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    try:
        c.execute("INSERT OR REPLACE INTO configurations (name, data) VALUES (?, ?)", (config_name, config_data))
        conn.commit()
        return jsonify({"status": "success", "message": "Configuration saved"})
    except Exception as e:
        logger.error(f"Error saving configuration: {e}")
        return jsonify({'error': 'Internal server error'}), 500
    finally:
        conn.close()

@app.route('/api/get_configs', methods=['GET'])
@login_required
def get_configs():
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    try:
        c.execute("SELECT name, data FROM configurations")
        configs = {row[0]: json.loads(row[1]) for row in c.fetchall()}
        return jsonify(configs)
    except Exception as e:
        logger.error(f"Error getting configurations: {e}")
        return jsonify({'error': 'Internal server error'}), 500
    finally:
        conn.close()

@app.route('/api/delete_config', methods=['POST'])
@login_required
def delete_config():
    config_name = request.json.get('name')
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    try:
        c.execute("DELETE FROM configurations WHERE name = ?", (config_name,))
        conn.commit()
        return jsonify({"status": "success", "message": "Configuration deleted"})
    except Exception as e:
        logger.error(f"Error deleting configuration: {e}")
        return jsonify({'error': 'Internal server error'}), 500
    finally:
        conn.close()

@app.route('/api/delete_clients', methods=['POST'])
@login_required
def delete_clients():
    client_ids = request.json.get('client_ids', [])
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    try:
        for client_id in client_ids:
            c.execute("DELETE FROM clients WHERE id = ?", (client_id,))
            # 同时删除内存中的命令执行状态
            client_command_status.pop(client_id, None)
        conn.commit()
        return jsonify({"status": "success", "message": f"{len(client_ids)} clients deleted"})
    except Exception as e:
        logger.error(f"Error deleting clients: {e}")
        return jsonify({'error': 'Internal server error'}), 500
    finally:
        conn.close()

@app.route('/api/test_db', methods=['GET'])
@login_required
def test_db():
    try:
        conn = sqlite3.connect(DB_NAME)
        c = conn.cursor()
        c.execute("SELECT name FROM sqlite_master WHERE type='table';")
        tables = c.fetchall()
        conn.close()
        return jsonify({"status": "success", "tables": [table[0] for table in tables]})
    except Exception as e:
        return jsonify({"status": "error", "message": str(e)}), 500

if __name__ == '__main__':
    init_db()  # 确保在启动服务器之前初始化数据库
    socketio.run(app, debug=True, host='0.0.0.0', port=5003)