Update main.py

main.py

@@ -1,242 +1,238 @@
-from fastapi import FastAPI, HTTPException
-from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
-import uuid
-import json
-import os
-import uvicorn
-import shutil
-from pymongo import MongoClient
-from dotenv import load_dotenv
-from pathlib import Path
-
-# Import your scripts
-import generator
-import judge 
-
-# ==========================================
-# 1. SETUP ENV & DATABASE
-# ==========================================
-
-# Locate and load .env from backend folder
-current_file = Path(__file__).resolve()
-from dotenv import load_dotenv
-load_dotenv()  # HF will inject secrets automatically
-
-if env_path.exists():
-    load_dotenv(dotenv_path=env_path)
-    print("✅ [PS Service] .env loaded successfully.")
-else:
-    print(f"⚠️ [PS Service] Warning: .env not found at {env_path}")
-MONGO_URI = os.getenv("MONGO_URI")
-
-# Initialize MongoDB
-problems_collection = None
-if MONGO_URI:
-    try:
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+import uuid
+import json
+import os
+import uvicorn
+import shutil
+from pymongo import MongoClient
+from dotenv import load_dotenv
+from pathlib import Path
+
+# Import your scripts
+import generator
+import judge 
+
+# ==========================================
+# 1. SETUP ENV & DATABASE
+# ==========================================
+
+# Locate and load .env from backend folder
+
+from dotenv import load_dotenv
+load_dotenv()  # HF will inject secrets automatically
+
+
+MONGO_URI = os.getenv("MONGO_URI")
+
+# Initialize MongoDB
+problems_collection = None
+if MONGO_URI:
+    try:
         client = MongoClient(MONGO_URI, serverSelectionTimeoutMS=5000)
