import os
import re
from flask import Flask, request, jsonify, render_template
import requests
from PIL import Image
import pytesseract

app = Flask(__name__)

# --- 2026 Updated Gemini Setup ---
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")

# 1. Update to the current stable model (Gemini 3.1 Flash)
model = "gemini-3.1-flash"

# 2. Use the stable 'v1' endpoint with the correct model path
GEMINI_URL = "https://generativelanguage.googleapis.com/v1/models/" + model + ":generateContent?key=" + GEMINI_API_KEY

BLACKLISTED_UPIS = [
    "fraudster@upi", "scam123@okaxis", "fake_pay@ybl", 
    "urgent_pay@okhdfcbank", "win_cash@sbi"
]
SUSPICIOUS_KEYWORDS = [
    "collect", "request", "urgent", "win", "lottery", 
    "cashback", "claim", "reward", "fee"
]

# Gemini function
def call_gemini(prompt, system_instruction=""):
    if not GEMINI_API_KEY:
        return "Gemini API key is missing."

    payload = {
        "contents": [
            {
                "parts": [
                    {"text": f"{system_instruction}\n\nUser: {prompt}"}
                ]
            }
        ]
    }
    headers = {"Content-Type": "application/json"}
    
    try:
        response = requests.post(GEMINI_URL, headers=headers, json=payload)
        response.raise_for_status()
        data = response.json()
        return data["candidates"][0]["content"]["parts"][0]["text"]
    except Exception as e:
        return f"Error contacting Gemini: {str(e)}"

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

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

    file = request.files["image"]
    if file.filename == "":
        return jsonify({"error": "No file selected"}), 400

    try:
        img = Image.open(file.stream)
        text = pytesseract.image_to_string(img).lower()
        
        amount_match = re.search(r'(?:rs\.?|inr|₹|amount)\s*(\d+(?:,\d+)*(?:\.\d{1,2})?)', text)
        amount_str = amount_match.group(1).replace(",", "") if amount_match else "0"
        amount = float(amount_str) if amount_str else 0.0
        
        upi_match = re.search(r'[\w.-]+@[\w.-]+', text)
        upi_id = upi_match.group(0) if upi_match else "Not found"
        
        found_keywords = [kw for kw in SUSPICIOUS_KEYWORDS if kw in text]
        risk_level = "Safe"
        reasons = []
        
        if "request" in text and "success" in text:
            reasons.append("Conflicting terms detected.")
            risk_level = "High Risk"
        elif "received" in text and "pay" in text:
            reasons.append("Conflicting transaction terms.")
            risk_level = "High Risk"
            
        if found_keywords:
            reasons.append(f"Suspicious keywords: {', '.join(found_keywords)}.")
            if risk_level != "High Risk":
                risk_level = "Suspicious"
                
        if amount > 5000:
            reasons.append(f"High amount (₹{amount}).")
            if risk_level == "Safe":
                risk_level = "Suspicious"
                
        if not amount_match and not upi_match:
            reasons.append("Missing important details.")
            risk_level = "Suspicious"
            
        if not reasons:
            reasons.append("No major issues found.")
            
        prompt = f"Transaction text: {text[:200]}, Amount: {amount}, UPI: {upi_id}, Risk: {risk_level}. Explain in Hinglish."
        gemini_explanation = call_gemini(prompt, "You are a UPI fraud detection assistant. Give short Hinglish advice.")
        
        return jsonify({
            "risk_level": risk_level,
            "reasons": reasons,
            "extracted_text": text[:500],
            "explanation": gemini_explanation,
            "amount": amount,
            "upi_id": upi_id
        })
    except Exception as e:
        return jsonify({"error": str(e)}), 500

@app.route("/check_upi", methods=["POST"])
def check_upi():
    data = request.json
    upi_id = data.get("upi_id", "").strip().lower()

    if not upi_id or "@" not in upi_id:
        return jsonify({"error": "Invalid UPI ID"}), 400
        
    risk_level = "Safe"
    reason = "UPI format looks valid."
    
    if upi_id in BLACKLISTED_UPIS:
        risk_level = "High Risk"
        reason = "UPI ID is blacklisted."
    elif any(kw in upi_id for kw in ["cashback", "offer", "win", "reward", "urgent"]):
        risk_level = "Suspicious"
        reason = "Suspicious words in UPI ID."
        
    prompt = f"UPI: {upi_id}, Risk: {risk_level}, Reason: {reason}. Should user proceed?"
    gemini_explanation = call_gemini(prompt, "Give short Hinglish advice.")
    
    return jsonify({
        "risk_level": risk_level,
        "reason": reason,
        "explanation": gemini_explanation
    })

@app.route("/ask", methods=["POST"])
def ask():
    data = request.json
    message = data.get("message", "")

    if not message:
        return jsonify({"error": "Empty message"}), 400
        
    system_instruction = "You are a UPI Fraud Detection Assistant. Give Hinglish answers. Warn about fraud."
    response = call_gemini(message, system_instruction)
    
    return jsonify({"response": response})

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