from flask import Flask, render_template, request, jsonify, make_response
from modules.model import summarize
from flask_cors import CORS, cross_origin

import __main__

app = Flask(__name__)
CORS(app, supports_credentials=True)
app.config['CORS_ALLOW_HEADERS'] = 'Content-Type'
app.config['CORS_SUPPORTS_CREDENTIALS'] = 'True'
# shortTokenizer = BartTokenizer.from_pretrained('sshleifer/distilbart-xsum-12-6')
# shortModel = BartForConditionalGeneration.from_pretrained('sshleifer/distilbart-xsum-12-6')

# longTokenizer = BartTokenizer.from_pretrained('sshleifer/distilbart-cnn-12-6')
# longModel = BartForConditionalGeneration.from_pretrained('sshleifer/distilbart-cnn-12-6')

@app.route('/')
def home():
    return render_template('index.html')


@app.route('/summarize', methods=['POST', 'OPTIONS'])
@cross_origin()
def recommend():
    if request.method == "POST":
        # Get form data
        request_data = request.get_json()
        input_text = request_data['input_text']

        # Call the function summarize to run the text summarization
        try:
            short_output_summary, long_output_summary = summarize(input_text)
            response = jsonify({'short': short_output_summary.strip(), 'long': long_output_summary.strip()})
            # Pass output summary to the output template
            response.headers.add('Access-Control-Allow-Origin', 'http://0.0.0.0:7860')
            response.headers.add('Access-Control-Allow-Headers', 'Content-Type, Authorization')
            response.headers.add('Access-Control-Allow-Methods', 'GET,PUT,POST,DELETE,OPTIONS')
            response.headers.add('Access-Control-Allow-Credentials', 'true')
            return response

        except Exception as e:
            return render_template('index.html', query=e)

    elif request.method == "OPTIONS":
        response = make_response("OK")
        response.status_code = 201
        response.headers.add('Access-Control-Allow-Origin', 'http://0.0.0.0:7860')
        response.headers.add('Access-Control-Allow-Headers', 'Content-Type, Authorization')
        response.headers.add('Access-Control-Allow-Methods', 'GET,PUT,POST,DELETE,OPTIONS')
        response.headers.add('Access-Control-Allow-Credentials', 'true')
        # Debug
        print("Right")
        return response

    pass



def main():
    # app.config['TEMPLATES_AUTO_RELOAD'] = True
    # app.config['SEND_FILE_MAX_AGE_DEFAULT'] = 0
    # app.run(debug=True)
    app.run(host='0.0.0.0', port=7860)

if __name__ == '__main__':
    print("Loading BART model and tokenzier . . .")
    main()