Update main.py

main.py

@@ -2,21 +2,148 @@ from flask import Flask, jsonify, request
 from flask_cors import CORS
 from pymongo.mongo_client import MongoClient
 from pymongo.server_api import ServerApi
-from bson.objectid import ObjectId
+from bson import json_util, ObjectId
 import google.generativeai as genai
 import urllib.parse
-from models import UserSchema
 from flask_bcrypt import Bcrypt
+import jwt
 from flask_jwt_extended import JWTManager, create_access_token
-from middleware.authUser import auth_user
+# from middleware.authUser import auth_user
 from datetime import timedelta 
-from controllers.demo import get_initial_data
-from controllers.mindmap import saveMindmap, getMindmap, deleteMindmap, getMindmapByid
+# from controllers.demo import get_initial_data
+# from controllers.mindmap import saveMindmap, getMindmap, deleteMindmap, getMindmapByid
 import json
-from dotenv import load_dotenv
+# from dotenv import load_dotenv
 import os
 
-load_dotenv()
+# load_dotenv()
+
+def auth_user(next_function):
+    def middleware_function(*args, **kwargs):
+        try:
+            token = request.headers.get('Authorization', '').split(" ")[1]
+            is_custom_auth = len(token) < 500
+            decoded_data = None
+
+            if token and is_custom_auth:
+                secret_key = os.getenv('JWT_SECRET')
+                decoded_data = jwt.decode(token, secret_key, algorithms=["HS256"])
+                print(decoded_data)
+                dat = decoded_data.get('sub')
+                request.email = dat.get('email')
+                request.userId = dat.get('id')
+            else:
+                #google auth
+                decoded_data = jwt.decode(token)
+                request.email = decoded_data.get('email')
+                request.userId = decoded_data.get('sub')
+
+            return next_function(*args, **kwargs)
+
+        except Exception as e:
+            print(e)  # Log the error if needed
+            return jsonify({"error": "Unauthorized"}), 401
+
+    return middleware_function
+
+def get_initial_data():
+    initial_nodes = [
+        {
+            "id": "data-input",
+            "position": {"x": 0, "y": 0},
+            "data": {"label": "Data Input"}
+        },
+        {
+            "id": "data-preprocessing",
+            "position": {"x": 200, "y": 0},
+            "data": {"label": "Data Preprocessing"}
+        },
+        {
+            "id": "model-training",
+            "position": {"x": 400, "y": 0},
+            "data": {"label": "Model Training"}
+        },
+        {
+            "id": "model-evaluation",
+            "position": {"x": 0, "y": 200},
+            "data": {"label": "Model Evaluation"}
+        },
+        {
+            "id": "prediction",
+            "position": {"x": 200, "y": 200},
+            "data": {"label": "Prediction"}
+        },
+        {
+            "id": "data-visualization",
+            "position": {"x": 400, "y": 200},
+            "data": {"label": "Data Visualization"}
+        },
+    ]
+
+    initial_edges = [
+        {"id": "data-input-to-preprocessing", "source": "data-input", "target": "data-preprocessing"},
+        {"id": "preprocessing-to-training", "source": "data-preprocessing", "target": "model-training"},
+        {"id": "training-to-evaluation", "source": "model-training", "target": "model-evaluation"},
+        {"id": "training-to-prediction", "source": "model-training", "target": "prediction"},
+        {"id": "evaluation-to-visualization", "source": "model-evaluation", "target": "data-visualization"},
+        {"id": "prediction-to-visualization", "source": "prediction", "target": "data-visualization"}
+    ]
+
+    return jsonify({
+        "initialNodes": initial_nodes,
+        "initialEdges": initial_edges
+    })
+
+def saveMindmap(data, userId, savedMindmap):
+    try:
+        result = savedMindmap.insert_one({
+            "userId": userId,
+            "data": data,
+        })
+
+        print(result.inserted_id)
+        return jsonify({"result": str(result.inserted_id)}), 201
+    except Exception as e:
+        print(e) 
+        return jsonify({"error": "An error occurred"}), 500
+
+def getMindmap(userId, savedMindmap):
+    try:
+        results = savedMindmap.find({"userId": userId})
+        mindmaps = [{"_id": str(result["_id"]), "data": result["data"]} for result in results]
+
+        if mindmaps:
+            # Convert ObjectId to string for JSON serialization
+            return json_util.dumps({"data": mindmaps}), 200
+        else:
+            return jsonify({"msg": "No Mindmap stored"}), 404
+    except Exception as e:
+        print(e)  
+        return jsonify({"error": "An error occurred"}), 500
+    
+def getMindmapByid(userId, id, savedMindmap):
+    try:
+        object_id = ObjectId(id)
+        result = savedMindmap.find_one({"userId": userId, "_id": object_id})
+        if result:
+            return json_util.dumps({"data": result}), 200
+        else:
+            return jsonify({"msg": "Mindmap not found"}), 404
+    except Exception as e:
+        print(e)
+        return jsonify({"error": "An error occurred"}), 500
+    
+def deleteMindmap(userId, data, savedMindmap):
+    try:
+        object_id = ObjectId(data["_id"])
+        result = savedMindmap.delete_one({"userId": userId, "_id": object_id})
+        if result.deleted_count == 1:
+            return jsonify({"result": True}), 200
+        else:
+            return jsonify({"result": False, "message": "Mindmap not found"}), 404
+    except Exception as e:
+        print(e)
+        return jsonify({"message": "Something went wrong"}), 500
 
 app = Flask(__name__)
 
@@ -158,8 +285,6 @@ def tree():
         # return temp 
 
 
-
-
 @app.route('/tree/demo', methods=["POST"])
 def treeDemo():
     if request.method == 'POST':
@@ -484,7 +609,7 @@ def signup():
         "password": hashed_password
     })
 
-    print(result);
+    print(result)
 
     expires = timedelta(days=7)
     access_token = create_access_token(identity={"email": email, "id": str(result.inserted_id)}, expires_delta=expires)
@@ -561,7 +686,7 @@ def mindmapDelete():
 @app.route('/mindmap/demo', methods=['POST'])
 def mindmapDemo():
     data = request.json
-    print(data);
+    print(data)
     return get_initial_data(), 200