from news_extractor import get_news
from db_operations.db_write import DBWrite
from db_operations.db_read import DBRead
from news_category_prediction import predict_news_category
import json
from flask import Flask, Response
from flask_cors import cross_origin, CORS
import logging
import tensorflow as tf 
import cloudpickle
from transformers import DistilBertTokenizerFast
import os

app = Flask(__name__)
CORS(app)
logging.warning('Initiated')


def load_model():
    interpreter = tf.lite.Interpreter(model_path=os.path.join("models/news_classification_hf_distilbert.tflite"))
    with open("models/news_classification_labelencoder.bin", "rb") as model_file_obj:
        label_encoder = cloudpickle.load(model_file_obj)
        
    model_checkpoint = "distilbert-base-uncased"
    tokenizer = DistilBertTokenizerFast.from_pretrained(model_checkpoint)
    return interpreter, label_encoder, tokenizer

interpreter, label_encoder, tokenizer = load_model()


@app.route("/")
@cross_origin()
def update_news():
    status_json = "{'status':'success'}"
    status_code = 200
    try:
        db_read = DBRead()
        db_write = DBWrite()
        old_news = db_read.read_news_from_db()
        new_news = get_news()
        news_df = predict_news_category(old_news, new_news, interpreter, label_encoder, tokenizer)
        news_json = [*json.loads(news_df.reset_index(drop=True).to_json(orient="index")).values()]
        db_write.insert_news_into_db(news_json)
    except:
        status_json = "{'status':'failure'}"
        status_code = 500
    return Response(status_json, status=status_code, mimetype='application/json')


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860, timeout=120, workers=1, threads=1)