import threading
import json
from flask import Flask, jsonify, request, send_from_directory
import sqlite3
import socket
import secrets
import time
import os
import logging
import re  # 添加正则表达式模块

# 在文件顶部添加常量定义
AVATAR_BASE_DIR = "avatar"  # 头像存储基础目录
DEFAULT_AVATAR = "default_avatar.png"  # 默认头像文件名

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

app = Flask(__name__)
socket_server = socket.socket()
socket_server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)

# 修改数据结构：使用用户名作为主键
active_connections = {}  # {username: {'conn': conn, 'ip': ip, 'last_active': timestamp}}
chat_connections = []    # 所有活跃连接列表
tokens = {}              # 令牌管理

def get_db_connection():
    conn = sqlite3.connect("usr.db")
    conn.row_factory = sqlite3.Row
    return conn

def isuserxist(name):
    cn = get_db_connection()
    csr = cn.cursor()
    csr.execute('SELECT * FROM users WHERE name = ?', (name,))
    rst = csr.fetchone()
    cn.close()
    return rst is not None

def ispsswdright(name, passwd):
    cn = get_db_connection()
    csr = cn.cursor()
    csr.execute("SELECT passwd FROM users WHERE name = ?", (name,))
    result = csr.fetchone()
    cn.close()
    return result and result[0] == passwd

# 用户名验证函数
def validate_username(username):
    """验证用户名是否符合规则"""
    # 长度在2-20个字符
    if len(username) < 2 or len(username) > 20:
        return False
    
    # 只能包含特定字符：字母、数字、_!@#$%^&+-~?
    if not re.match(r'^[a-zA-Z0-9_!@#$%^&+\-~?]+$', username):
        return False
    
    # 不能以数字开头
    if username[0].isdigit():
        return False
    
    return True

# 修改获取头像函数
def get_avatar(username):
    cn = get_db_connection()
    csr = cn.cursor()
    csr.execute("SELECT avatar FROM users WHERE name = ?", (username,))
    result = csr.fetchone()
    cn.close()
    if result:
        avatar_file = result[0]
        # 返回HTTP URL格式的头像路径
        return f"/avatar/{username}/{avatar_file}"
    
    # 返回默认头像URL
    return f"/avatar/default/{DEFAULT_AVATAR}"

def register_user(usr, pwd, avatar="default_avatar.png"):
    conn = get_db_connection()
    try:
        cursor = conn.cursor()
        # 创建用户专属头像目录
        user_avatar_dir = os.path.join(AVATAR_BASE_DIR, usr)
        os.makedirs(user_avatar_dir, exist_ok=True)
        cursor.execute("INSERT INTO users (name, passwd, avatar) VALUES (?, ?, ?)",
                       (usr, pwd, avatar))
        conn.commit()
        return {"type": "register_1", "error": False, "message": "User registered successfully"}
    except sqlite3.IntegrityError:
        return {"type": "register_0", "success": False, "message": "Username already exists"}
    except sqlite3.Error as e:
        return {"type": "register_-1", "error": True, "message": str(e)}
    finally:
        conn.close()

def generate_token(username):
    token = secrets.token_hex(16)
    tokens[token] = {'username': username, 'timestamp': time.time()}
    return token

def validate_token(token):
    if token in tokens:
        if time.time() - tokens[token]['timestamp'] < 3600:
            tokens[token]['timestamp'] = time.time()
            return tokens[token]['username']
    return None

@app.route("/api/register", methods=['POST'])
def register1():
    vl = request.get_json()
    usr = vl.get('username')
    pwd = vl.get('password')
    avatar = vl.get('avatar', 'default_avatar.png')
    
    # 验证用户名
    if not validate_username(usr):
        return jsonify({
            "type": "register_-2",
            "success": False,
            "message": "Invalid username. Must be 2-20 characters, start with a letter or symbol, and contain only letters, numbers, or symbols: _!@#$%^&+-~?"
        }), 400

    if avatar and not (avatar.endswith('.png') or avatar.endswith('.jpg')):
        return jsonify({
            "type": "register_0",
            "success": False,
            "message": "Invalid avatar format. Only .png or .jpg allowed"
        }), 400

    result = register_user(usr, pwd, avatar)
    if result.get('success', True):  # 成功注册时返回的字典有'success'键
        return jsonify(result)
    else:
        return jsonify(result), 403 if result['message'] == "Username already exists" else 500

