routes/auth.py

from flask import Blueprint, request, jsonify, session, render_template, redirect, url_for
from bson import ObjectId
from datetime import datetime
import logging
import re

from utils.extensions import ext
from utils.auth import hash_password, verify_password, validate_phone, validate_email
from utils.email import generate_otp, generate_random_password, send_email
from utils.query import can_make_query, update_query_count

auth_bp = Blueprint('auth', __name__)

@auth_bp.route('/register', methods=['GET', 'POST'])
def register():
    if request.method == 'GET':
        return render_template('register.html')
    
    data = request.get_json(silent=True) or {}
    username = data.get('username', '').strip()
    email = data.get('email', '').strip()
    password = data.get('password', '').strip()
    phone = data.get('phone', '').strip()
    account_type = data.get('account_type', 'limited').strip()
    
    if not all([username, email, password, phone]):
        return jsonify({'error': 'Thiếu thông tin bắt buộc'}), 400
    
    if not validate_phone(phone):
        return jsonify({'error': 'Số điện thoại không hợp lệ. Định dạng: +84xxxxxxxxx hoặc 0xxxxxxxxx'}), 400
    
    if not validate_email(email):
        return jsonify({'error': 'Email không hợp lệ'}), 400
    
    if account_type not in ['limited', 'unlimited']:
        return jsonify({'error': 'Loại tài khoản không hợp lệ'}), 400
    
    try:
        if ext.db.users.find_one({'$or': [{'email': email}, {'phone': phone}]}):
            return jsonify({'error': 'Email hoặc số điện thoại đã tồn tại'}), 400
        
        otp = generate_otp()
        password_hash = hash_password(password)
        
        user = {
            'username': username,
            'email': email,
            'phone': phone,
            'password_hash': password_hash,
            'otp': otp,
            'is_active': False,
            'created_at': datetime.utcnow(),
            'is_admin': False,
            'account_type': account_type,
            'query_limit': 3 if account_type == 'limited' else None,
            'query_count': 0,
            'last_reset': datetime.utcnow()
        }
        
        result = ext.db.users.insert_one(user)
        
        if send_email(
            email,
            'Mã OTP xác thực tài khoản',
            f'Mã OTP của bạn là: {otp}. Vui lòng sử dụng mã này để xác thực tài khoản.'
        ):
            return jsonify({
                'message': 'Đăng ký thành công, vui lòng kiểm tra email để lấy mã OTP',
                'user_id': str(result.inserted_id)
            }), 201
        else:
            ext.db.users.delete_one({'_id': result.inserted_id})
            return jsonify({'error': 'Lỗi khi gửi OTP, vui lòng thử lại'}), 500
    except Exception as e:
        logging.error(f"Error in register: {e}")
        return jsonify({'error': 'Lỗi hệ thống'}), 500

@auth_bp.route('/verify_otp', methods=['GET', 'POST'])
def verify_otp():
    if request.method == 'GET':
        user_id = request.args.get('user_id')
        if not user_id:
            return jsonify({'error': 'Thiếu user_id'}), 400
        try:
            user = ext.db.users.find_one({'_id': ObjectId(user_id)})
            if not user:
                return jsonify({'error': 'Người dùng không tồn tại'}), 404
            return render_template('verify_otp.html', user_id=user_id)
        except Exception as e:
            logging.error(f"Invalid user_id: {e}")
            return jsonify({'error': 'user_id không hợp lệ'}), 400
    
    elif request.method == 'POST':
        data = request.get_json(silent=True) or {}
        user_id = data.get('user_id', '').strip()
        otp = data.get('otp', '').strip()
        
        if not user_id or not otp:
            return jsonify({'error': 'Thiếu user_id hoặc OTP'}), 400
        
        try:
            user = ext.db.users.find_one({'_id': ObjectId(user_id)})
            if not user:
                return jsonify({'error': 'Người dùng không tồn tại'}), 404
            
            if user.get('otp') != otp:
                return jsonify({'error': 'Mã OTP không đúng'}), 400
            
            ext.db.users.update_one(
                {'_id': ObjectId(user_id)},
                {'$set': {'is_active': True, 'otp': None}}
            )
            
            return jsonify({'message': 'Xác thực tài khoản thành công'}), 200
        except Exception as e:
            logging.error(f"Error verifying OTP: {e}")
            return jsonify({'error': 'Lỗi hệ thống, vui lòng thử lại'}), 500

