main.py

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__":
    uvicorn.run(app, host="0.0.0.0", port=8003)