from fastapi import FastAPI, UploadFile, File, HTTPException, Query
from fastapi.responses import JSONResponse, StreamingResponse
import uvicorn
import io
import json
import os
import tempfile
import numpy as np
import cv2
from PIL import Image
from pdf2image import convert_from_bytes

GENAI_API_KEY = os.getenv("GENAI_API_KEY")
if not GENAI_API_KEY:
    raise Exception("GENAI_API_KEY not set in .env file.")

# Import the Google GenAI client libraries.
from google import genai
from google.genai import types

# Initialize the GenAI client with the API key from .env.
client = genai.Client(api_key=GENAI_API_KEY)

app = FastAPI(title="Student Result Card API")

# Use the system temporary directory
TEMP_FOLDER = tempfile.gettempdir()

# -----------------------------
# Preprocessing Methods
# -----------------------------
def preprocess_candidate_info(image_cv):
    """
    Preprocess the image to extract the candidate information region.
    Region is defined by a mask covering the top-left portion.
    """
    height, width = image_cv.shape[:2]
    mask = np.zeros((height, width), dtype="uint8")
    margin_top = int(height * 0.10)
    margin_bottom = int(height * 0.25)
    cv2.rectangle(mask, (0, margin_top), (width, height - margin_bottom), 255, -1)
    masked = cv2.bitwise_and(image_cv, image_cv, mask=mask)
    coords = cv2.findNonZero(mask)
    x, y, w, h = cv2.boundingRect(coords)
    cropped = masked[y:y+h, x:x+w]
    return Image.fromarray(cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB))

def preprocess_mcq(image_cv):
    """
    Preprocess the image to extract the MCQ answers region (questions 1 to 10).
    Region is defined by a mask on the left side of the page.
    """
    height, width = image_cv.shape[:2]
    mask = np.zeros((height, width), dtype="uint8")
    margin_top = int(height * 0.27)
    margin_bottom = int(height * 0.23)
    right_boundary = int(width * 0.35)
    cv2.rectangle(mask, (0, margin_top), (right_boundary, height - margin_bottom), 255, -1)
    masked = cv2.bitwise_and(image_cv, image_cv, mask=mask)
    coords = cv2.findNonZero(mask)
    x, y, w, h = cv2.boundingRect(coords)
    cropped = masked[y:y+h, x:x+w]
    return Image.fromarray(cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB))

def preprocess_free_response(image_cv):
    """
    Preprocess the image to extract the free-response answers region (questions 11 to 15).
    Region is defined by a mask on the middle-right part of the page.
    """
    height, width = image_cv.shape[:2]
    mask = np.zeros((height, width), dtype="uint8")
    margin_top = int(height * 0.27)
    margin_bottom = int(height * 0.38)
    left_boundary = int(width * 0.35)
    right_boundary = int(width * 0.68)
    cv2.rectangle(mask, (left_boundary, margin_top), (right_boundary, height - margin_bottom), 255, -1)
    masked = cv2.bitwise_and(image_cv, image_cv, mask=mask)
    coords = cv2.findNonZero(mask)
    x, y, w, h = cv2.boundingRect(coords)
    cropped = masked[y:y+h, x:x+w]
    return Image.fromarray(cv2.cvtColor(cropped, cv2.COLOR_BGR2RGB))

def preprocess_full_answers(image_cv):
    """
    For extracting the correct answer key, we assume the entire page contains the answers.
    """
    return Image.fromarray(cv2.cvtColor(image_cv, cv2.COLOR_BGR2RGB))

# -----------------------------
# Extraction Methods using Gemini
# -----------------------------
def extract_json_from_output(output_str):
    """
    Extracts a JSON object from a string containing extra text.
    """
    start = output_str.find('{')
    end = output_str.rfind('}')
    if start == -1 or end == -1:
        return None
    json_str = output_str[start:end+1]
    try:
        return json.loads(json_str)
    except json.JSONDecodeError:
        return None

