from flask import Flask, request, session, render_template
from flask_sqlalchemy import SQLAlchemy
from flask_cors import CORS
from datetime import timedelta
import json

db = SQLAlchemy()
app = Flask(__name__)
app.secret_key = "UAZFCNfzfdGFNfhjtLIFGEbjytdazdwfQFYDS"
app.config['SQLALCHEMY_DATABASE_URI'] = 'mysql+pymysql://project6110user:inse6110DB@localhost:3306/project_6110'
# Second DB
app.config['SQLALCHEMY_BINDS'] = {
    'keys_db': 'mysql+pymysql://project6110user:inse6110DB@192.168.100.161:3306/project_6110'
}
app.config['SESSION_COOKIE_SAMESITE'] = 'None'
app.config['SESSION_COOKIE_SECURE'] = False
app.permanent_session_lifetime = timedelta(days=7) # Session lasts for 7 days

CORS(app, origins=["*"], supports_credentials=True)

# Default DB (db1)

from models import *

db.init_app(app)

# AUTHENTICATION

@app.route('/login', methods=['POST'])
def login():

    username = request.form.get('username')
    hashed_password = request.form.get('hashed_password')

    if not username or not hashed_password:
        return "Missing username or password\n", 400

    user = User.query.filter_by(user_name=username, user_password=hashed_password).first()

    if not user:
        return "Invalid username or password\n", 403

    session['user_id'] = user.user_id
    session['username'] = user.user_name
 

    return f"User {username} logged in successfully!\n", 200


@app.route('/logout', methods=['POST'])
def Logout():
    session.pop('user_id', None)
    session.pop('username', None)
    session.clear()
    return "User logged out successfully!\n", 200

# HELPER FUNCTIONS
# will clean dead or broken chats and users

def is_logged_in():
    # return session['user_id'], session['username']
    return 26, "test1" # for testing purposes only

def cleanup_db():
    # check if there are users with no private or public key and delete them from both databases
    users = User.query.all()

    for user in users:

        user_keys = UserKey.query.filter_by(user_id=user.user_id).first()

        if not user_keys:
            db.session.delete(user)

    db.session.commit()

    # check if there are chats with no members or no keys and delete them from both db
    chats = Chat.query.all()

    for chat in chats:

        chat_members = chat_members = ChatMember.query.filter_by(chat_id=chat.chat_id).all()
        chat_keys = ChatKey.query.filter_by(chat_id=chat.chat_id).all()
        
        if not chat_members or not chat_keys:
            db.session.delete(chat)

    db.session.commit()

def add_Chat_Members(chat_id: int, user_ids: list, owner = None): # used in init_chat
    print("test", chat_id, user_ids[0], owner)
    for user_id in user_ids:
        new_member = ChatMember (
            chat_id=chat_id,
            user_id=user_id,
            user_role = "owner" if owner else "member" # if the owner parameter is set the user gets added as the owner
        )
        db.session.add(new_member)
    db.session.commit()




# ACTUAL APP



# WROK PERFECTLY
@app.route("/register", methods=['POST'])
def add_user():
    # get all the arguments needed to create a user

    username = request.form.get('username')
    userpassword = request.form.get('password')
    userdisplayname = request.form.get('displayname') if request.form.get('displayname') else username
    userage = request.form.get('age') if request.form.get('age') else 0
    protected_enc_privatekey = request.form.get('enc_privatekey')
    protected_sign_privatekey = request.form.get('sign_privatekey')
    enc_publickey = request.form.get('public_key_enc')
    sign_publickey = request.form.get('public_key_sign')

    if not username or not userpassword:
        return "Missing User's fields\n", 400
    
    if not protected_enc_privatekey or not enc_publickey or not protected_sign_privatekey or not sign_publickey:
        return "Missing User Keys user cannot be created\n", 400
    
    print(protected_enc_privatekey)

    # protected_enc_privatekey = json.loads(protected_enc_privatekey)
    # protected_sign_privatekey = json.loads(protected_sign_privatekey)

    try:
        # adding user to the main user database
        new_user = User(
            user_name=str(username),
            user_password=str(userpassword),
            user_displayname=str(userdisplayname),
            user_age=int(userage)
        )

        db.session.add(new_user)
        db.session.commit()

        new_user_id = User.query.filter_by(user_name=str(username)).first()

        # adding the user's minimum info and keys to the keys database
        new_key = UserKey(
            user_name=str(username),
            user_id=int(new_user_id.user_id),
            protected_enc_privatekey=str(protected_enc_privatekey),
            enc_publickey=str(enc_publickey),
            protected_sign_privatekey=str(protected_sign_privatekey),
            sign_publickey=str(sign_publickey),
        )

        db.session.add(new_key)
        db.session.commit()

    except Exception as e:
        return f"Error while adding user to database: {str(e)}\n", 500
    
    return f"User {username} added successfully!\n", 200