@app.route("/api/login", methods=['POST'])
def login():
    data = request.get_json()
    username = data['username']

    # 检查账号是否已登录
    if username in active_connections:
        # 检查连接是否仍然活跃
        conn_info = active_connections[username]
        try:
            # 发送测试消息检查连接是否有效
            conn_info['conn'].sendall(json.dumps({"type": "ping"}).encode('utf-8'))
            logger.info(f"用户 {username} 的连接仍然活跃")
            return jsonify({
                "type": "login",
                "status": "error_0",
                "message": "Account already logged in"
            }), 409
        except:
            # 连接已失效，清理旧连接
            logger.warning(f"清理无效连接: {username}")
            if username in active_connections:
                del active_connections[username]
            if conn_info['conn'] in chat_connections:
                chat_connections.remove(conn_info['conn'])

    if isuserxist(username) and ispsswdright(username, data['password']):
        token = generate_token(username)
        avatar = get_avatar(username)
        return jsonify({
            "type": "login_1",
            "status": "success",
            "token": token,
            "avatar": avatar
        })
    return jsonify({
        "type": "login_0",
        "status": "error",
        "message": "Invalid credentials"
    }), 401

# 添加头像静态路由
@app.route("/avatar/<username>/<filename>", methods=['GET'])
def serve_avatar(username, filename):
    try:
        avatar_dir = os.path.join(AVATAR_BASE_DIR, username)
        return send_from_directory(avatar_dir, filename)
    except FileNotFoundError:
        # 如果找不到头像，返回默认头像
        return send_from_directory(AVATAR_BASE_DIR, DEFAULT_AVATAR)

@app.route("/api/chat", methods=['POST'])
def chat():
    token = request.headers.get('Authorization')
    username = validate_token(token)
    if not username:
        return jsonify({"type": "chat", "status": "error"}), 401

    data = request.get_json()
    message = {
        "type": "chat",
        "user": username,
        "message": data['message'],
        "avatar": get_avatar(username)
    }
    broadcast_message(message)
    return jsonify({"type": "chat", "status": "success"})

def broadcast_message(message, sender=None):
    for conn in chat_connections[:]:  # 使用副本迭代
        try:
            conn.sendall(json.dumps(message).encode('utf-8'))
        except:
            # 连接异常时移除
            for uname, info in list(active_connections.items()):
                if info['conn'] == conn:
                    logger.warning(f"广播时移除无效连接: {uname}")
                    del active_connections[uname]
                    break
            if conn in chat_connections:
                chat_connections.remove(conn)

def handle_socket_message(data, addr, conn):
    try:
        action = data.get('type')
        if action == 'register':
            username = data.get('username')
            password = data.get('password')
            avatar = data.get('avatar', 'default_avatar.png')
            
            # 验证用户名
            if not validate_username(username):
                response = {
                    "type": "register_-2",
                    "success": False,
                    "message": "Invalid username format"
                }
                conn.sendall(json.dumps(response).encode('utf-8'))
                return response

            if avatar and not (avatar.endswith('.png') or avatar.endswith('.jpg')):
                response = {
                    "type": "register_-3",
                    "success": False,
                    "message": "Invalid avatar format. Only .png or .jpg allowed"
                }
                conn.sendall(json.dumps(response).encode('utf-8'))
                return response

            result = register_user(username, password, avatar)
            conn.sendall(json.dumps(result).encode('utf-8'))
            return result

        elif action == 'login':
            username = data['username']
            password = data['password']

            # 检查账号是否已登录且连接有效
            if username in active_connections:
                conn_info = active_connections[username]
                try:
                    # 测试连接是否仍然有效
                    conn_info['conn'].sendall(json.dumps({"type": "ping"}).encode('utf-8'))
                    logger.info(f"用户 {username} 尝试登录但已有活跃连接")
                    response = {
                        "type": "login",
                        "status": "error_-1",
                        "message": "Account already logged in"
                    }
                    conn.sendall(json.dumps(response).encode('utf-8'))
                    return response
                except:
                    # 连接已失效，清理旧连接
                    logger.warning(f"清理无效连接后允许登录: {username}")
                    if username in active_connections:
                        del active_connections[username]
                    if conn_info['conn'] in chat_connections:
                        chat_connections.remove(conn_info['conn'])

            if isuserxist(username) and ispsswdright(username, password):
                # 添加新连接
                active_connections[username] = {'conn': conn, 'ip': addr[0], 'last_active': time.time()}
                if conn not in chat_connections:
                    chat_connections.append(conn)

                token = generate_token(username)
                avatar = get_avatar(username)
                response = {
                    "type": "login",
                    "status": "success",
                    "message": "Login successful",
                    "token": token,
                    "username": username,
                    "avatar": avatar
                }
                conn.sendall(json.dumps(response).encode('utf-8'))
                logger.info(f"用户 {username} 登录成功")
                return response
            else:
                response = {
                    "type": "login",
                    "status": "error_0",
                    "message": "Invalid credentials"
                }
                conn.sendall(json.dumps(response).encode('utf-8'))
                return response

        elif action == 'chat':
            token = data.get('token')
            if not token:
                response = {
                    "type": "chat",
                    "status": "error_Mt",
                    "message": "Missing token"
                }
                conn.sendall(json.dumps(response).encode('utf-8'))
                return response

            username = validate_token(token)
            if not username:
                response = {
                    "type": "chat",
                    "status": "error_It",
                    "message": "Invalid token"
                }
                conn.sendall(json.dumps(response).encode('utf-8'))
                return response

            if username not in active_connections:
                response = {
                    "type": "chat",
                    "status": "error_Nli",
                    "message": "Not logged in"
                }
                conn.sendall(json.dumps(response).encode('utf-8'))
                return response

            # 更新最后活跃时间
            active_connections[username]['last_active'] = time.time()

            message = {
                "type": "chat",
                "user": username,
                "message": data['message'],
                "avatar": get_avatar(username)
            }
            broadcast_message(message)
            response = {"type": "chat", "status": "success"}
            conn.sendall(json.dumps(response).encode('utf-8'))
            return response

        elif action == 'heartbeat':
            # 心跳检测
            token = data.get('token')
            if token:
                username = validate_token(token)
                if username and username in active_connections:
                    # 更新最后活跃时间
                    active_connections[username]['last_active'] = time.time()
                    response = {"type": "heartbeat", "status": "success"}
                    conn.sendall(json.dumps(response).encode('utf-8'))
                    return response
            return {"type": "heartbeat", "status": "error"}

    except Exception as e:
        logger.error(f"处理消息时出错: {str(e)}")
        response = {
            "status": "error",
            "message": str(e)
        }
        try:
            conn.sendall(json.dumps(response).encode('utf-8'))
        except:
            pass
        return response

