from flask import Flask, request, jsonify
from flask_cors import CORS
import torch
import os
from transformers import BertTokenizer, BertModel, AutoModelForCausalLM, AutoTokenizer

# ✅ Initialize Flask app
app = Flask(__name__)
CORS(app)  # Enable CORS for cross-origin requests

# ✅ Detect GPU
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

# ✅ Load Tokenizer for Yelp Model
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")

# ✅ Load Models (if available)
models = {}

def load_model(model_name, class_definition):
    """Load model if it exists, otherwise return None."""
    model_path = f"models/{model_name}.pth"
    if os.path.exists(model_path):
        model = class_definition().to(device)
        model.load_state_dict(torch.load(model_path, map_location=device))
        model.eval()
        return model
    else:
        print(f"⚠️ {model_name} model not found!")
        return None

# ✅ Define Star Rating Model
class StarRatingModel(torch.nn.Module):
    def __init__(self):
        super(StarRatingModel, self).__init__()
        self.bert = BertModel.from_pretrained("bert-base-uncased")
        self.fc_stars = torch.nn.Linear(768, 5)

    def forward(self, input_ids, attention_mask):
        text_output = self.bert(input_ids=input_ids, attention_mask=attention_mask).pooler_output
        return self.fc_stars(text_output)

models["star_model"] = load_model("star_rating_model", StarRatingModel)

# ✅ Define Usefulness Model
class UsefulnessModel(torch.nn.Module):
    def __init__(self):
        super(UsefulnessModel, self).__init__()
        self.bert = BertModel.from_pretrained("bert-base-uncased")
        self.fc_useful = torch.nn.Linear(768, 1)

    def forward(self, input_ids, attention_mask):
        text_output = self.bert(input_ids=input_ids, attention_mask=attention_mask).pooler_output
        return torch.clamp(torch.round(self.fc_useful(text_output)), 1, 5)

models["usefulness_model"] = load_model("usefulness_model", UsefulnessModel)

# ✅ Load Chatbot Response Model
chatbot_model_name = "facebook/opt-1.3b"
if os.path.exists("models/response_model.pth"):
    chatbot_tokenizer = AutoTokenizer.from_pretrained(chatbot_model_name)
    chatbot_model = AutoModelForCausalLM.from_pretrained(chatbot_model_name).to(device)
    chatbot_model.load_state_dict(torch.load("models/response_model.pth", map_location=device))
    chatbot_model.eval()
    models["response_model"] = chatbot_model
else:
    print("⚠️ Response model not found!")

@app.route("/")
def home():
    return jsonify({"message": "Welcome to Yelp Review AI Predictor!"})

@app.route("/predict", methods=["POST"])
def predict_rating():
    try:
        # ✅ Get JSON data from request
        data = request.get_json()

        if not data:
            return jsonify({"error": "No input data provided"}), 400

        text = data.get("text", "").strip()

        # ✅ Ensure text is provided
        if not text:
            return jsonify({"error": "Missing 'text' field"}), 400

        # ✅ Tokenize input
        tokens = tokenizer(
            text, truncation=True, padding="max_length", max_length=256, return_tensors="pt"
        )
        input_ids = tokens["input_ids"].to(device)
        attention_mask = tokens["attention_mask"].to(device)

        # ✅ Make predictions (if models exist)
        results = {}

        # ⭐ Predict Stars
        if models["star_model"]:
            with torch.no_grad():
                stars_logits = models["star_model"](input_ids, attention_mask)
                predicted_stars = torch.argmax(stars_logits, dim=1).item() + 1  # Convert to 1-5
            results["predicted_stars"] = predicted_stars
        else:
            results["predicted_stars"] = "AI has not been trained and created"

        # 📈 Predict Usefulness
        if models["usefulness_model"]:
            with torch.no_grad():
                predicted_usefulness = models["usefulness_model"](input_ids, attention_mask).item()
            results["predicted_usefulness"] = int(predicted_usefulness)  # Convert to integer (1-5)
        else:
            results["predicted_usefulness"] = "AI has not been trained and created"

        # 📝 Generate AI Response (if usefulness > 1)
        if "response_model" in models and results["predicted_usefulness"] != "AI has not been trained and created":
            if results["predicted_usefulness"] > 1:
                response_input = chatbot_tokenizer(text, return_tensors="pt").to(device)
                response_output = models["response_model"].generate(**response_input, max_length=50)
                ai_response = chatbot_tokenizer.decode(response_output[0], skip_special_tokens=True)
            else:
                ai_response = "No response (low usefulness)."
            results["ai_response"] = ai_response
        else:
            results["ai_response"] = "AI has not been trained and created"

        return jsonify(results)

    except Exception as e:
        return jsonify({"error": str(e)}), 500

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=8000, debug=True)  # Listen on all interfaces