from flask import Flask, request, jsonify
from flask_cors import CORS
import secrets
from ai import AI


class API:
    def __init__(self):
        self.app = Flask(__name__)
        self.ai_response = {}
        self.ai = AI()
        CORS(self.app)

    def run(self):
        @self.app.route('/interstellar/api/ai_create', methods=['GET'])
        def create_ai():
            access_token = secrets.token_urlsafe(4096)
            self.ai_response[access_token] = ""
            return jsonify({'status': 200, 'access_token': access_token})

        @self.app.route('/interstellar/api/ai_send', methods=['POST'])
        def send_ai():
            data = request.get_json()
            messages = data.get('messages')
            ai_model = data.get('ai_model')
            access_token = data.get('access_token')
            if access_token not in self.ai_response:
                return jsonify({'status': 401, 'error': 'Invalid access token'})
            self.ai.process_local(ai_model, messages, self, access_token)
            return jsonify({'status': 200})

        @self.app.route('/interstellar/api/ai_get', methods=['GET'])
        def get_ai():
            data = request.args.get('access_token')
            if data not in self.ai_response:
                return jsonify({'status': 401, 'error': 'Invalid access token'})
            return jsonify({'status': 200, 'response': self.ai_response[data]})

        ssl_context = ('cert.pem', 'key.pem')
        self.app.run(debug=True, host='0.0.0.0', port=5000, ssl_context=ssl_context)


if __name__ == '__main__':
    api = API()
    api.run()