@auth_bp.route('/get_masked_phone', methods=['POST'])
def get_masked_phone():
    data = request.get_json(silent=True) or {}
    email = data.get('email', '').strip()
    
    if not email:
        return jsonify({'error': 'Thiếu email'}), 400
    
    try:
        user = ext.db.users.find_one({'email': email})
        if not user:
            return jsonify({'error': 'Email không tồn tại'}), 404
        
        phone = user.get('phone', '')
        if len(phone) <= 4:
            masked_phone = '****'
        else:
            masked_phone = phone[:-4] + '****'
        
        return jsonify({'masked_phone': masked_phone}), 200
    except Exception as e:
        logging.error(f"Error getting masked phone: {e}")
        return jsonify({'error': 'Lỗi hệ thống'}), 500

@auth_bp.route('/forgot_password', methods=['GET', 'POST'])
def forgot_password():
    if request.method == 'GET':
        return render_template('forgot_password.html')
    
    elif request.method == 'POST':
        data = request.get_json(silent=True) or {}
        email = data.get('email', '').strip()
        last_four_digits = data.get('last_four_digits', '').strip()
        
        if not email or not last_four_digits:
            return jsonify({'error': 'Thiếu email hoặc 4 số cuối của số điện thoại'}), 400
        
        try:
            user = ext.db.users.find_one({'email': email})
            if not user:
                return jsonify({'error': 'Email không tồn tại'}), 404
            
            phone = user.get('phone', '')
            if len(phone) < 4 or not phone[-4:] == last_four_digits:
                return jsonify({'error': '4 số cuối của số điện thoại không khớp'}), 400
            
            new_password = generate_random_password(12)
            new_password_hash = hash_password(new_password)
            
            ext.db.users.update_one(
                {'_id': user['_id']},
                {'$set': {'password_hash': new_password_hash}}
            )
            
            if send_email(
                email,
                'Mật khẩu mới',
                f'Mật khẩu mới của bạn là: {new_password}. Vui lòng đổi mật khẩu sau khi đăng nhập.'
            ):
                return jsonify({'message': 'Mật khẩu mới đã được gửi qua email'}), 200
            else:
                return jsonify({'error': 'Lỗi khi gửi mật khẩu mới'}), 500
        except Exception as e:
            logging.error(f"Error in forgot_password: {e}")
            return jsonify({'error': 'Lỗi hệ thống'}), 500

@auth_bp.route('/change_password', methods=['GET'])
def change_password_get():
    if 'user_id' not in session:
        return redirect(url_for('auth.login_page'))
    return render_template('change_password.html')

@auth_bp.route('/change_password', methods=['POST'])
def change_password():
    if 'user_id' not in session:
        return jsonify({'error': 'Vui lòng đăng nhập để đổi mật khẩu'}), 401
    
    data = request.get_json(silent=True) or {}
    current_password = data.get('current_password', '').strip()
    new_password = data.get('new_password', '').strip()
    
    if not current_password or not new_password:
        return jsonify({'error': 'Thiếu mật khẩu hiện tại hoặc mật khẩu mới'}), 400
    
    user_id = session['user_id']
    try:
        user = ext.db.users.find_one({'_id': ObjectId(user_id)})
        if not user:
            return jsonify({'error': 'Người dùng không tồn tại'}), 404
        
        if not verify_password(user['password_hash'], current_password):
            return jsonify({'error': 'Mật khẩu hiện tại không đúng'}), 401
        
        new_password_hash = hash_password(new_password)
        ext.db.users.update_one(
            {'_id': ObjectId(user_id)},
            {'$set': {'password_hash': new_password_hash}}
        )
        
        # Send confirmation email
        if send_email(
            user['email'],
            'Xác nhận đổi mật khẩu',
            f'Mật khẩu của bạn đã được thay đổi thành công vào lúc {datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S")}.'
        ):
            return jsonify({'message': 'Đổi mật khẩu thành công, email xác nhận đã được gửi'}), 200
        else:
            return jsonify({'error': 'Đổi mật khẩu thành công nhưng lỗi khi gửi email xác nhận'}), 200
    except Exception as e:
        logging.error(f"Error in change_password: {e}")
        return jsonify({'error': 'Lỗi hệ thống'}), 500

