from flask import Flask, request, jsonify
import torch
import os
from transformers import DistilBertTokenizerFast, DistilBertForSequenceClassification

app = Flask(__name__)

# Hugging Face repo and token
REPO = "aaronmrls/distilBERT-maintenance-priority-scorer"
HF_TOKEN = os.environ.get("HUGGINGFACE_TOKEN")  # optional (only needed if repo is private)

print("Loading tokenizer and model from Hugging Face Hub...")
print("Using Hugging Face token:", "Yes" if HF_TOKEN else "No")

# Load tokenizer & model directly from Hub
tokenizer = DistilBertTokenizerFast.from_pretrained(REPO, use_auth_token=HF_TOKEN)
model = DistilBertForSequenceClassification.from_pretrained(REPO, use_auth_token=HF_TOKEN)
model.eval()

@app.route("/predict", methods=["POST"])
def predict():
    try:
        data = request.get_json()
        results = []

        for item in data:
            # Build input text
            input_text = f"{item['category']} - {item['subcategory']} in {item['area']}. {item.get('comments', '')}"
            text_lower = input_text.lower()

            # Critical and low-urgency keyword lists
            critical_keywords = [
                "umuusok", "sunog", "amoy sunog", "spark", "kuryente",
                "leak", "baha", "gas", "short circuit", "smoke"
            ]

            low_keywords = [
                "lightbulb", "bumbilya", "ilaw", "palitan", "replace bulb",
                "painting", "door knob", "hinge", "minor", "cosmetic"
            ]

            # ✅ Priority rules
            if any(word in text_lower for word in critical_keywords):
                results.append({"priority_score": 5})
                continue
            elif any(word in text_lower for word in low_keywords):
                results.append({"priority_score": 2})
                continue

            # Tokenize & predict
            inputs = tokenizer(input_text, return_tensors="pt", truncation=True, padding=True)

            with torch.no_grad():
                outputs = model(**inputs)
                predicted_class = torch.argmax(outputs.logits, dim=1).item()

            # ✅ Convert back to 1–5 scale
            priority_score = predicted_class + 1

            # Soft correction: cap too-high or too-low predictions
            if priority_score >= 5 and not any(word in text_lower for word in critical_keywords):
                priority_score = 4  # reduce overly high unless it's critical
            elif priority_score <= 1:
                priority_score = 2  # lift up if too low

            results.append({"priority_score": priority_score})

        return jsonify(results)

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


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860)