from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.image import load_img, img_to_array
import numpy as np
from flask import Flask, request, render_template, jsonify
from werkzeug.utils import secure_filename
import os

# Initialize Flask app
app = Flask(__name__)

# Load model
model_path = os.path.join(os.path.dirname(__file__), "SoilNet.keras")
SoilNet = load_model(model_path)

# Classes dictionary
classes = {
    0: "Alluvial Soil:-{ Rice, Wheat, Sugarcane, Maize, Cotton, Soyabean, Jute }",
    1: "Black Soil:-{ Virginia, Wheat, Jowar, Millets, Linseed, Castor, Sunflower }",
    2: "Clay Soil:-{ Rice, Lettuce, Chard, Broccoli, Cabbage, Snap Beans }",
    3: "Red Soil:-{ Cotton, Wheat, Pulses, Millets, Oil Seeds, Potatoes }"
}

# API Key (set this securely in prod)
API_KEY = "your-secret-api-key-1234"

# Prediction function
def model_predict(image_path, model):
    image = load_img(image_path, target_size=(224, 224))
    image = img_to_array(image) / 255.0
    image = np.expand_dims(image, axis=0)

    result = np.argmax(model.predict(image), axis=-1)[0]
    prediction = classes[result]

    if result == 0:
        return "Alluvial", "Alluvial.html"
    elif result == 1:
        return "Black", "Black.html"
    elif result == 2:
        return "Clay", "Clay.html"
    elif result == 3:
        return "Red", "Red.html"

# Route: Home (form)
@app.route('/', methods=['GET'])
def index():
    return render_template('index.html')

# Route: Form-based upload + result display
@app.route('/predict', methods=['POST'])
def predict():
    file = request.files.get('image')
    if not file or file.filename == '':
        return "No image uploaded", 400

    # Validate extension
    filename = secure_filename(file.filename)
    if not filename.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.gif')):
        return "Unsupported file type", 400

    # Save image
    upload_folder = os.path.join(os.path.dirname(__file__), 'static', 'user_uploaded')
    os.makedirs(upload_folder, exist_ok=True)
    file_path = os.path.join(upload_folder, filename)
    file.save(file_path)

    # Check image is valid
    try:
        _ = load_img(file_path)
    except Exception as e:
        os.remove(file_path)
        return f"Invalid image file: {e}", 400

    pred, output_page = model_predict(file_path, SoilNet)
    user_image_path = os.path.join('static', 'user_uploaded', filename)

    return render_template(output_page, pred_output=pred, user_image=user_image_path)

# Route: API endpoint with API key
@app.route('/api/predict', methods=['POST'])
def api_predict():
    # Check for API key
    key = request.headers.get('x-api-key')
    if key != API_KEY:
        return jsonify({"error": "Unauthorized"}), 401

    # Validate image
    file = request.files.get('image')
    if not file or file.filename == '':
        return jsonify({"error": "No image uploaded"}), 400

    filename = secure_filename(file.filename)
    if not filename.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.gif')):
        return jsonify({"error": "Unsupported file type"}), 400

    # Save temp image
    api_temp_folder = os.path.join(os.path.dirname(__file__), 'static', 'api_temp')
    os.makedirs(api_temp_folder, exist_ok=True)
    file_path = os.path.join(api_temp_folder, filename)
    file.save(file_path)

    try:
        _ = load_img(file_path)
    except Exception as e:
        os.remove(file_path)
        return jsonify({"error": f"Invalid image: {str(e)}"}), 400

    pred, _ = model_predict(file_path, SoilNet)
    os.remove(file_path)  # Optional: delete after prediction

    return jsonify({"soil_type": pred})

# Start the app
if __name__ == '__main__':
    app.run(debug=True, threaded=False)