from flask import Flask, request, jsonify
from flask_cors import CORS
import tensorflow as tf
import numpy as np
from PIL import Image
import io

app = Flask(__name__)
CORS(app)

interpreter = tf.lite.Interpreter(model_path="final_resnet.tflite")
interpreter.allocate_tensors()

input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
class_names = {
    0: 'FreshApple', 1: 'FreshBanana', 2: 'FreshGrape', 3: 'FreshGuava',
    4: 'FreshJujube', 5: 'FreshOrange', 6: 'FreshPomegranate', 7: 'FreshStrawberry',
    8: 'RottenApple', 9: 'RottenBanana', 10: 'RottenGrape', 11: 'RottenGuava',
    12: 'RottenJujube', 13: 'RottenOrange', 14: 'RottenPomegranate', 15: 'RottenStrawberry'
}

def preprocess_image(img_bytes):
    img = Image.open(io.BytesIO(img_bytes)).resize((224, 224))
    img = np.array(img) / 255.0
    return np.expand_dims(img, axis=0)

@app.route('/')
def home():
    return "API is running!"

@app.route('/favicon.ico')
def favicon():
    return '', 204  # or serve a real favicon if needed
    

@app.route('/predict', methods=['POST'])
def predict():
    file = request.files.get('image')
    if not file:
        return jsonify({'error': 'No image provided'}), 400

    img = preprocess_image(file.read()).astype(np.float32)
    interpreter.set_tensor(input_details[0]['index'], img)
    interpreter.invoke()
    pred = interpreter.get_tensor(output_details[0]['index'])[0]

    sorted_indices = np.argsort(pred)[::-1]
    top1, top2 = sorted_indices[:2]
    top1_label = class_names[top1]
    top2_label = class_names[top2]
    top1_conf = pred[top1] * 100
    top2_conf = pred[top2] * 100

    result = {}
    
    # Extract categories
    top1_cat = "Fresh" if "Fresh" in top1_label else "Rotten"
    top2_cat = "Fresh" if "Fresh" in top2_label else "Rotten"

    # Format prediction output (without fruit name)
    prediction_text = f"{top1_cat} ({top1_conf:.1f}%)"

    # Rule 1: If prediction > 80%
    if top1_conf >= 80:
        result['prediction'] = prediction_text
        result['message'] = "✅ Safe to eat!" if top1_cat == "Fresh" else "⚠️ Not safe to eat!"
    
    # Rule 3: < 80%, top two different categories
    elif top1_cat != top2_cat:
        result['prediction'] = f"{top1_cat} ({top1_conf:.1f}%) vs {top2_cat} ({top2_conf:.1f}%)"
        if top1_cat == "Rotten":
            result['message'] = "⚠️ Do not eat this fruit."
        else:
            result['message'] = "🍃 Seems fresh, but be cautious."
    
    # Rule 4: < 80%, same category (e.g., FreshApple vs FreshBanana)
    elif top1_cat == top2_cat:
        result['prediction'] = prediction_text
        if top1_cat == "Rotten":
            result['message'] = "⚠️ Do not eat this fruit."
        else:
            result['message'] = "🍃 Seems fresh, but be cautious."

    # Fallback message (if none matched)
    else:
        result['prediction'] = prediction_text
        result['message'] = "⚠️ Unable to confidently predict freshness."

    # Optional extra info
    result['confidence'] = round(top1_conf, 2)
    result['fruitType'] = top1_label
    result['fresh'] = 'Fresh' in top1_label

    return jsonify(result)