import os
import torch
import easyocr
from flask import Flask, request, jsonify, render_template
from flask_cors import CORS
import os
from PIL import Image
from transformers import AutoImageProcessor, AutoModelForImageClassification

app = Flask(__name__)
CORS(app)

MODEL_NAME = "prithivMLmods/deepfake-detector-model-v1"
UPLOAD_FOLDER = "uploads"
os.makedirs(UPLOAD_FOLDER, exist_ok=True)

print("Loading AI Models...")

processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
model = AutoModelForImageClassification.from_pretrained(MODEL_NAME)
reader = easyocr.Reader(['en'], gpu=False)

print("All Models Loaded Successfully!")

@app.route("/health")
def health():
    return {"status": "ok"}, 200

@app.route("/about")
def about():
    return render_template("about.html")

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

@app.route("/terms")
def terms():
    return render_template("terms.html")

@app.route("/analyze", methods=["POST"])
def analyze():
    file_path = None
    try:
        if "image" not in request.files:
            return jsonify({"error": "No image uploaded"}), 400

        file = request.files["image"]

        if file.filename == "":
            return jsonify({"error": "Empty file"}), 400

        file_path = os.path.join(UPLOAD_FOLDER, file.filename)
        file.save(file_path)

        img = Image.open(file_path)
        width, height = img.size

        bytes_size = os.path.getsize(file_path)
        size_mb = round(bytes_size / (1024 * 1024), 2)

        img_rgb = img.convert("RGB")
        result = reader.readtext(file_path)
        text = " ".join([r[1] for r in result]).strip()

        inputs = processor(images=img_rgb, return_tensors="pt")
        with torch.no_grad():
            outputs = model(**inputs)
            probs = torch.nn.functional.softmax(outputs.logits, dim=1).squeeze().tolist()

        fake_score = probs[0]
        real_score = probs[1]
        is_real = real_score > fake_score
        confidence = round((real_score if is_real else fake_score) * 100, 2)

        return jsonify({
            "analysis": {
                "is_real": bool(is_real),
                "confidence": confidence,
                "reason": "Likely Real" if is_real else "Likely AI Generated"
            },
            "text": text,
            "metadata": {
                "width": int(width),
                "height": int(height),
                "size_mb": float(size_mb)
            },
            "labels": ["deepfake-check", "ocr"]
        })

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

    finally:
        if file_path and os.path.exists(file_path):
            try:
                os.remove(file_path)
            except:
                pass

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 5000))
    app.run(host='0.0.0.0', port=port)