chamthi.py

import cv2
import numpy as np
import imutils
from imutils.contours import sort_contours
import json
import base64
import logging
import os
import uuid
import google.generativeai as genai
from fastapi import APIRouter, File, UploadFile, Form
from supabase import create_client, Client

# Cấu hình Log
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

router = APIRouter(prefix="/api/omr", tags=["ChamThiOMR"])

# --- CẤU HÌNH GEMINI API ---
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
if GEMINI_API_KEY:
    genai.configure(api_key=GEMINI_API_KEY)

# --- KẾT NỐI SUPABASE ---
SUPABASE_URL = os.environ.get("SUPABASE_URL")
SUPABASE_KEY = os.environ.get("SUPABASE_KEY")
supabase: Client = None

if SUPABASE_URL and SUPABASE_KEY:
    try:
        supabase = create_client(SUPABASE_URL, SUPABASE_KEY)
        logger.info("✅ OMR Module: Đã kết nối Supabase")
    except Exception as e:
        logger.error(f"❌ OMR Module: Lỗi kết nối Supabase: {e}")

# --- BIẾN TOÀN CỤC (CACHE RAM) ---
sessions = {}

# --- CÁC HÀM XỬ LÝ ẢNH ---

def order_points(pts):
    """Sắp xếp 4 điểm góc"""
    rect = np.zeros((4, 2), dtype="float32")
    s = pts.sum(axis=1)
    rect[0] = pts[np.argmin(s)]
    rect[2] = pts[np.argmax(s)]
    diff = np.diff(pts, axis=1)
    rect[1] = pts[np.argmin(diff)]
    rect[3] = pts[np.argmax(diff)]
    return rect

def four_point_transform(image, pts):
    """Warp ảnh"""
    rect = order_points(pts)
    (tl, tr, br, bl) = rect
    maxWidth = max(int(np.linalg.norm(br - bl)), int(np.linalg.norm(tr - tl)))
    maxHeight = max(int(np.linalg.norm(tr - br)), int(np.linalg.norm(tl - bl)))
    dst = np.array([[0, 0],[maxWidth - 1, 0],[maxWidth - 1, maxHeight - 1],[0, maxHeight - 1]], dtype="float32")
    M = cv2.getPerspectiveTransform(rect, dst)
    return cv2.warpPerspective(image, M, (maxWidth, maxHeight))