def get_student_info(image_input):
    """
    Extracts candidate information from an image.
    """
    output_format = """
Answer in the following JSON format. Do not write anything else:
{
  "Candidate Info": {
    "Name": "<name>",
    "Number": "<number>",
    "Country": "<country>",
    "Level": "<level>"
  }
}
"""
    prompt = f"""
You are an assistant that extracts candidate information from an image.
The image contains details including name, candidate number, country, and level.
Extract the information accurately and provide the result in JSON using the format below:
{output_format}
"""
    response = client.models.generate_content(model="gemini-2.0-flash", contents=[prompt, image_input])
    return extract_json_from_output(response.text)

def get_mcq_answers(image_input):
    """
    Extracts multiple-choice answers (questions 1 to 10) from an image.
    """
    output_format = """
Answer in the following JSON format do not write anything else:
{
  "Answers": {
    "1": "<option>",
    "2": "<option>",
    "3": "<option>",
    "4": "<option>",
    "5": "<option>",
    "6": "<option>",
    "7": "<option>",
    "8": "<option>",
    "9": "<option>",
    "10": "<option>"
  }
}
"""
    prompt = f"""
You are an assistant that extracts MCQ answers from an image.
The image is a screenshot of a 10-question multiple-choice answer sheet.
Extract which option is marked for each question (1 to 10) and provide the answers in JSON using the format below:
{output_format}
"""
    response = client.models.generate_content(model="gemini-2.0-flash", contents=[prompt, image_input])
    return extract_json_from_output(response.text)

def get_free_response_answers(image_input):
    """
    Extracts free-text answers (questions 11 to 15) from an image.
    """
    output_format = """
Answer in the following JSON format. Do not write anything else:
{
  "Free Answers": {
    "11": "<answer for question 11>",
    "12": "<answer for question 12>",
    "13": "<answer for question 13>",
    "14": "<answer for question 14>",
    "15": "<answer for question 15>"
  }
}
"""
    prompt = f"""
You are an assistant that extracts free-text answers from an image.
The image contains responses for questions 11 to 15.
Extract the answers accurately and provide the result in JSON using the format below:
{output_format}
"""
    response = client.models.generate_content(model="gemini-2.0-flash", contents=[prompt, image_input])
    return extract_json_from_output(response.text)

def get_all_answers(image_input):
    """
    Extracts all answers (questions 1 to 15) from an image of the correct answer key.
    """
    output_format = """
Answer in the following JSON format. Do not write anything else:
{
  "Answers": {
    "1": "<option or text>",
    "2": "<option or text>",
    "3": "<option or text>",
    "4": "<option or text>",
    "5": "<option or text>",
    "6": "<option or text>",
    "7": "<option or text>",
    "8": "<option or text>",
    "9": "<option or text>",
    "10": "<option or text>",
    "11": "<free-text answer>",
    "12": "<free-text answer>",
    "13": "<free-text answer>",
    "14": "<free-text answer>",
    "15": "<free-text answer>"
  }
}
"""
    prompt = f"""
You are an assistant that extracts answers from an image.
The image is a screenshot of an answer sheet containing 15 questions.
For questions 1 to 10, the answers are multiple-choice selections.
For questions 11 to 15, the answers are free-text responses.
Extract the answer for each question and provide the result in JSON using the format below:
{output_format}
"""
    response = client.models.generate_content(model="gemini-2.0-flash", contents=[prompt, image_input])
    return extract_json_from_output(response.text)

