import socket
import select
import uuid
import pickle
import os
import threading
from flask import Flask, request, jsonify
from flask_socketio import SocketIO, emit, join_room, leave_room

# Flask app for REST API and WebSocket
app = Flask(__name__)
socketio = SocketIO(app, cors_allowed_origins="*")  # Allow all origins

CHATPORT = 8181
BUFFER_SIZE = 1024
MAX_CONNECT_QUEUE = 10
DATA_FILE = 'server_data.pkl'
UUID_FILE = 'uuids.pkl'

# Global variables for chat server
socket_list = []
client_info = {}
chat_rooms = {}
client_states = {}
uuid_dict = {}
message_history = {}  # New: Store message history for each chat room
server_socket = None
quitFlag = True

# Utility functions
def save_data(client_info, chat_rooms, client_states):
    serializable_client_info = {info["user_id"]: info for sock, info in client_info.items()}
    serializable_chat_rooms = {info["user_id"]: chat_rooms[sock] for sock, info in client_info.items()}
    serializable_client_states = {info["user_id"]: client_states[sock] for sock, info in client_info.items()}
    with open(DATA_FILE, 'wb') as f:
        pickle.dump((serializable_client_info, serializable_chat_rooms, serializable_client_states), f)

def load_data():
    if os.path.exists(DATA_FILE):
        with open(DATA_FILE, 'rb') as f:
            serializable_client_info, serializable_chat_rooms, serializable_client_states = pickle.load(f)
            client_info = {uuid.UUID(user_id): info for user_id, info in serializable_client_info.items()}
            chat_rooms = {uuid.UUID(user_id): room for user_id, room in serializable_chat_rooms.items()}
            client_states = {uuid.UUID(user_id): state for user_id, state in serializable_client_states.items()}
            return client_info, chat_rooms, client_states
    return {}, {}, {}

def save_uuids(uuid_dict):
    with open(UUID_FILE, 'wb') as f:
        pickle.dump(uuid_dict, f)

def load_uuids():
    if os.path.exists(UUID_FILE):
        with open(UUID_FILE, 'rb') as f:
            return pickle.load(f)
    return {}

def get_client_uuid(client_socket, uuid_dict):
    client_address = client_socket.getpeername()
    client_key = f"{client_address[0]}_{client_address[1]}"

    if client_key in uuid_dict:
        user_id = uuid_dict[client_key]
    else:
        user_id = str(uuid.uuid4())
        uuid_dict[client_key] = user_id
        save_uuids(uuid_dict)

    return user_id

def handle_message(username, chat_room, message):
    """Store the message in history and broadcast it to all clients."""
    # Store the message in the chat room's history
    if chat_room not in message_history:
        message_history[chat_room] = []
    message_history[chat_room].append(f"{username}: {message}")

    # Broadcast the message to WebSocket clients in the chat room
    socketio.emit('message', {'username': username, 'message': message}, room=chat_room)

    # Broadcast the message to socket-based clients in the chat room
    send_message = f"{username} >> {message}".encode('utf-8')
    for client_sock in socket_list:
        if client_sock != server_socket and chat_rooms.get(client_sock) == chat_room:
            client_sock.send(send_message)

# REST API Endpoints
@app.route('/send_message', methods=['POST'])
def send_message():
    """Send a message to a specific chat room."""
    data = request.json
    username = data.get('username')
    chat_room = data.get('chat_room', 'default')
    message = data.get('message')

    if not username or not message:
        return jsonify({"error": "Username and message are required"}), 400

    # Use the centralized message handling function
    handle_message(username, chat_room, message)

    return jsonify({"status": "Message sent and saved in history"}), 200

@app.route('/create_chatroom', methods=['POST'])
def create_chatroom():
    """Create a new chat room."""
    data = request.json
    chat_room = data.get('chat_room')

    if not chat_room:
        return jsonify({"error": "Chat room name is required"}), 400

    # Add logic to create a chat room (if needed)
    return jsonify({"status": f"Chat room '{chat_room}' created"}), 200

@app.route('/get_messages', methods=['GET'])
def get_messages():
    """Retrieve messages for a specific user (not implemented here)."""
    # You can implement a message queue or buffer for each user to store messages
    return jsonify({"error": "Not implemented"}), 501