@auth_bp.route('/logins', methods=['POST'])
def login():
    data = request.get_json(silent=True) or {}
    email = data.get('email', '').strip()
    password = data.get('password', '').strip()
    
    try:
        user = ext.db.users.find_one({'email': email})
        if not user:
            logging.error(f"No user found for email: {email}")
            return jsonify({'error': 'Email hoặc mật khẩu không đúng'}), 401
        
        if not user.get('is_active', False):
            return jsonify({
                'error': 'Tài khoản chưa được kích hoạt. Vui lòng xác thực OTP.',
                'id_user': str(user['_id'])
            }), 401
        
        if not verify_password(user['password_hash'], password):
            return jsonify({'error': 'Email hoặc mật khẩu không đúng'}), 401
        
        session['user_id'] = str(user['_id'])
        session['username'] = user['username']
        session['is_admin'] = user.get('is_admin', False)
        session['query_limit'] = user.get('query_limit', 3)
        session['query_count'] = user.get('query_count', 0)
        session['account_type'] = user.get('account_type', 'limited')
        
        # Broadcast initial query count and limit
        if str(user['_id']) in ext.connected_clients:
            ext.socketio.emit('query_update', {
                'query_count': user.get('query_count', 0),
                'query_limit': user.get('query_limit', 3 if user.get('account_type') == 'limited' else None)
            }, room=ext.connected_clients[str(user['_id'])])
            logging.info(f"Broadcasted initial query update to user {user['_id']}")

        return jsonify({
            'message': 'Đăng nhập thành công',
            'username': user['username'],
            'is_admin': user.get('is_admin', False),
            'account_type': user.get('account_type', 'limited'),
            'query_limit': user.get('query_limit', 3),
            'query_count': user.get('query_count', 0),
        }), 200
    except Exception as e:
        logging.error(f"Error in login: {e}")
        return jsonify({'error': 'Lỗi hệ thống'}), 500

@auth_bp.route('/logout', methods=['POST'])
def logout():
    user_id = session.get('user_id')
    if user_id in ext.connected_clients:
        del ext.connected_clients[user_id]
        logging.info(f"User {user_id} removed from connected clients on logout")
    
    session.pop('user_id', None)
    session.pop('username', None)
    session.pop('is_admin', None)
    session.pop('account_type', None)
    session.pop('query_limit', None)
    session.pop('query_count', None)
    
    return jsonify({'message': 'Đăng xuất thành công'}), 200

@auth_bp.route('/check_session', methods=['GET'])
def check_session():
    if 'user_id' in session:
        try:
            user = ext.db.users.find_one({'_id': ObjectId(session['user_id'])})
            if user:
                session['query_limit'] = user.get('query_limit', 3)
                session['query_count'] = user.get('query_count', 0)
                session['account_type'] = user.get('account_type', 'limited')
                
                if session['user_id'] in ext.connected_clients:
                    ext.socketio.emit('query_update', {
                        'query_count': user.get('query_count', 0),
                        'query_limit': user.get('query_limit', 3 if user.get('account_type') == 'limited' else None)
                    }, room=ext.connected_clients[session['user_id']])
                
                return jsonify({
                    'logged_in': True,
                    'username': session['username'],
                    'query_limit': session['query_limit'],
                    'query_count': session['query_count'],
                    'is_admin': session.get('is_admin', False),
                    'account_type': session['account_type']
                }), 200
        except Exception as e:
            logging.error(f"Error in check_session: {e}")
    
    return jsonify({'logged_in': False}), 200

@auth_bp.route('/resend_otp', methods=['POST'])
def resend_otp():
    data = request.get_json(silent=True) or {}
    user_id = data.get('user_id', '').strip()
    
    if not user_id:
        return jsonify({'error': 'Thiếu user_id'}), 400
    
    try:
        user = ext.db.users.find_one({'_id': ObjectId(user_id)})
        if not user:
            return jsonify({'error': 'Người dùng không tồn tại'}), 404
        
        if user.get('is_active', False):
            return jsonify({'error': 'Tài khoản đã được kích hoạt'}), 400
        
        new_otp = generate_otp()
        ext.db.users.update_one(
            {'_id': ObjectId(user_id)},
            {'$set': {'otp': new_otp}}
        )
        
        if send_email(
            user['email'],
            'Mã OTP mới',
            f'Mã OTP mới của bạn là: {new_otp}. Vui lòng sử dụng mã này để xác thực tài khoản.'
        ):
            return jsonify({'message': 'Mã OTP mới đã được gửi qua email'}), 200
        else:
            return jsonify({'error': 'Lỗi khi gửi OTP'}), 500
    except Exception as e:
        logging.error(f"Error resending OTP: {e}")
        return jsonify({'error': 'Lỗi hệ thống'}), 500