# -----------------------------
# Method to calculate result card
# -----------------------------
def calculate_result(student_info, student_mcq, student_free, correct_answers):
    """
    Compares student's answers with the correct answers, calculates marks and percentage,
    and returns a result card in JSON.
    """
    student_all = {}
    if student_mcq and "Answers" in student_mcq:
        student_all.update(student_mcq["Answers"])
    if student_free and "Free Answers" in student_free:
        student_all.update(student_free["Free Answers"])
    
    correct_all = correct_answers.get("Answers", {})
    total_questions = 15
    marks = 0
    detailed = {}
    
    for q in map(str, range(1, total_questions + 1)):
        student_ans = student_all.get(q, "").strip()
        correct_ans = correct_all.get(q, "").strip()
        if student_ans == correct_ans:
            marks += 1
            detailed[q] = {"Student": student_ans, "Correct": correct_ans, "Result": "Correct"}
        else:
            detailed[q] = {"Student": student_ans, "Correct": correct_ans, "Result": "Incorrect"}
    
    percentage = (marks / total_questions) * 100
    result_card = {
        "Candidate Info": student_info.get("Candidate Info", {}),
        "Total Marks": marks,
        "Total Questions": total_questions,
        "Percentage": percentage,
        "Detailed Results": detailed
    }
    return result_card

# -----------------------------
# API Endpoint to process PDFs and return student result cards
# -----------------------------
@app.post("/process")
async def process_pdfs(
    student_pdf: UploadFile = File(...),
    answer_key_pdf: UploadFile = File(...),
    download: bool = Query(True, description="Set to true to download result card list as a JSON file")
):
    try:
        # Read student PDF bytes and convert to images
        student_bytes = await student_pdf.read()
        student_images = convert_from_bytes(student_bytes)
        
        # Read answer key PDF bytes and convert to images; assume correct key is in the last page.
        answer_key_bytes = await answer_key_pdf.read()
        answer_key_images = convert_from_bytes(answer_key_bytes)
        last_page = answer_key_images[-1]
        last_page_cv = np.array(last_page)
        last_page_cv = cv2.cvtColor(last_page_cv, cv2.COLOR_RGB2BGR)
        correct_image = preprocess_full_answers(last_page_cv)
        correct_answers = get_all_answers(correct_image)
        
        student_result_cards = []
        
        # Process each student page.
        for idx, page in enumerate(student_images):
            page_cv = np.array(page)
            page_cv = cv2.cvtColor(page_cv, cv2.COLOR_RGB2BGR)
            student_info_image = preprocess_candidate_info(page_cv)
            mcq_image = preprocess_mcq(page_cv)
            free_image = preprocess_free_response(page_cv)
            
            student_info = get_student_info(student_info_image)
            student_mcq  = get_mcq_answers(mcq_image)
            student_free = get_free_response_answers(free_image)
            
            result_card = calculate_result(student_info, student_mcq, student_free, correct_answers)
            result_card["Student Index"] = idx + 1
            student_result_cards.append(result_card)
        
        response_data = {"result_cards": student_result_cards}
        
        if download:
            # Create downloadable JSON file and save to system temp folder
            json_bytes = json.dumps(response_data, indent=2).encode("utf-8")
            file_path = os.path.join(TEMP_FOLDER, "result_cards.json")
            with open(file_path, "wb") as f:
                f.write(json_bytes)
            return StreamingResponse(
                io.BytesIO(json_bytes),
                media_type="application/json",
                headers={"Content-Disposition": "attachment; filename=result_cards.json"}
            )
        else:
            return JSONResponse(content=response_data)
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# -----------------------------
# New Download Endpoint
# -----------------------------
@app.get("/download")
async def download_result_cards():
    """
    Returns the previously generated result_cards.json file from the system temporary folder.
    """
    file_path = os.path.join(TEMP_FOLDER, "result_cards.json")
    if not os.path.exists(file_path):
         raise HTTPException(status_code=404, detail="File not found")
    return StreamingResponse(
         open(file_path, "rb"),
         media_type="application/json",
         headers={"Content-Disposition": "attachment; filename=result_cards.json"}
    )

@app.get("/")
async def root():
    return {
        "message": "Welcome to the Student Result Card API.",
        "usage": "POST PDFs to /process with 'student_pdf' and 'answer_key_pdf' fields. Use ?download=true for file download or GET /download to re-download the JSON file."
    }

if __name__ == "__main__":
    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True)