def read_bubbles(roi, cols, rows, draw_on_me=None, offset=(0,0), bubble_thresh=55):
    """Hàm đọc bong bóng. bubble_thresh=55: Ngưỡng thấp để bắt được bút chì mờ/mực xanh."""
    gray = cv2.cvtColor(roi, cv2.COLOR_BGR2GRAY)
    thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU)[1]
    
    cnts = cv2.findContours(thresh.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    cnts = imutils.grab_contours(cnts)
    bubbles = []
    
    h_img, w_img = roi.shape[:2]
    min_w = w_img // (cols * 6) 
    
    for c in cnts:
        (x, y, w, h) = cv2.boundingRect(c)
        ar = w / float(h)
        if w >= min_w and 0.5 <= ar <= 1.6:
            bubbles.append(c)
    
    if not bubbles: return None

    try:
        bubbles = sort_contours(bubbles, method="left-to-right")[0]
        columns = []
        current_col = []
        prev_x = -1000
        
        for c in bubbles:
            (x, y, w, h) = cv2.boundingRect(c)
            if x - prev_x > w:
                if current_col: 
                    current_col = sort_contours(current_col, method="top-to-bottom")[0]
                    columns.append(current_col)
                current_col = [c]
                prev_x = x + w/2 
            else:
                current_col.append(c)
        
        if current_col: 
            current_col = sort_contours(current_col, method="top-to-bottom")[0]
            columns.append(current_col)
            
        if len(columns) > cols: columns = columns[:cols]
        
        result_str = ""
        for col in columns:
            filled_idx = -1
            max_pixel = 0
            for i, c in enumerate(col):
                mask = np.zeros(thresh.shape, dtype="uint8")
                cv2.drawContours(mask, [c], -1, 255, -1)
                mask = cv2.bitwise_and(thresh, thresh, mask=mask)
                total = cv2.countNonZero(mask)
                if draw_on_me is not None:
                    (x, y, w, h) = cv2.boundingRect(c)
                    cv2.rectangle(draw_on_me, (x+offset[0], y+offset[1]), (x+w+offset[0], y+h+offset[1]), (0, 255, 0), 1)
                if total > max_pixel:
                    max_pixel = total
                    filled_idx = i
            if max_pixel > bubble_thresh:
                result_str += str(filled_idx) if filled_idx < 10 else "?"
                if draw_on_me is not None and filled_idx != -1 and filled_idx < len(col):
                    c = col[filled_idx]
                    (x, y, w, h) = cv2.boundingRect(c)
                    cv2.circle(draw_on_me, (int(x+w/2+offset[0]), int(y+h/2+offset[1])), int(w/2), (0, 0, 255), 3)
            else:
                result_str += "?"
        return result_str
    except Exception as e:
        logger.error(f"Sort Error: {e}")
        return None

# --- AI VISION HELPER ---
def read_handwriting_with_gemini(roi_image):
    """Dùng Gemini Flash để đọc SBD/Mã đề viết tay"""
    if not GEMINI_API_KEY: return None, None
    try:
        _, buffer = cv2.imencode('.jpg', roi_image)
        model = genai.GenerativeModel('gemini-1.5-flash')
        prompt = """
        Analyze this OMR sheet header. 
        Identify the handwritten Student ID (SBD - usually 6 digits) and Exam Code (Mã đề - usually 3 or 4 digits).
        Ignore the printed text or bubbles.
        Return JSON format: {"sbd": "...", "code": "..."}
        If you cannot find them, return empty strings.
        """
        response = model.generate_content([{'mime_type': 'image/jpeg', 'data': buffer.tobytes()}, prompt])
        text = response.text.strip()
        
        clean_text = text
        if "json" in clean_text:
            clean_text = clean_text.replace("```json", "").replace("```", "").strip()
        elif "{" in clean_text: 
             start = clean_text.find("{")
             end = clean_text.rfind("}") + 1
             if start != -1 and end != -1: clean_text = clean_text[start:end]

        data = json.loads(clean_text)
        return data.get("sbd", ""), data.get("code", "")
    except Exception as e:
        logger.error(f"AI Vision Error: {e}")
        return None, None

def process_omr(image_bytes, all_keys, use_ai=False, pre_warped=False):
    try:
        nparr = np.frombuffer(image_bytes, np.uint8)
        image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        
        if pre_warped:
            warped = image
            if warped.shape[1] != 1600: warped = imutils.resize(warped, width=1600)
        else:
            if image.shape[1] > 1600: image = imutils.resize(image, width=1600)
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
            blurred = cv2.GaussianBlur(gray, (5, 5), 0)
            edged = cv2.Canny(blurred, 75, 200)
            cnts = cv2.findContours(edged.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            cnts = imutils.grab_contours(cnts)
            docCnt = None
            if len(cnts) > 0:
                cnts = sorted(cnts, key=cv2.contourArea, reverse=True)
                for c in cnts:
                    peri = cv2.arcLength(c, True)
                    approx = cv2.approxPolyDP(c, 0.02 * peri, True)
                    if len(approx) == 4: docCnt = approx; break
            if docCnt is not None:
                warped = four_point_transform(image, docCnt.reshape(4, 2))
            else:
                warped = image

        h, w = warped.shape[:2]
        draw_img = warped.copy()

        # === 1. XỬ LÝ ID (SBD + MÃ ĐỀ) ===
        split_x = int(w * 0.36) 
        id_roi_y = int(h * 0.12) 
        left_roi = warped[id_roi_y:h-50, 0:split_x]
        
        id_result = read_bubbles(left_roi, 10, 10, draw_on_me=draw_img, offset=(0, id_roi_y), bubble_thresh=55)
        
        sbd = "AI_READ"
        code = "DEFAULT"
        
        if id_result and len(id_result) >= 6:
            sbd_omr = id_result[:6].replace("?", "")
            code_omr = ""
            if len(id_result) >= 10: code_omr = id_result[6:10].replace("?", "")
            elif len(id_result) >= 9: code_omr = id_result[6:].replace("?", "")
            sbd = sbd_omr
            code = code_omr if code_omr else "DEFAULT"
        
        # --- AI VISION LAYER ---
        if use_ai:
            header_roi = warped[0:int(h*0.4), 0:split_x]
            ai_sbd, ai_code = read_handwriting_with_gemini(header_roi)
            if ai_sbd and (len(sbd) < 6 or "?" in id_result[:6]): sbd = ai_sbd
            if ai_code and (code == "DEFAULT" or len(code) < 3): code = ai_code

        if not sbd: sbd = "UNKNOWN"
        if not code: code = "DEFAULT"

        # === 2. XỬ LÝ CÂU HỎI ===
        ans_roi_y = int(h * 0.20)
        
        avail_codes = list(all_keys.keys())
        if code not in avail_codes and avail_codes: 
            if "DEFAULT" in avail_codes: code = "DEFAULT"
            else: code = str(avail_codes[0])
            
        key_data = all_keys.get(code, {})
        total_q = len(key_data)
        num_cols_page = 4 if total_q > 30 else 3 
        col_width = (w - split_x) // num_cols_page
        
        score = 0
        correct_count = 0
        
        for c_idx in range(num_cols_page):
            c_x_start = split_x + c_idx * col_width
            c_roi = warped[ans_roi_y:h-50, c_x_start : c_x_start + col_width]
            
            gray_c = cv2.cvtColor(c_roi, cv2.COLOR_BGR2GRAY)
            thresh_c = cv2.threshold(gray_c, 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU)[1]
            cnts_c = cv2.findContours(thresh_c, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            cnts_c = imutils.grab_contours(cnts_c)
            
            bubbles_c = []
            for b in cnts_c:
                (bx, by, bw, bh) = cv2.boundingRect(b)
                if bw > 5 and bw < 100: bubbles_c.append(b)
            if not bubbles_c: continue
            try: bubbles_c = sort_contours(bubbles_c, method="top-to-bottom")[0]
            except: continue
            
            rows = []
            temp_row = []
            prev_y = -1000
            
            for b in bubbles_c:
                (bx, by, bw, bh) = cv2.boundingRect(b)
                if by - prev_y > bh * 0.8:
                    if temp_row:
                        if len(temp_row) == 4:
                            temp_row = sort_contours(temp_row, method="left-to-right")[0]
                            rows.append(temp_row)
                    temp_row = [b]
                    prev_y = by + bh/2
                else: temp_row.append(b)
            if len(temp_row) == 4: 
                temp_row = sort_contours(temp_row, method="left-to-right")[0]
                rows.append(temp_row)
                
            for r_idx, row in enumerate(rows):
                questions_per_col = 10 
                q_num = c_idx * questions_per_col + r_idx + 1
                if q_num > total_q: break
                
                filled_opt = -1
                max_p = 0
                for o_idx, b in enumerate(row):
                    mask = np.zeros(thresh_c.shape, dtype="uint8")
                    cv2.drawContours(mask, [b], -1, 255, -1)
                    mask = cv2.bitwise_and(thresh_c, thresh_c, mask=mask)
                    total = cv2.countNonZero(mask)
                    (bx, by, bw, bh) = cv2.boundingRect(b)
                    gx, gy = c_x_start + bx, ans_roi_y + by
                    cv2.rectangle(draw_img, (gx, gy), (gx+bw, gy+bh), (0, 255, 0), 1)
                    if total > max_p:
                        max_p = total
                        filled_opt = o_idx
                
                correct_char = key_data.get(q_num, key_data.get(str(q_num))) 
                if correct_char:
                    correct_idx = ord(correct_char) - 65 
                    if max_p > 55:
                        b = row[filled_opt]
                        (bx, by, bw, bh) = cv2.boundingRect(b)
                        gx, gy = c_x_start + bx, ans_roi_y + by
                        if filled_opt == correct_idx:
                            correct_count += 1
                            cv2.circle(draw_img, (gx+bw//2, gy+bh//2), bw//2, (0, 255, 0), -1)
                        else:
                            cv2.circle(draw_img, (gx+bw//2, gy+bh//2), bw//2, (0, 0, 255), -1)
                            if correct_idx < len(row):
                                b_correct = row[correct_idx]
                                (bx_c, by_c, bw_c, bh_c) = cv2.boundingRect(b_correct)
                                gx_c, gy_c = c_x_start + bx_c, ans_roi_y + by_c
                                cv2.circle(draw_img, (gx_c+bw_c//2, gy_c+bh_c//2), 5, (255, 0, 0), -1)

        if total_q > 0: score = round((correct_count / total_q) * 10, 2)

        cv2.rectangle(draw_img, (0, 0), (w, 100), (255, 255, 255), -1)
        cv2.putText(draw_img, f"SBD: {sbd} | CODE: {code}", (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0,0,0), 2)
        cv2.putText(draw_img, f"DIEM: {score} ({correct_count}/{total_q})", (20, 90), cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0,0,255), 3)

        _, buffer = cv2.imencode('.jpg', draw_img)
        b64 = base64.b64encode(buffer).decode('utf-8')
        
        return {
            "student_id": sbd, "exam_code": code, "score": score,
            "correct_count": correct_count, "total_questions": total_q,
            "wrong_count": total_q - correct_count, "image_base64": b64,
            "status": "success",
            "ai_used": use_ai
        }

    except Exception as e:
        logger.error(f"OMR Error: {str(e)}")
        return {"score": 0, "error": str(e)}

# --- DB HELPERS ---
def save_result_to_db(data, filename, session_id):
    if not supabase: return
    try:
        payload = {
            "session_id": session_id,
            "student_id": data.get("student_id"),
            "exam_code": data.get("exam_code"),
            "score": data.get("score"),
            "correct_count": data.get("correct_count"),
            "wrong_count": data.get("wrong_count"),
            "total_questions": data.get("total_questions"),
            "filename": filename
        }
        supabase.table("omr_results").insert(payload).execute()
    except Exception as e: logger.error(f"DB Save Error: {e}")

def get_key_from_db(session_id):
    if not supabase: return {}
    try:
        res = supabase.table("omr_sessions").select("answer_key_json").eq("session_id", session_id).execute()
        if res.data: return json.loads(res.data[0]['answer_key_json'])
    except: pass
    return {}

# --- API ENDPOINTS ---
@router.post("/create_session")
async def create_session(name: str = Form(...), answer_key_json: str = Form(...)):
    session_id = str(uuid.uuid4())[:8].upper()
    try: all_keys = json.loads(answer_key_json)
    except: all_keys = {}
    sessions[session_id] = {"keys": all_keys, "results": []}
    if supabase:
        try: supabase.table("omr_sessions").insert({"session_id": session_id, "session_name": name, "answer_key_json": answer_key_json}).execute()
        except Exception as e: logger.error(f"Create Session DB Error: {e}")
    return {"session_id": session_id, "name": name}

@router.post("/grade")
async def grade_exam(file: UploadFile = File(...), answer_key_json: str = Form(...), handwriting: str = Form(None), pre_warped: str = Form(None)):
    try: all_keys = json.loads(answer_key_json)
    except: return {"error": "JSON Error"}
    content = await file.read()
    use_ai = (handwriting == "true")
    is_pre_warped = (pre_warped == "true")
    result = process_omr(content, all_keys, use_ai=use_ai, pre_warped=is_pre_warped)
    save_result_to_db(result, file.filename, "Direct_Upload")
    return result

@router.post("/grade_mobile")
async def grade_mobile(session_id: str = Form(...), file: UploadFile = File(...), pre_warped: str = Form(None)):
    if session_id not in sessions:
        db_keys = get_key_from_db(session_id)
        if db_keys: sessions[session_id] = {"keys": db_keys, "results": []}
        else: return {"error": "Session not found"}
    content = await file.read()
    is_pre_warped = (pre_warped == "true")
    result = process_omr(content, sessions[session_id]["keys"], use_ai=True, pre_warped=is_pre_warped)
    result["filename"] = file.filename
    sessions[session_id]["results"].append(result)
    save_result_to_db(result, file.filename, session_id)
    return result

@router.get("/poll_results/{session_id}")
async def poll_results(session_id: str):
    if session_id not in sessions: return []
    res = sessions[session_id]["results"]
    sessions[session_id]["results"] = [] 
    return res

@router.get("/history")
def get_history():
    if not supabase: return []
    try: return supabase.table("omr_sessions").select("*").order("created_at", desc=True).execute().data
    except: return []

@router.get("/history/{session_id}")
def get_history_detail(session_id: str):
    if not supabase: return []
    try: return supabase.table("omr_results").select("*").eq("session_id", session_id).order("created_at", desc=True).execute().data
    except: return []

@router.put("/update_session/{session_id}")
async def update_session(session_id: str, name: str = Form(...)):
    if supabase:
        try:
            supabase.table("omr_sessions").update({"session_name": name}).eq("session_id", session_id).execute()
            return {"status": "updated", "name": name}
        except Exception as e: return {"error": str(e)}
    return {"status": "local_updated"}

@router.delete("/delete_session/{session_id}")
async def delete_session(session_id: str):
    if session_id in sessions: del sessions[session_id]
    if supabase:
        try:
            supabase.table("omr_results").delete().eq("session_id", session_id).execute()
            supabase.table("omr_sessions").delete().eq("session_id", session_id).execute()
            return {"status": "deleted"}
        except Exception as e: return {"error": str(e)}
    return {"status": "local_deleted"}

@router.post("/init_session")
async def init_session(answer_key_json: str = Form(...)):
    return await create_session(name="Quick Session", answer_key_json=answer_key_json)