def check_inactive_connections():
    """定期检查不活跃的连接并清理"""
    while True:
        time.sleep(60)  # 每分钟检查一次
        current_time = time.time()
        inactive_users = []

        for username, info in list(active_connections.items()):
            # 5分钟无活动视为不活跃
            if current_time - info['last_active'] > 300:
                logger.warning(f"检测到不活跃用户: {username}, 最后活跃: {current_time - info['last_active']}秒前")
                inactive_users.append(username)

        for username in inactive_users:
            info = active_connections[username]
            try:
                info['conn'].close()
            except:
                pass
            if username in active_connections:
                del active_connections[username]
            if info['conn'] in chat_connections:
                chat_connections.remove(info['conn'])
            logger.info(f"已清理不活跃用户: {username}")

def run_socket_server():
    socket_server.bind(("0.0.0.0", 8889))
    socket_server.listen()
    logger.info("Socket server running on port 8889")

    # 启动连接检查线程
    threading.Thread(target=check_inactive_connections, daemon=True).start()

    while True:
        conn, addr = socket_server.accept()
        logger.info(f"Socket client connected: {addr}")
        try:
            while True:
                data = conn.recv(1024)
                if not data:
                    break

                try:
                    decoded_data = data.decode('utf-8', errors='ignore')
                    json_data = json.loads(decoded_data)
                    response = handle_socket_message(json_data, addr, conn)
                except json.JSONDecodeError:
                    response = {
                        "type": "error",
                        "status": "error",
                        "message": "Invalid JSON"
                    }
                    conn.sendall(json.dumps(response).encode('utf-8'))
                except Exception as e:
                    logger.error(f"处理数据时出错: {str(e)}")
                    response = {
                        "type": "error",
                        "status": "error",
                        "message": str(e)
                    }
                    conn.sendall(json.dumps(response).encode('utf-8'))
        except (ConnectionResetError, BrokenPipeError):
            logger.warning(f"Client {addr} disconnected abruptly")
        finally:
            # 清理断开的连接
            for username, info in list(active_connections.items()):
                if info['conn'] == conn:
                    del active_connections[username]
                    logger.info(f"用户 {username} 断开连接")
                    break

            if conn in chat_connections:
                chat_connections.remove(conn)

            try:
                conn.close()
            except:
                pass
            logger.info(f"Connection closed for {addr}")

if __name__ == '__main__':
    with get_db_connection() as conn:
        conn.execute('''CREATE TABLE IF NOT EXISTS users
                      (name TEXT PRIMARY KEY, 
                       passwd TEXT,
                       avatar TEXT DEFAULT 'default_avatar.png')''')
    # 确保头像目录存在
    os.makedirs(AVATAR_BASE_DIR, exist_ok=True)
    threading.Thread(target=run_socket_server, daemon=True).start()
    app.run(port=5001, host='0.0.0.0')