from datetime import datetime as dt, timezone as tz
from flask import request, redirect, url_for, render_template, Blueprint, session, current_app, jsonify
from flask_security import hash_password, roles_required, roles_accepted
from sqlalchemy.exc import SQLAlchemyError
from flask_jwt_extended import create_access_token, jwt_required, get_jwt_identity
from flask_socketio import emit, join_room, leave_room
import ast


from common.models.user import User, Tenant
from common.models.interaction import ChatSession, Interaction, InteractionEmbedding
from common.models.document import Embedding
from common.extensions import db, socketio, kms_client
from common.utils.database import Database

chat_bp = Blueprint('chat_bp', __name__, url_prefix='/chat')


@chat_bp.route('/register_client', methods=['POST'])
def register_client():
    tenant_id = request.json.get('tenant_id')
    api_key = request.json.get('api_key')

    # Validate tenant_id and api_key here (e.g., check against the database)
    if validate_tenant(tenant_id, api_key):
        access_token = create_access_token(identity={'tenant_id': tenant_id, 'api_key': api_key})
        current_app.logger.debug(f'Tenant Registration: Tenant {tenant_id} registered successfully')
        return jsonify({'token': access_token}), 200
    else:
        current_app.logger.debug(f'Tenant Registration: Invalid tenant_id ({tenant_id}) or api_key ({api_key})')
        return jsonify({'message': 'Invalid credentials'}), 401


@socketio.on('connect', namespace='/chat')
@jwt_required()
def handle_connect():
    current_tenant = get_jwt_identity()
    current_app.logger.debug(f'Tenant {current_tenant["tenant_id"]} connected')


@socketio.on('message', namespace='/chat')
@jwt_required()
def handle_message(data):
    current_tenant = get_jwt_identity()
    current_app.logger.debug(f'Tenant {current_tenant["tenant_id"]} sent a message: {data}')
    # Store interaction in the database
    emit('response', {'data': 'Message received'}, broadcast=True)


def validate_tenant(tenant_id, api_key):
    tenant = Tenant.query.get_or_404(tenant_id)
    encrypted_api_key = ast.literal_eval(tenant.encrypted_chat_api_key)

    decrypted_api_key = kms_client.decrypt_api_key(encrypted_api_key)

    return decrypted_api_key == api_key



# @chat_bp.route('/', methods=['GET', 'POST'])
# def chat():
#     return render_template('chat.html')
#
#
# @chat.record_once
# def on_register(state):
#     # TODO: write initialisation code when the blueprint is registered (only once)
#     # socketio.init_app(state.app)
#     pass
#
#
# @socketio.on('message', namespace='/chat')
# def handle_message(message):
#     # TODO: write message handling code to actually realise chat
#     # print('Received message:', message)
#     # socketio.emit('response', {'data': message}, namespace='/chat')
#     pass