import os
import secrets
from flask import Flask, render_template, request, jsonify
from PIL import Image
import numpy as np

# Try to import ML libraries, but don't crash if they are missing
HAS_ML = False
try:
    import tensorflow as tf
    from tensorflow.keras.models import load_model
    # Check if model exists
    model_path = os.path.join(os.path.dirname(__file__), 'cifar10_cnn_v1.h5')
    if os.path.exists(model_path):
        model = load_model(model_path)
        HAS_ML = True
        print("ML Model loaded successfully.")
except Exception as e:
    print(f"ML mode disabled: {e}")

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'uploads'
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB

# CIFAR-10 classes
CLASSES = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/predict', methods=['POST'])
def predict():
    if 'file' not in request.files:
        return jsonify({'success': False, 'error': 'No file part'})
    
    file = request.files['file']
    if file.filename == '':
        return jsonify({'success': False, 'error': 'No selected file'})

    try:
        # Save file
        filename = secrets.token_hex(8) + "_" + file.filename
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(filepath)

        # Process image
        img = Image.open(filepath).convert('RGB')
        img_resized = img.resize((32, 32)) 
        
        if HAS_ML:
            # Real Inference
            img_array = np.array(img_resized) / 255.0
            img_array = np.expand_dims(img_array, axis=0)
            predictions = model.predict(img_array)
            class_idx = np.argmax(predictions[0])
            confidence = float(predictions[0][class_idx])
            class_name = CLASSES[class_idx]
        else:
            # Mock Inference for demonstration if environment is broken
            # We use the filename hash to pick a "random" but consistent class for the same image
            hash_val = sum(ord(c) for c in filename)
            class_idx = hash_val % len(CLASSES)
            class_name = CLASSES[class_idx]
            confidence = 0.85 + (hash_val % 15) / 100.0

        return jsonify({
            'success': True,
            'class': class_name,
            'confidence': confidence,
            'mode': 'real' if HAS_ML else 'mock'
        })

    except Exception as e:
        return jsonify({'success': False, 'error': str(e)})

if __name__ == '__main__':
    if not os.path.exists(app.config['UPLOAD_FOLDER']):
        os.makedirs(app.config['UPLOAD_FOLDER'])
    app.run(debug=True, port=5000)