@app.route('/get_history', methods=['GET'])
def get_history():
    """Retrieve message history for a specific chat room."""
    chat_room = request.args.get('chat_room', 'default')
    if chat_room not in message_history:
        return jsonify({"error": f"No history found for chat room '{chat_room}'"}), 404

    return jsonify({"chat_room": chat_room, "messages": message_history[chat_room]}), 200

# WebSocket event handlers
@socketio.on('connect')
def handle_connect():
    print("WebSocket client connected")

@socketio.on('disconnect')
def handle_disconnect():
    print("WebSocket client disconnected")

@socketio.on('join')
def handle_join(data):
    print(f"Received join event: {data}")
    username = data.get('username')
    chat_room = data.get('chat_room', 'default')

    if not username:
        emit('error', {'message': 'Username is required'})
        return

    join_room(chat_room)
    emit('message', {'username': 'Server', 'message': f"{username} has joined the chat room '{chat_room}'."}, room=chat_room)

@socketio.on('leave')
def handle_leave(data):
    username = data.get('username')
    chat_room = data.get('chat_room', 'default')
    leave_room(chat_room)
    emit('message', {'message': f"{username} has left the chat room '{chat_room}'."}, room=chat_room)

@socketio.on('send_message')
def handle_send_message(data):
    print(f"Received send_message event: {data}")
    username = data.get('username')
    chat_room = data.get('chat_room', 'default')
    message = data.get('message')

    if not username or not message:
        emit('error', {'message': 'Username and message are required'})
        return

    # Use the centralized message handling function
    handle_message(username, chat_room, message)

