"""
autograder.py

Given a multi-page PDF (one answer sheet per page):

  - Convert each page to an image
  - Detect fiducials, warp, read bubbles + QR version label
  - Look up the correct answers for that version from a CSV
  - Grade each page

The answer-key CSV should be generated by make_exams.py and have columns:

  version_label, Q1, Q2, Q3, ...

Example row:

  01234567,a,b,c,d,a,...

Dependencies:
  - pdf2image
  - answer_sheet_reader.LayoutConfig
  - answer_sheet_reader.read_answer_sheet_with_version
"""

from __future__ import annotations

from pathlib import Path

from typing import Dict, List, Optional, Tuple, Any

import argparse
import os
import csv
import tempfile
import logging
import numpy as np

from pdf2image import convert_from_path

from .answersheetreader import AnswerSheetReader, LayoutConfig

logger = logging.getLogger(__name__)

class Autograder:
    """
    Autograder for multi-page PDF answer sheets.
    """
    def __init__(self, layout_config: LayoutConfig):
        self.layout = layout_config
        self.version_keys: Dict[str, List[str]] = {}

    def load_version_keys_csv(self, keys_csv_path: Path | str):
        if not os.path.exists(keys_csv_path):
            raise FileNotFoundError(f"Answer-key CSV not found: {keys_csv_path}")
        current_number_of_keys = len(self.version_keys)
        with open(keys_csv_path, newline="", encoding="utf-8") as f:
            reader = csv.DictReader(f)
            fieldnames = reader.fieldnames or []
            # Identify question columns and sort them by number
            qcols = [col for col in fieldnames if col.startswith("Q")]
            qcols.sort(key=lambda name: int(name[1:]))  # "Q10" -> 10

            for row in reader:
                version_label = (row.get("version_label") or "").strip()
                if not version_label:
                    continue

                answers: List[str] = []
                for col in qcols:
                    val = row.get(col)
                    if val == 'True':
                        answers.append('a')
                    elif val == 'False':
                        answers.append('b')
                    else:
                        answers.append(val.lower() if val else "")
                self.version_keys[version_label] = answers
        new_number_of_keys = len(self.version_keys)
        logger.info(f"Loaded {new_number_of_keys - current_number_of_keys} new version keys from {keys_csv_path}")
        logger.info(f"Total version keys loaded: {new_number_of_keys}")
        
    def grade_pdf(self, 
        pdf_file_path: Path | str, 
        output_dir_path: [Path | str] = Path.cwd(),
        gradesheet_output_csv_path: Optional[Path | str] = None,
        question_tally_csv_path: Optional[Path | str] = None,
        debug_output_dir_path: Optional[Path | str] = None):
        """
        Grade the given PDF of answer sheets.
        
        Parameters
        ----------
        pdf_file_path : Path | str
            Path to the input PDF file.
        output_dir_path : [Path | str]
            Path of directory in which to save individual graded images and the output CSV report. If None, no file is written.
        debug_output_dir_path : Optional[Path | str]
            Directory to save debug output images. If None, no debug images are saved.
        """
        if isinstance(pdf_file_path, str):
            pdf_file_path = Path(pdf_file_path)
        if not pdf_file_path.exists():
            raise FileNotFoundError(f"PDF not found: {pdf_file_path}")

        if debug_output_dir_path is not None:
            debug_output_dir_path = Path(debug_output_dir_path)
            if not debug_output_dir_path.exists():
                debug_output_dir_path.mkdir(parents=True, exist_ok=True)

        pages = convert_from_path(str(pdf_file_path), dpi=300, fmt='png')

        results: List[Dict[str, Any]] = []

        for page_index, page in enumerate(pages, start=1):
            img = np.array(page)
            reader = AnswerSheetReader(img, layout_config=self.layout, debug_output_dir=debug_output_dir_path)
            read_results = reader.read()

            page_result: Dict[str, Any] = {
                "page_index": page_index,
                "version_label": read_results['version'],
                "student_id": read_results['student_id_bubbles'],
                "answers_detected": read_results['answers'],
                "answer_key": None,
                "per_question": {},
                "num_questions": None,
                "num_correct": None,
                "score_fraction": None,
                "status": "ok",
                "overlay_path": output_dir_path / f"graded_page_{page_index:03d}.png" if output_dir_path is not None else None,
            }

            key = self.version_keys.get(read_results['version']) if read_results['version'] else None
            logger.debug(f"Grading page {page_index}: version={read_results['version']}, student_id={read_results['student_id_bubbles']}, key_found={key is not None}")
            logger.debug(f'key={key}, detected_answers={read_results["answers"]}')
            if not read_results['version']:
                page_result["status"] = "no_qr"
                results.append(page_result)
                continue
            if key is None:
                page_result["status"] = "unknown_version"
                results.append(page_result)
                continue

            page_result["answer_key"] = key
            num_q = len(key)
            page_result["num_questions"] = num_q

            num_correct = 0
            per_q: Dict[int, Dict[str, Any]] = {}
            for qnum in range(1, num_q + 1):
                correct = (key[qnum - 1] or "").lower()
                detected = read_results['answers'].get(qnum)
                detected_norm = (detected or "").lower() if detected else ""
                is_correct = bool(correct) and (detected_norm == correct)
                if is_correct:
                    num_correct += 1
                per_q[qnum] = {
                    "correct": correct or None,
                    "detected": detected_norm or None,
                    "is_correct": is_correct,
                }
                logger.debug(f' Q{qnum}: correct="{correct}", detected="{detected_norm}", is_correct={is_correct}')

            page_result["per_question"] = per_q
            page_result["num_correct"] = num_correct
            page_result["score_fraction"] = (
                float(num_correct) / num_q if num_q > 0 else None
            )

            logger.debug(f' Graded page {page_index}: num_correct={num_correct}/{num_q}, score_fraction={page_result["score_fraction"]}')
            logger.debug(f' Writing graded overlay image for page {page_index} to {output_dir_path / f"graded_{read_results["student_id_bubbles"]}_{read_results["version"]}.png"}...')
            if output_dir_path is not None:
                reader.write_graded_annotations(
                    per_question_results=per_q,
                    score_fraction=page_result["score_fraction"],
                    overlay_path=output_dir_path / f"graded_{read_results['student_id_bubbles']}_{read_results['version']}.png",
                )

            results.append(page_result)

        if gradesheet_output_csv_path is not None:
            gradesheet_output_csv_path = Path(gradesheet_output_csv_path)
            # if this file exists, read it in and append to it if this ID-version combo is not already present
            if gradesheet_output_csv_path.exists():
                existing_rows: List[Dict[str, Any]] = []
                with open(gradesheet_output_csv_path, "r", newline="", encoding="utf-8") as f:
                    reader = csv.DictReader(f)
                    for row in reader:
                        existing_rows.append(row)
                existing_combos = set((row["student_id"], row["version_label"]) for row in existing_rows)
                new_results = []
                for res in results:
                    combo = (res["student_id"], res["version_label"])
                    if combo not in existing_combos:
                        new_results.append(res)
                results = existing_rows + new_results
            if len(new_results) > 0:
                with open(gradesheet_output_csv_path, "w", newline="", encoding="utf-8") as f:
                    fieldnames = [
                        "page_index",
                        "version_label",
                        "student_id",
                        "num_questions",
                        "num_correct",
                        "score",
                        "status",
                    ]
                    writer = csv.DictWriter(f, fieldnames=fieldnames)
                    writer.writeheader()
                    for res in results:
                        writer.writerow({
                            "page_index": res["page_index"],
                            "version_label": res["version_label"],
                            "student_id": res["student_id"],
                            "num_questions": res["num_questions"],
                            "num_correct": res["num_correct"],
                            "score": f'{res["score_fraction"]*100:.1f}',
                            "status": res["status"],
                        })
                logger.info(f"Wrote grading results CSV to {gradesheet_output_csv_path}")

        if question_tally_csv_path is not None:
            # for each student ID, write a row that contains the exam version and their answers
            question_tally_csv_path = Path(question_tally_csv_path)
            if question_tally_csv_path.exists():
                # read existing rows and append only new ones
                existing_rows: List[Dict[str, Any]] = []
                with open(question_tally_csv_path, "r", newline="", encoding="utf-8") as f:
                    reader = csv.DictReader(f)
                    for row in reader:
                        existing_rows.append(row)
                existing_combos = set((row["student_id"], row["version_label"]) for row in existing_rows)
                new_results = []
                for res in results:
                    combo = (res["student_id"], res["version_label"])
                    if combo not in existing_combos:
                        new_results.append(res)
                results = existing_rows + new_results
            if len(new_results) > 0:
                with open(question_tally_csv_path, "w", newline="", encoding="utf-8") as f:
                    fieldnames = [
                        "student_id",
                        "version_label",
                    ] + [f"Q{qnum}" for qnum in range(1, max(len(res["answer_key"] or []) for res in results) + 1)]
                    writer = csv.DictWriter(f, fieldnames=fieldnames)
                    writer.writeheader()
                    for res in results:
                        row = {
                            "student_id": res["student_id"],
                            "version_label": res["version_label"],
                        }
                        answers = res["answers_detected"]
                        for qnum in range(1, len(res["answer_key"] or []) + 1):
                            row[f"Q{qnum}"] = answers.get(qnum)
                        writer.writerow(row)
                logger.info(f"Wrote question tally CSV to {question_tally_csv_path}")