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
import requests
import uuid
import traceback

# 设置日志
logging.basicConfig(level=logging.DEBUG, 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'
UPDATE_SERVER_URL = "https://update.uqdm.com"

# 设置访问密码和令牌
ACCESS_PASSWORD = "131417aa"
ACCESS_TOKEN = "131417"  # 这应该是一个安全的随机字符串

def init_db():
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    c.execute('''CREATE TABLE IF NOT EXISTS clients
                 (unique_id TEXT PRIMARY KEY, client_id TEXT, group_name TEXT, last_seen TIMESTAMP, 
                  system_info TEXT, command_status 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
                 (unique_id TEXT, 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'), ('离线')")
    
    conn.commit()
    conn.close()
    logger.info("数据库初始化完成")

# 登录检查装饰器
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

# 令牌验证装饰器
def token_required(f):
    @wraps(f)
    def decorated_function(*args, **kwargs):
        token = request.headers.get('Authorization')
        if token != f'Bearer {ACCESS_TOKEN}':
            return jsonify({"error": "无效或缺失的令牌"}), 401
        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"客户端连接: {request.sid}")

@socketio.on('disconnect')
def on_disconnect():
    logger.info(f"客户端断开连接: {request.sid}")

@socketio.on('register')
def on_register(data):
    client_id = data['client_id']
    unique_id = data.get('unique_id')
    system_info = json.dumps(data.get('system_info', {}))
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    try:
        if unique_id:
            c.execute("SELECT group_name FROM clients WHERE unique_id = ?", (unique_id,))
        else:
            c.execute("SELECT unique_id, group_name FROM clients WHERE client_id = ?", (client_id,))
        existing = c.fetchone()
        
        if existing:
            if unique_id:
                existing_group = existing[0]
            else:
                unique_id, existing_group = existing
            c.execute("""
                UPDATE clients 
                SET last_seen = ?, system_info = ?, command_status = ?, program_status = ?, client_id = ?
                WHERE unique_id = ?
            """, (datetime.now().isoformat(), system_info, json.dumps({'status': 'N/A'}), 'N/A', client_id, unique_id))
        else:
            unique_id = str(uuid.uuid4())
            existing_group = 'default'
            c.execute("""
                INSERT INTO clients 
                (unique_id, client_id, group_name, last_seen, system_info, command_status, program_status) 
                VALUES (?, ?, ?, ?, ?, ?, ?)
            """, (unique_id, client_id, existing_group, datetime.now().isoformat(), system_info, json.dumps({'status': 'N/A'}), 'N/A'))

        conn.commit()
        logger.info(f"客户端注册或更新: {client_id}, unique_id: {unique_id}, 分组: {existing_group}")
        
        # 将客户端加入到以其 unique_id 命名的房间
        socketio.server.enter_room(request.sid, unique_id)
        
        # 通知客户端其 unique_id 和分组
        socketio.emit('update_client_info', {'unique_id': unique_id, 'group': existing_group}, room=request.sid)
        
    except Exception as e:
        logger.error(f"注册客户端 {client_id} 时出错: {e}")
    finally:
        conn.close()

@socketio.on('heartbeat')
def on_heartbeat(data):
    unique_id = data.get('unique_id')
    client_id = data.get('client_id')
    system_info = json.dumps(data.get('system_info', {}))
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    try:
        if unique_id:
            c.execute("""
                UPDATE clients 
                SET last_seen = ?, system_info = ? 
                WHERE unique_id = ?
            """, (datetime.now().isoformat(), system_info, unique_id))
        elif client_id:
            c.execute("""
                UPDATE clients 
                SET last_seen = ?, system_info = ? 
                WHERE client_id = ?
            """, (datetime.now().isoformat(), system_info, client_id))
        else:
            logger.error("收到的心跳数据中没有 unique_id 或 client_id")
        conn.commit()
    except Exception as e:
        logger.error(f"更新心跳时出错: {e}")
    finally:
        conn.close()

@socketio.on('status_update')
def on_status_update(data):
    unique_id = data.get('unique_id')
    client_id = data.get('client_id')
    system_info = json.dumps(data.get('system_info', {}))
    program_status = data.get('program_status', 'unknown')
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    try:
        if unique_id:
            c.execute("UPDATE clients SET system_info = ?, program_status = ?, last_seen = ? WHERE unique_id = ?", 
                      (system_info, program_status, datetime.now().isoformat(), unique_id))
        elif client_id:
            c.execute("UPDATE clients SET system_info = ?, program_status = ?, last_seen = ? WHERE client_id = ?", 
                      (system_info, program_status, datetime.now().isoformat(), client_id))
        else:
            logger.error("收到的状态更新中没有 unique_id 或 client_id")
        conn.commit()
        logger.info(f"更新了客户端状态: {unique_id or client_id}")
        socketio.emit('client_status_updated', {
            'unique_id': unique_id,
            'client_id': client_id,
            'system_info': json.loads(system_info),
            'program_status': program_status
        })
    except Exception as e:
        logger.error(f"更新状态时出错: {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("PRAGMA table_info(command_results)")
        columns = [column[1] for column in c.fetchall()]
        c.execute("SELECT * FROM command_results ORDER BY timestamp DESC LIMIT 10")
        results = c.fetchall()
        return jsonify([dict(zip(columns, row)) for row in results])
    except sqlite3.OperationalError as e:
        logger.error(f"获取命令结果时数据库错误: {e}")
        return jsonify({'error': '数据库错误', 'message': str(e)}), 500
    except Exception as e:
        logger.error(f"获取命令结果时意外错误: {e}")
        return jsonify({'error': '意外的服务器错误', 'message': str(e)}), 500
    finally:
        conn.close()

@socketio.on('command_result')
def on_command_result(data):
    unique_id = data.get('unique_id')
    client_id = data.get('client_id')
    command = data.get('command')
    result = data.get('result', {})
    logger.info(f"收到来自 {client_id} (unique_id: {unique_id}) 的命令结果: {command} - {result}")
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    try:
        command_status_json = json.dumps(result)
        c.execute("UPDATE clients SET command_status = ?, program_status = ? WHERE unique_id = ?", 
                  (command_status_json, result.get('program_status', 'N/A'), unique_id))
        
        c.execute("INSERT INTO command_results (unique_id, client_id, command, status, output, error, program_status) VALUES (?, ?, ?, ?, ?, ?, ?)",
                  (unique_id, client_id, command, result.get('status'), result.get('output'), result.get('error'), result.get('program_status')))
        
        conn.commit()
        socketio.emit('update_client_status', {'unique_id': unique_id, 'client_id': client_id, 'command_status': result, 'program_status': result.get('program_status', 'N/A')}, broadcast=True)
        logger.info(f"广播了客户端 {client_id} (unique_id: {unique_id}) 的状态更新")
    except Exception as e:
        logger.error(f"更新客户端 {client_id} (unique_id: {unique_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 unique_id, client_id, group_name, last_seen FROM clients")
        offline_clients = []
        for unique_id, client_id, group_name, last_seen in c.fetchall():
            logger.debug(f"检查客户端: {client_id}, last_seen: {last_seen}, group: {group_name}")
            if last_seen:
                try:
                    last_seen_date = datetime.fromisoformat(last_seen)
                    if last_seen_date < offline_threshold:
                        if group_name != '离线':
                            c.execute("UPDATE clients SET group_name = '离线' WHERE unique_id = ?", (unique_id,))
                        offline_clients.append(client_id)
                except ValueError as e:
                    logger.error(f"解析 last_seen 时出错 (client_id: {client_id}): {e}")
                    if group_name != '离线':
                        c.execute("UPDATE clients SET group_name = '离线' WHERE unique_id = ?", (unique_id,))
                    offline_clients.append(client_id)
            else:
                logger.warning(f"客户端 {client_id} 的 last_seen 为 None")
                if group_name != '离线':
                    c.execute("UPDATE clients SET group_name = '离线' WHERE unique_id = ?", (unique_id,))
                offline_clients.append(client_id)
        conn.commit()
        logger.info(f"检测到 {len(offline_clients)} 个离线客户端")
        return offline_clients
    except Exception as e:
        logger.error(f"检查离线客户端时出错: {e}")
        logger.error(traceback.format_exc())
        return []
    finally:
        conn.close()

@app.route('/api/clients', methods=['GET'])
@login_required
def get_clients():
    try:
        offline_clients = check_offline_clients()
        conn = sqlite3.connect(DB_NAME)
        c = conn.cursor()
        c.execute("SELECT unique_id, client_id, group_name, system_info, command_status, program_status FROM clients")
        clients = {}
        for row in c.fetchall():
            unique_id, client_id, group_name, system_info, command_status, program_status = row
            logger.debug(f"处理客户端: {client_id}, group: {group_name}, system_info: {system_info}, command_status: {command_status}, program_status: {program_status}")
            
            # 添加类型检查
            if not isinstance(unique_id, str):
                logger.warning(f"客户端 {client_id} 的 unique_id 不是字符串: {type(unique_id)}")
                unique_id = str(unique_id) if unique_id is not None else ""
            
            if not isinstance(group_name, str):
                logger.warning(f"客户端 {client_id} 的 group_name 不是字符串: {type(group_name)}")
                group_name = str(group_name) if group_name is not None else "default"
            
            try:
                system_info_json = json.loads(system_info) if system_info else {}
            except json.JSONDecodeError:
                system_info_json = {}
                logger.warning(f"客户端 {client_id} 的系统信息 JSON 无效")
            
            try:
                command_status_json = json.loads(command_status) if command_status else {'status': 'N/A'}
            except json.JSONDecodeError:
                command_status_json = {'status': 'N/A'}
                logger.warning(f"客户端 {client_id} 的命令状态 JSON 无效")
            
            clients[client_id] = {
                'unique_id': unique_id,
                'group': group_name or 'default',
                '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()]
        
        logger.info(f"获取到 {len(clients)} 个客户端和 {len(groups)} 个分组")
        
        return jsonify({
            'clients': clients,
            'groups': groups,
            'offline_clients': offline_clients
        })
    except Exception as e:
        logger.error(f"获取客户端时出错: {e}")
        logger.error(traceback.format_exc())
        return jsonify({'error': '服务器内部错误', 'message': str(e)}), 500
    finally:
        if 'conn' in locals():
            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"执行命令: {command}, 目标客户端: {client_ids}")
        logger.info(f"命令参数: {params}")
        
        if command == 'force_update':
            try:
                response = requests.get(f"{UPDATE_SERVER_URL}/check_update/0.0.0")
                if response.status_code == 200:
                    update_info = response.json()
                    latest_version = update_info.get('latest_version')
                    if latest_version:
                        params['update_url'] = f"{UPDATE_SERVER_URL}/download/{latest_version}"
                    else:
                        return jsonify({"status": "error", "message": "无法获取最新版本"}), 400
                else:
                    return jsonify({"status": "error", "message": "无法检查更新"}), 400
            except requests.RequestException as e:
                logger.error(f"检查更新时出错: {e}")
                return jsonify({"status": "error", "message": "无法检查更新"}), 500
        
        conn = sqlite3.connect(DB_NAME)
        c = conn.cursor()
        try:
            for client_id in client_ids:
                c.execute("SELECT unique_id FROM clients WHERE client_id = ?", (client_id,))
                result = c.fetchone()
                if result:
                    unique_id = result[0]
                    socketio.emit('execute_command', {'command': command, 'params': params}, room=unique_id)
                    affected_clients += 1
                    logger.info(f"命令 '{command}' 已发送给客户端 {client_id} (unique_id: {unique_id})")
                else:
                    logger.warning(f"数据库中未找到客户端 {client_id}")
        finally:
            conn.close()
        
        return jsonify({"status": "命令已发送", "affected_clients": affected_clients})
    except Exception as e:
        logger.error(f"执行命令时出错: {e}")
        logger.error(traceback.format_exc())
        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": "设置已更新"})

@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:
        c.execute("INSERT OR IGNORE INTO groups (name) VALUES (?)", (new_group,))
        
        for client_id in client_ids:
            c.execute("UPDATE clients SET group_name = ? WHERE client_id = ?", (new_group, client_id))
        conn.commit()
        return jsonify({"status": "客户端分组已更新"})
    except Exception as e:
        logger.error(f"更新客户端分组时出错: {e}")
        logger.error(traceback.format_exc())
        return jsonify({'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": "分组已添加"})
    except Exception as e:
        logger.error(f"添加分组时出错: {e}")
        logger.error(traceback.format_exc())
        return jsonify({'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": "无法删除默认或离线分组"}), 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": "分组已删除"})
    except Exception as e:
        logger.error(f"删除分组时出错: {e}")
        logger.error(traceback.format_exc())
        return jsonify({'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": "配置已保存"})
    except Exception as e:
        logger.error(f"保存配置时出错: {e}")
        logger.error(traceback.format_exc())
        return jsonify({'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"获取配置时出错: {e}")
        logger.error(traceback.format_exc())
        return jsonify({'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": "配置已删除"})
    except Exception as e:
        logger.error(f"删除配置时出错: {e}")
        logger.error(traceback.format_exc())
        return jsonify({'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 client_id = ?", (client_id,))
        conn.commit()
        return jsonify({"status": "success", "message": f"{len(client_ids)} 个客户端已删除"})
    except Exception as e:
        logger.error(f"删除客户端时出错: {e}")
        logger.error(traceback.format_exc())
        return jsonify({'error': '服务器内部错误'}), 500
    finally:
        conn.close()

@app.route('/api/check_db', methods=['GET'])
@login_required
def check_db():
    conn = sqlite3.connect(DB_NAME)
    c = conn.cursor()
    try:
        c.execute("SELECT * FROM clients")
        clients = c.fetchall()
        c.execute("PRAGMA table_info(clients)")
        columns = [column[1] for column in c.fetchall()]
        
        client_data = [dict(zip(columns, client)) for client in clients]
        
        return jsonify({
            'status': 'success',
            'client_count': len(clients),
            'columns': columns,
            'sample_data': client_data[:5] if client_data else []
        })
    except Exception as e:
        logger.error(f"检查数据库时出错: {e}")
        logger.error(traceback.format_exc())
        return jsonify({'error': '服务器内部错误', 'message': str(e)}), 500
    finally:
        conn.close()

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