# Chat server logic
def chat_server():
    global quitFlag, server_socket, socket_list, client_info, chat_rooms, client_states, uuid_dict

    # Create and configure the server socket
    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
    server_socket.bind(('', CHATPORT))
    server_socket.listen(MAX_CONNECT_QUEUE)
    server_socket.setblocking(False)

    socket_list.append(server_socket)
    print(f"Chat server listening on port {CHATPORT}")

    while quitFlag:
        read_sockets, _, _ = select.select(socket_list, [], [])
        for sock in read_sockets:
            if sock == server_socket:
                # Accept new client connections
                client_socket, client_address = server_socket.accept()
                client_socket.setblocking(False)
                socket_list.append(client_socket)
                print(f"Accepted connection from {client_address}")
                client_socket.send(b"Welcome to the chat server!\n")
                client_socket.send(b"Type in a username: \n")
                client_states[client_socket] = "WAITING_FOR_USERNAME"
                user_id = get_client_uuid(client_socket, uuid_dict)
                client_info[client_socket] = {"user_id": user_id}
            else:
                # Handle client messages
                try:
                    message = sock.recv(BUFFER_SIZE).decode('utf-8')
                    if message:
                        state = client_states.get(sock, "NORMAL")
                        if state == "WAITING_FOR_USERNAME":
                            username = message.strip()
                            if username == "admin":
                                sock.send(b"\nEnter admin password: \n")
                                client_states[sock] = "WAITING_FOR_ADMIN_PASSWORD"
                                client_info[sock]["username"] = username
                            else:
                                client_info[sock]["username"] = username
                                chat_rooms[sock] = "default"
                                sock.send(b"\nUsername accepted. You are now in the 'default' chatroom.\n")
                                client_states[sock] = "NORMAL"
                        elif state == "WAITING_FOR_ADMIN_PASSWORD":
                            password = message.strip()
                            if password == "adminpass":
                                sock.send(b"\nAdmin authenticated successfully.\n")
                                client_info[sock]["is_admin"] = True
                                chat_rooms[sock] = "default"
                                client_states[sock] = "NORMAL"
                            else:
                                sock.send(b"\nInvalid admin password. Please enter a username: \n")
                                client_states[sock] = "WAITING_FOR_USERNAME"
                        elif state == "WAITING_FOR_NEW_USERNAME":
                            client_info[sock]["username"] = message.strip()
                            sock.send(b"\nUsername changed successfully.\n")
                            print(f"\nClient switched username to {client_info[sock]['username']}.")
                            client_states[sock] = "NORMAL"
                        elif state == "WAITING_FOR_NEW_CHATROOM":
                            old_chatroom = chat_rooms[sock]
                            chat_rooms[sock] = message.strip()
                            sock.send(f"\nSwitched to chatroom '{chat_rooms[sock]}.'\n".encode('utf-8'))
                            print(f"\nClient switched to chatroom '{chat_rooms[sock]}.'")
                            if old_chatroom != "default" and old_chatroom not in chat_rooms.values():
                                print(f"\nChatroom '{old_chatroom}' is empty and removed.")
                            client_states[sock] = "NORMAL"
                        else:
                            if message.strip() == "%!":
                                sock.close()
                                socket_list.remove(sock)
                                del client_info[sock]
                                del chat_rooms[sock]
                                del client_states[sock]
                                print("\nClient has closed connection.")
                            elif message.strip() == "%?":
                                print("\nClient requested help.")
                                sock.send(b"Help: (%?)\n")
                                sock.send(b"Send messages like normal. Use enter to send a message.\n")
                                sock.send(b"Available Commands:\n")
                                sock.send(b"\t%? : Show this help prompt.\n")
                                sock.send(b"\t%! : Exit from the Chat Server.\n")
                                sock.send(b"\t%* : Switch to another chatroom.\n")
                                sock.send(b"\t%$ : Change Username.\n")
                                sock.send(b"\t%@ <username> <message> : Send a direct message to a user.\n")
                                if client_info[sock].get("is_admin"):
                                    sock.send(b"\t%stop : Stop the server.\n")
                            elif message.strip() == "%*":
                                sock.send(b"\nEnter chatroom name: \n")
                                client_states[sock] = "WAITING_FOR_NEW_CHATROOM"
                            elif message.strip() == "%$":
                                sock.send(b"\nEnter new username: \n")
                                client_states[sock] = "WAITING_FOR_NEW_USERNAME"
                            elif message.strip().startswith("%@ "):
                                parts = message.strip().split(' ', 2)
                                if len(parts) < 3:
                                    sock.send(b"\nInvalid direct message format. Use: %@ <username> <message>\n")
                                else:
                                    target_username = parts[1]
                                    dm_message = parts[2]
                                    target_sock = None
                                    for s, info in client_info.items():
                                        if info.get("username") == target_username:
                                            target_sock = s
                                            break
                                    if target_sock:
                                        sender_username = client_info[sock]["username"]
                                        target_sock.send(f"DM from {sender_username}: {dm_message}\n".encode('utf-8'))
                                        sock.send(b"\nDirect message sent.\n")
                                    else:
                                        sock.send(b"\nUser not found.\n")
                            elif message.strip() == "%stop" and client_info[sock].get("is_admin"):
                                print("\nAdmin has stopped the server.")
                                quitFlag = False
                            else:
                                username = client_info[sock]["username"]
                                chat_room = chat_rooms[sock]

                                # Use the centralized message handling function
                                handle_message(username, chat_room, message)
                    else:
                        sock.close()
                        socket_list.remove(sock)
                        if sock in client_info:
                            del client_info[sock]
                        if sock in chat_rooms:
                            old_chatroom = chat_rooms[sock]
                            del chat_rooms[sock]
                            if old_chatroom != "default" and old_chatroom not in chat_rooms.values():
                                print(f"\nChatroom '{old_chatroom}' is empty and has been removed.")
                        if sock in client_states:
                            del client_states[sock]
                        print("\nClient has closed connection.")
                except Exception as err:
                    print(f"Error: {err}")
                    sock.close()
                    socket_list.remove(sock)
                    if sock in client_info:
                        del client_info[sock]
                    if sock in chat_rooms:
                        old_chatroom = chat_rooms[sock]
                        del chat_rooms[sock]
                        if old_chatroom != "default" and old_chatroom not in chat_rooms.values():
                            print(f"\nChatroom '{old_chatroom}' is empty and has been removed.")
                    if sock in client_states:
                        del client_states[sock]
                    print("\nClient has closed connection.")

    # Save data before shutting down
    save_data(client_info, chat_rooms, client_states)
    server_socket.close()

if __name__ == "__main__":
    # Start the chat server in a separate thread
    server_thread = threading.Thread(target=chat_server, daemon=True)
    server_thread.start()

    # Start the Flask REST API and WebSocket server
    socketio.run(app, host='0.0.0.0', port=5000)