# works and was tested without actual app and auth
# starts a new chat with one or multiple users
@app.route('/init_chat', methods=['POST'])
def init_chat():

    #check for authentication
    user_id, username = is_logged_in()
    if not user_id: 
        return "User not logged in\n", 403

    target_userID = request.form.get('target_userID')

    if not target_userID:
        return "Missing target user for chat init\n", 400

    target_userID = target_userID.split(',') if ',' in target_userID else int(target_userID) # if there are multiple users the target_userID is a list of int otherwise it is a single int

    # if the target is an int (only one value) then is_group is false if it is an array is_group is True
    if target_userID.__class__ == list: is_group = True
    else: is_group = False

    chat_name = request.form.get('chat_name') if request.form.get('chat_name') else "New Chat Group" #replace by the concat of at least 2 users

    try:
        new_chat = Chat(
            chat_name=chat_name,
            is_group=is_group
        )
        db.session.add(new_chat)
        db.session.commit()

        new_chat_id = Chat.query.order_by(Chat.chat_id.desc()).first()

        if not new_chat_id:
            return "Error while retrieving new chat's ID\n", 500

        # add the creator of the chat as a member
        add_Chat_Members(new_chat_id.chat_id, [user_id], owner = True)

        # add the target user(s) as member(s)
        if is_group:
            add_Chat_Members(new_chat_id.chat_id, target_userID)
        else:
            add_Chat_Members(new_chat_id.chat_id, [target_userID])


    except Exception as e:
        return f"Error while creating chat: {str(e)}\n", 500

    return f"Chat created successfully!\n", 200

@app.route('/send_message', methods=['POST'])
def send_message():
    user_id, username = is_logged_in()
    if not user_id:
        return "User not logged in\n", 403

    enc_message = request.form.get('enc_message')
    signature = request.form.get('signature')
    chat_id = request.form.get('chat_id')
    recipients_json = request.form.get('recipients')  # JSON string of [{user_id, encrypted_key}, ...]

    if not enc_message or not signature or not chat_id or not recipients_json:
        return "Missing fields to send message\n", 400

    chat_member_check = ChatMember.query.filter_by(chat_id=chat_id, user_id=user_id).first()
    if not chat_member_check:
        return "User not in chat\n", 403

    try:
        # Parse recipients
        recipients = json.loads(recipients_json)
        if not recipients or not isinstance(recipients, list):
            return "Invalid recipients format\n", 400

        # 1. Create ONE Message row
        new_message = Message(
            chat_id=int(chat_id),
            user_id=user_id,
            key_id=0,
            enc_content=str(enc_message),
            signature=signature
        )
        db.session.add(new_message)
        db.session.flush()  # Get message_id

        # Use message_id as key_id
        key_id = new_message.message_id
        new_message.key_id = key_id

        # 2. Create ChatKey row for EACH recipient
        for recipient in recipients:
            recipient_user_id = int(recipient['user_id'])
            encrypted_key = recipient['encrypted_key']

            recipient_user = User.query.filter_by(user_id=recipient_user_id).first()
            if not recipient_user:
                db.session.rollback()
                return f"Recipient {recipient_user_id} does not exist\n", 400

            new_key = ChatKey(
                user_name=recipient_user.user_name,
                user_id=recipient_user_id,
                chat_id=int(chat_id),
                key_id=int(key_id),
                encrypted_symetrickey=str(encrypted_key)
            )
            db.session.add(new_key)

        # 3. Commit everything at once
        db.session.commit()

        return {"status": "ok", "message_id": key_id}, 200

    except json.JSONDecodeError:
        db.session.rollback()
        return "Invalid recipients JSON\n", 400
    except Exception as e:
        db.session.rollback()
        return f"Error while sending message: {str(e)}\n", 500

# works and was tested without actual app and auth
# get all the chats info of a user (exept messages and symetric key)
@app.route('/get_chats', methods=['POST'])
def get_chats():

    #check for authentication
    user_id, username = is_logged_in()
    if not user_id: 
        return "User not logged in\n", 403
    

    try: 
        chat_memberships = ChatMember.query.filter_by(user_id=user_id).all()
        chats_info = []
        for membership in chat_memberships:
            chat = Chat.query.filter_by(chat_id=membership.chat_id).first()
            if chat:
                chats_info.append({
                    'chat_id': chat.chat_id,
                    'chat_name': chat.chat_name,
                    'is_group': chat.is_group,
                    'creation_date': chat.creation_date,
                    'user_role': membership.user_role
                })
    
    except Exception as e:
        return f"Error while retrieving chats: {str(e)}\n", 500
    
    return {"chats": chats_info}, 200