-
-        db = client["acemock"]
-        problems_collection = db["coding_problems"]
-        print("✅ [PS Service] Connected to MongoDB.")
-    except Exception as e:
-        print(f"❌ [PS Service] MongoDB Connection Error: {e}")
-else:
-    print("❌ [PS Service] MONGO_URI not found. Database features will fail.")
-
-app = FastAPI(title="AceMock Problem Solving")
-
-# Enable CORS
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["http://localhost:3001", "http://127.0.0.1:3001"], 
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-# Ensure temp directory exists for running code
-os.makedirs("temp", exist_ok=True)
-
-# --- HELPER: Normalize Language ---
-def normalize_lang(lang: str):
-    lang = lang.lower()
-    if lang == "javascript": return "js"
-    if lang == "c++": return "cpp"
-    return lang
-
-# --- DATA MODELS ---
-class ProblemRequest(BaseModel):
-    level: str  # "Fresh", "Junior", "Senior"
-    language: str # "cpp", "python", "javascript"
-
-class ExecuteRequest(BaseModel):
-    language: str
-    code: str
-    input: str = ""  # User's custom input
-
-class SubmissionRequest(BaseModel):
-    problem_id: str
-    user_code: str
-    language: str
-
-# --- ENDPOINT 1: GENERATE PROBLEM (MongoDB) ---
-@app.post("/generate-problem")
-def get_problem(req: ProblemRequest):
-    # ✅ FIX: Explicit None check for PyMongo 4+ compatibility
-    if problems_collection is None:
-        raise HTTPException(status_code=500, detail="Database connection not available.")
-
-    print(f"Generating {req.level} problem in {req.language}...")
-    
-    problem_data = None
-    # Retry logic (AI can fail occasionally)
-    for _ in range(3): 
-        try:
-            problem_data = generator.generate_problem_for_api(req.level, normalize_lang(req.language))
-            if problem_data: break
-        except Exception as e:
-            print(f"Generation attempt failed: {e}")
-    
-    if not problem_data:
-        raise HTTPException(status_code=500, detail="AI generation failed. Please try again.")
-
-    # Generate a unique ID for this session
-    problem_id = str(uuid.uuid4())
-    
-    # Add ID to data
-    problem_data["_id"] = problem_id
-    problem_data["level"] = req.level
-    problem_data["language"] = req.language
-    problem_data["created_at"] = str(uuid.uuid1()) # Timestamp roughly
-
-# # ✅ SAVE TO LOCAL JSON FILE
-#     try:
-#         # Ensure the 'problems' directory exists
-#         problems_dir = os.path.join(os.path.dirname(__file__), "problems")
-#         os.makedirs(problems_dir, exist_ok=True)
-        
-#         # Save the file using the problem_id as the name
-#         file_path = os.path.join(problems_dir, f"{problem_id}.json")
-#         with open(file_path, "w", encoding="utf-8") as f:
-#             json.dump(problem_data, f, indent=4, ensure_ascii=False)
-#         print(f"✅ Problem saved to local JSON: {file_path}")
-#     except Exception as e:
-#         print(f"⚠️ Failed to save problem to JSON file: {e}")
-
-    # ✅ SAVE TO MONGODB
-    try:
-        problems_collection.insert_one(problem_data)
-        print(f"✅ Problem saved to MongoDB with ID: {problem_id}")
-    except Exception as e:
-        print(f"❌ Failed to save to DB: {e}")
-        raise HTTPException(status_code=500, detail="Database save failed.")
-
-    # Return ONLY what the user needs to see (Hide test cases!)
-    return {
-        "problem_id": problem_id,
-        "title": problem_data['title'],
-        "description": problem_data['description'],
-        "input_format": problem_data['input_format'],
-        "output_format": problem_data['output_format'],
-        "samples": problem_data['samples'],
-        "solution_code": problem_data.get('solution_code') 
-    }
-
-# --- ENDPOINT 2: EXECUTE CODE (Run Button - No DB needed) ---
-@app.post("/execute")
-def execute_code(req: ExecuteRequest):
-    lang_key = normalize_lang(req.language)
-    
-    # Create temp file
-    ext_map = {"cpp": ".cpp", "python": ".py", "js": ".js"}
-    filename = f"run_{uuid.uuid4()}{ext_map.get(lang_key, '.txt')}"
-    filepath = f"temp/{filename}"
-    
-    with open(filepath, "w") as f:
-        f.write(req.code)
-        
-    try:
-        # Run using the exact same Judge logic as submit
-        result = judge.run_test_case(filepath, lang_key, req.input)
-        
-        # Cleanup
-        if os.path.exists(filepath): os.remove(filepath)
-        
-        return result 
-        
-    except Exception as e:
-        if os.path.exists(filepath): os.remove(filepath)
-        return {"status": "Error", "output": str(e)}
-
-# --- ENDPOINT 3: SUBMIT SOLUTION (Grading via MongoDB) ---
-@app.post("/submit")
-def submit_solution(req: SubmissionRequest):
-    if problems_collection is None:
-        raise HTTPException(status_code=500, detail="Database connection not available.")
-
-    # 1. Fetch Problem from MongoDB
-    problem_data = problems_collection.find_one({"_id": req.problem_id})
-    
-    if not problem_data:
-        raise HTTPException(status_code=404, detail="Problem ID not found in database.")
-
-    lang_key = normalize_lang(req.language)
-    ext_map = {"cpp": ".cpp", "python": ".py", "js": ".js"}
-    filename = f"sub_{req.problem_id}{ext_map.get(lang_key, '.txt')}"
-    user_file_path = f"temp/{filename}"
-    
-    with open(user_file_path, "w") as f:
-        f.write(req.user_code)
-
-    # 3. Run Test Cases (Retrieved from DB)
-    results = []
-    passed_count = 0
-    total_tests = len(problem_data['test_cases'])
-
-    for test in problem_data['test_cases']:
-        run_res = judge.run_test_case(
-            filepath=user_file_path, 
-            lang=lang_key, 
-            input_str=test['input']
-        )
-        
-        if run_res['status'] == "Success":
-            if run_res['output'].strip() == test['expected_output'].strip():
-                results.append({"id": test['id'], "status": "Passed"})
-                passed_count += 1
-            else:
-                results.append({
-                    "id": test['id'], 
-                    "status": "Wrong Answer", 
-                    "your_output": run_res['output'][:50], 
-                    "expected": test['expected_output'].strip()[:100] 
-                })
-        else:
-            results.append({"id": test['id'], "status": run_res['status'], "details": run_res['output']})
-
-# 4. Cleanup (Delete temp file)
-    if os.path.exists(user_file_path):
-        os.remove(user_file_path)
-
-    db_solutions = problem_data.get("solution_code", {})
-    
-    if isinstance(db_solutions, dict):
-        # Fetch the requested language. Fallback to python if it's missing.
-        final_solution = db_solutions.get(lang_key, db_solutions.get("python", "# Solution missing"))
-    else:
-        # Fallback logic just in case an old problem (string format) is tested
-        final_solution = str(db_solutions)
-
-    # 5. Return Score AND The Match Solution
-    return {
-        "passed": passed_count,
-        "total": total_tests,
-        "score": round((passed_count / total_tests) * 100, 1),
-        "details": results,
-        "solution_code": final_solution
-    }
-
-# --- RUN ON PORT 8003 ---
-if __name__ == "__main__":
+
+        db = client["acemock"]
+        problems_collection = db["coding_problems"]
+        print("✅ [PS Service] Connected to MongoDB.")
+    except Exception as e:
+        print(f"❌ [PS Service] MongoDB Connection Error: {e}")
+else:
+    print("❌ [PS Service] MONGO_URI not found. Database features will fail.")
+
+app = FastAPI(title="AceMock Problem Solving")
+
+# Enable CORS
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["http://localhost:3001", "http://127.0.0.1:3001"], 
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Ensure temp directory exists for running code
+os.makedirs("temp", exist_ok=True)
+
+# --- HELPER: Normalize Language ---
+def normalize_lang(lang: str):
+    lang = lang.lower()
+    if lang == "javascript": return "js"
+    if lang == "c++": return "cpp"
+    return lang
+
+# --- DATA MODELS ---
+class ProblemRequest(BaseModel):
+    level: str  # "Fresh", "Junior", "Senior"
+    language: str # "cpp", "python", "javascript"
+
+class ExecuteRequest(BaseModel):
+    language: str
+    code: str
+    input: str = ""  # User's custom input
+
+class SubmissionRequest(BaseModel):
+    problem_id: str
+    user_code: str
+    language: str
+
+# --- ENDPOINT 1: GENERATE PROBLEM (MongoDB) ---
+@app.post("/generate-problem")
+def get_problem(req: ProblemRequest):
+    # ✅ FIX: Explicit None check for PyMongo 4+ compatibility
+    if problems_collection is None:
+        raise HTTPException(status_code=500, detail="Database connection not available.")
+
+    print(f"Generating {req.level} problem in {req.language}...")
+    
+    problem_data = None
+    # Retry logic (AI can fail occasionally)
+    for _ in range(3): 
+        try:
+            problem_data = generator.generate_problem_for_api(req.level, normalize_lang(req.language))
+            if problem_data: break
+        except Exception as e:
+            print(f"Generation attempt failed: {e}")
+    
+    if not problem_data:
+        raise HTTPException(status_code=500, detail="AI generation failed. Please try again.")
+
+    # Generate a unique ID for this session
+    problem_id = str(uuid.uuid4())
+    
+    # Add ID to data
+    problem_data["_id"] = problem_id
+    problem_data["level"] = req.level
+    problem_data["language"] = req.language
+    problem_data["created_at"] = str(uuid.uuid1()) # Timestamp roughly
+
+# # ✅ SAVE TO LOCAL JSON FILE
+#     try:
+#         # Ensure the 'problems' directory exists
+#         problems_dir = os.path.join(os.path.dirname(__file__), "problems")
+#         os.makedirs(problems_dir, exist_ok=True)
+        
+#         # Save the file using the problem_id as the name
+#         file_path = os.path.join(problems_dir, f"{problem_id}.json")
+#         with open(file_path, "w", encoding="utf-8") as f:
+#             json.dump(problem_data, f, indent=4, ensure_ascii=False)
+#         print(f"✅ Problem saved to local JSON: {file_path}")
+#     except Exception as e:
+#         print(f"⚠️ Failed to save problem to JSON file: {e}")
+
+    # ✅ SAVE TO MONGODB
+    try:
+        problems_collection.insert_one(problem_data)
+        print(f"✅ Problem saved to MongoDB with ID: {problem_id}")
+    except Exception as e:
+        print(f"❌ Failed to save to DB: {e}")
+        raise HTTPException(status_code=500, detail="Database save failed.")
+
+    # Return ONLY what the user needs to see (Hide test cases!)
+    return {
+        "problem_id": problem_id,
+        "title": problem_data['title'],
+        "description": problem_data['description'],
+        "input_format": problem_data['input_format'],
+        "output_format": problem_data['output_format'],
+        "samples": problem_data['samples'],
+        "solution_code": problem_data.get('solution_code') 
+    }
+
+# --- ENDPOINT 2: EXECUTE CODE (Run Button - No DB needed) ---
+@app.post("/execute")
+def execute_code(req: ExecuteRequest):
+    lang_key = normalize_lang(req.language)
+    
+    # Create temp file
+    ext_map = {"cpp": ".cpp", "python": ".py", "js": ".js"}
+    filename = f"run_{uuid.uuid4()}{ext_map.get(lang_key, '.txt')}"
+    filepath = f"temp/{filename}"
+    
+    with open(filepath, "w") as f:
+        f.write(req.code)
+        
+    try:
+        # Run using the exact same Judge logic as submit
+        result = judge.run_test_case(filepath, lang_key, req.input)
+        
+        # Cleanup
+        if os.path.exists(filepath): os.remove(filepath)
+        
+        return result 
+        
+    except Exception as e:
+        if os.path.exists(filepath): os.remove(filepath)
+        return {"status": "Error", "output": str(e)}
+
+# --- ENDPOINT 3: SUBMIT SOLUTION (Grading via MongoDB) ---
+@app.post("/submit")
+def submit_solution(req: SubmissionRequest):
+    if problems_collection is None:
+        raise HTTPException(status_code=500, detail="Database connection not available.")
+
+    # 1. Fetch Problem from MongoDB
+    problem_data = problems_collection.find_one({"_id": req.problem_id})
+    
+    if not problem_data:
+        raise HTTPException(status_code=404, detail="Problem ID not found in database.")
+
+    lang_key = normalize_lang(req.language)
+    ext_map = {"cpp": ".cpp", "python": ".py", "js": ".js"}
+    filename = f"sub_{req.problem_id}{ext_map.get(lang_key, '.txt')}"
+    user_file_path = f"temp/{filename}"
+    
+    with open(user_file_path, "w") as f:
+        f.write(req.user_code)
+
+    # 3. Run Test Cases (Retrieved from DB)
+    results = []
+    passed_count = 0
+    total_tests = len(problem_data['test_cases'])
+
+    for test in problem_data['test_cases']:
+        run_res = judge.run_test_case(
+            filepath=user_file_path, 
+            lang=lang_key, 
+            input_str=test['input']
+        )
+        
+        if run_res['status'] == "Success":
+            if run_res['output'].strip() == test['expected_output'].strip():
+                results.append({"id": test['id'], "status": "Passed"})
+                passed_count += 1
+            else:
+                results.append({
+                    "id": test['id'], 
+                    "status": "Wrong Answer", 
+                    "your_output": run_res['output'][:50], 
+                    "expected": test['expected_output'].strip()[:100] 
+                })
+        else:
+            results.append({"id": test['id'], "status": run_res['status'], "details": run_res['output']})
+
+# 4. Cleanup (Delete temp file)
+    if os.path.exists(user_file_path):
+        os.remove(user_file_path)
+
+    db_solutions = problem_data.get("solution_code", {})
+    
+    if isinstance(db_solutions, dict):
+        # Fetch the requested language. Fallback to python if it's missing.
+        final_solution = db_solutions.get(lang_key, db_solutions.get("python", "# Solution missing"))
+    else:
+        # Fallback logic just in case an old problem (string format) is tested
+        final_solution = str(db_solutions)
+
+    # 5. Return Score AND The Match Solution
+    return {
+        "passed": passed_count,
+        "total": total_tests,
+        "score": round((passed_count / total_tests) * 100, 1),
+        "details": results,
+        "solution_code": final_solution
+    }
+
+# --- RUN ON PORT 8003 ---
+if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=8003)