# works and was tested without actual app and auth
# get messages from a chat and the user's encrypted symetric key for the chat 
@app.route('/get_messages', methods=['POST'])
def get_messages():
    # check for authentication
    user_id, username = is_logged_in()
    if not user_id:
        return "User not logged in\n", 403

    chat_id = request.form.get('chat_id')
    limit = int(request.form.get('limit')) if request.form.get('limit') else 50
    offset = int(request.form.get('offset')) if request.form.get('offset') else 0

    if not chat_id:
        return "Missing chat_id to get messages from\n", 400

    # ensure user is member of chat
    chat_member_check = ChatMember.query.filter_by(chat_id=chat_id, user_id=user_id).first()
    if not chat_member_check:
        return "User not in chat\n", 403

    try:
        # get messages for this chat
        messages = (
            Message.query
            .filter_by(chat_id=chat_id, is_deleted=False)
            .order_by(Message.sent_at)
            .limit(limit)
            .offset(offset)
            .all()
        )

        messages_info = []

        for message in messages:
            # for EACH message, find the ChatKey row for THIS user and this key_id
            key = ChatKey.query.filter_by(
                chat_id=chat_id,
                key_id=message.key_id,
                user_id=user_id
            ).first()

            if not key:
                # no key for this user => they cannot decrypt this message
                continue

            sender_keys = UserKey.query.filter_by(user_id=message.user_id).first()
            sign_publickey = sender_keys.sign_publickey if sender_keys else None

            messages_info.append({
                'message_id': message.message_id,
                'user_id': message.user_id,
                'key_id': key.key_id,
                'enc_content': message.enc_content,
                'sent_at': message.sent_at,
                'signature': message.signature,
                'encrypted_symetrickey': key.encrypted_symetrickey,
                'sign_publickey': sign_publickey
            })

    except Exception as e:
        return f"Error while retrieving messages: {str(e)}\n", 500

    return {"messages": messages_info}, 200


# works and was tested without actual app and auth
# retrives all the user's keys (private, public and symetric chat keys)
@app.route('/get_mykeys', methods=['POST'])
def get_mykeys():
    user_id, username = is_logged_in()
    if not user_id:
        return "User not logged in\n", 403

    user_exist = User.query.filter_by(user_id=user_id).first()
    if not user_exist:
        return "User does not exist\n", 400

    try:
        user_keys = UserKey.query.filter_by(user_id=user_id).first()
        if not user_keys:
            return "No keys found for this user\n", 500
    except Exception as e:
        return f"Error while retrieving the keys: {str(e)}\n", 500

    # These are used on *every login* to rebuild keys client-side
    return {
    "protected_enc_privatekey": user_keys.protected_enc_privatekey,
    "enc_publickey": user_keys.enc_publickey,
    "protected_sign_privatekey": user_keys.protected_sign_privatekey,
    "sign_publickey": user_keys.sign_publickey,
}, 200


# this is used to get the publickeys of each member of a chat, it will also return the last key id of said chat.
@app.route('/get_chat_publickeys', methods=['POST'])
def get_chat_publickeys():
    user_id, username = is_logged_in()
    if not user_id:
        return "User not logged in\n", 403

    chat_id = request.form.get('chat_id')

    chat_exist = Chat.query.filter_by(chat_id=chat_id).first()
    if not chat_exist:
        return "Chat does not exist\n", 400

    is_member = ChatMember.query.filter_by(chat_id=chat_id, user_id=user_id).first()
    if not is_member:
        return "User not in chat\n", 403

    chat_members = ChatMember.query.filter_by(chat_id=chat_id).all()

    try:
        users_keys_list = []

        # Include ALL chat members (including current user so they can decrypt their own messages)
        for member in chat_members:
            user_keys = UserKey.query.filter_by(user_id=member.user_id).first()
            if not user_keys:
                continue

            users_keys_list.append({
                "user_id": member.user_id,
                "enc_publickey": user_keys.enc_publickey,
                "sign_publickey": user_keys.sign_publickey
            })

        if not users_keys_list:
            return "No keys found for chat members\n", 404

    except Exception as e:
        return f"Error while retrieving the keys: {str(e)}\n", 500

    return {"users_keys": users_keys_list}, 200



# works and was tested without actual app and auth
# retrieves all the users except the requester
@app.route('/get_users', methods=['POST'])
def get_all_users():

    #check for authentication
    user_id, username = is_logged_in()
    if not user_id: 
        return "User not logged in\n", 403

    try:
        users = User.query.all()
        users_info = []

        for user in users:

            # add the entry to users_info only if the user is not the requester and add the user's public key

            if user.user_id == user_id:
                continue

            user_pub_keys = UserKey.query.filter_by(user_id=user.user_id).first()

            users_info.append({
                'user_id': user.user_id,
                'user_name': user.user_name,
                'user_displayname': user.user_displayname,
                'user_age': user.user_age,
                'public_key': user_pub_keys.enc_publickey if user_pub_keys else None
            })

        return {"users": users_info}, 200

    except Exception as e:
        return f"Error while retrieving users: {str(e)}\n", 500

if __name__ == '__main__':
    app.run(debug=True, port=8000, host='0.0.0.0')
