#!/usr/bin/env python3
import argparse
import base64
import json
import math
import re
import shutil
import subprocess
import sys
import tempfile
import xml.etree.ElementTree as ET
from pathlib import Path

from PIL import Image


OPTION_LETTERS = ["A", "B", "C", "D"]


def run_command(args):
    completed = subprocess.run(args, check=False, capture_output=True)
    if completed.returncode != 0:
        stdout = completed.stdout.decode("utf-8", errors="ignore").strip()
        stderr = completed.stderr.decode("utf-8", errors="ignore").strip()
        detail = stderr or stdout or "no command output"
        raise RuntimeError(f"{args[0]} failed with exit {completed.returncode}: {detail[:1400]}")
    return completed


def run_command_optional(args):
    return subprocess.run(args, check=False, capture_output=True)


def normalize_pdf_for_poppler(pdf_path):
    if not shutil.which("qpdf"):
        return pdf_path

    normalized_path = pdf_path.with_name(f"{pdf_path.stem}-normalized.pdf")
    completed = run_command_optional([
        "qpdf",
        "--object-streams=disable",
        str(pdf_path),
        str(normalized_path),
    ])
    if completed.returncode in (0, 3) and normalized_path.exists() and normalized_path.stat().st_size > 0:
        return normalized_path
    return pdf_path


def get_pdf_page_count(pdf_path):
    info = run_command(["pdfinfo", str(pdf_path)]).stdout.decode("utf-8", errors="ignore")
    match = re.search(r"^Pages:\s+(\d+)\s*$", info, re.MULTILINE)
    if not match:
        raise RuntimeError("Could not read page count from PDF.")
    return int(match.group(1))


def extract_questions(pdf_path, settings, output_dir=None):
    pdf_path = normalize_pdf_for_poppler(pdf_path)
    page_count = get_pdf_page_count(pdf_path)
    start_page = int(settings.get("startPage", 1))
    end_page = int(settings.get("endPage", start_page))
    if start_page < 1:
        start_page = 1
    if end_page > page_count:
        end_page = page_count
    if start_page > page_count or start_page > end_page:
        raise RuntimeError(f"Invalid page range {start_page}-{end_page}. PDF has {page_count} pages.")
    first_question = int(settings.get("firstQuestion", 1))
    expected_questions = int(settings.get("expectedQuestions", 75))
    answer_mode = settings.get("answerMode", "mcq")
    marker_limit = first_question + expected_questions - 1
    dpi = 216

    markers = []
    page_sizes = {}
    for page_number in range(start_page, end_page + 1):
        bbox_xml = run_command([
            "pdftotext",
            "-bbox",
            "-f",
            str(page_number),
            "-l",
            str(page_number),
            str(pdf_path),
            "-"
        ]).stdout.decode("utf-8", errors="ignore")
        page_markers, page_size = find_question_markers(bbox_xml, page_number, first_question, marker_limit)
        page_sizes[page_number] = page_size
        markers.extend(page_markers)

    markers = normalize_markers(markers, first_question, marker_limit)
    if not markers:
        raise RuntimeError("No question numbers were detected. Try changing the start page or end page.")

    with tempfile.TemporaryDirectory() as temp_dir:
        prefix = Path(temp_dir) / "page"
        run_command([
            "pdftoppm",
            "-f",
            str(start_page),
            "-l",
            str(end_page),
            "-r",
            str(dpi),
            "-png",
            str(pdf_path),
            str(prefix)
        ])

        page_images = {}
        for page_number in sorted({marker["pageNumber"] for marker in markers}):
            page_images[page_number] = Image.open(find_rendered_page(prefix, page_number)).convert("RGB")

        questions = []
        for index, marker in enumerate(markers[:expected_questions]):
            previous_marker = markers[index - 1] if index > 0 else None
            next_marker = markers[index + 1] if index + 1 < len(markers) else None
            output_path = Path(output_dir) / f"q-{marker['questionNo']}.png" if output_dir else None
            image = crop_question(page_images[marker["pageNumber"]], marker, previous_marker, next_marker, dpi, output_path)
            question = {
                "id": f"q-{marker['questionNo']}-{marker['pageNumber']}",
                "questionNo": marker["questionNo"],
                "pageNumber": marker["pageNumber"],
                "subject": subject_for_question(marker["questionNo"], expected_questions),
                "answerType": answer_type_for_question(marker["questionNo"], answer_mode),
                "width": image["width"],
                "height": image["height"]
            }
            if "dataUrl" in image:
                question["imageUrl"] = image["dataUrl"]
            if "imagePath" in image:
                question["imagePath"] = image["imagePath"]
            questions.append(question)

    return {"questions": questions}


def find_question_markers(bbox_xml, page_number, first_question, marker_limit):
    root = ET.fromstring(bbox_xml)
    page = root.find(".//{*}page")
    if page is None:
        return [], {"width": 612.0, "height": 792.0}

    page_width = float(page.attrib.get("width", "612"))
    page_height = float(page.attrib.get("height", "792"))
    markers = []

    for word in page.findall(".//{*}word"):
        text = re.sub(r"\s+", "", word.text or "").strip()
        match = re.match(r"^(\d{1,3})(\.)?$", text)
        if not match:
            continue

        question_no = int(match.group(1))
        if question_no < first_question or question_no > marker_limit:
            continue

        x_min = float(word.attrib["xMin"])
        y_min = float(word.attrib["yMin"])
        has_dot = bool(match.group(2))
        question_column_left = page_width * 0.055
        question_column_right = page_width * (0.12 if has_dot else 0.095)
        inside_question_column = question_column_left <= x_min <= question_column_right
        inside_content_band = page_height * 0.05 <= y_min <= page_height * 0.93
        if not inside_question_column or not inside_content_band:
            continue

        markers.append({
            "questionNo": question_no,
            "pageNumber": page_number,
            "x": x_min,
            "y": y_min,
            "pageWidth": page_width,
            "pageHeight": page_height
        })

    deduped = []
    for marker in sorted(markers, key=lambda item: item["y"]):
        previous = deduped[-1] if deduped else None
        if previous and previous["questionNo"] == marker["questionNo"] and abs(previous["y"] - marker["y"]) <= 7:
            continue
        deduped.append(marker)

    return deduped, {"width": page_width, "height": page_height}


def normalize_markers(raw_markers, first_question, marker_limit):
    accepted = []
    last_question = first_question - 1
    for marker in sorted(raw_markers, key=lambda item: (item["pageNumber"], item["y"])):
        if marker["questionNo"] <= last_question or marker["questionNo"] > marker_limit:
            continue
        accepted.append(marker)
        last_question = marker["questionNo"]
    return accepted


def find_rendered_page(prefix, page_number):
    candidates = [
        Path(f"{prefix}-{page_number}.png"),
        Path(f"{prefix}-{page_number:02d}.png"),
        Path(f"{prefix}-{page_number:03d}.png"),
    ]
    for candidate in candidates:
        if candidate.exists():
            return candidate
    raise FileNotFoundError(f"Rendered page image not found for page {page_number}")


def crop_question(page_image, marker, previous_marker, next_marker, dpi, output_path=None):
    pixels_per_point = dpi / 72
    width, height = page_image.size
    top_content = int(4 * pixels_per_point)
    bottom_content = int(height - 18 * pixels_per_point)
    top_padding = int(25 * pixels_per_point)
    bottom_gap = int(6 * pixels_per_point)
    marker_y = int(marker["y"] * pixels_per_point)
    next_marker_y = int(next_marker["y"] * pixels_per_point) if next_marker and next_marker["pageNumber"] == marker["pageNumber"] else None

    # Question markers are restricted to the real left-column "1." style labels,
    # so the crop can safely end just before the next marker.
    adjusted_top = clamp(marker_y - top_padding, top_content, height - 20)
    raw_bottom = next_marker_y - bottom_gap if next_marker_y is not None else bottom_content
    adjusted_bottom = clamp(raw_bottom, adjusted_top + int(110 * pixels_per_point), bottom_content)

    crop = page_image.crop((0, adjusted_top, width, adjusted_bottom))
    if output_path:
        output_path.parent.mkdir(parents=True, exist_ok=True)
        crop.save(output_path, format="PNG", compress_level=1)
        return {
            "imagePath": str(output_path),
            "width": crop.width,
            "height": crop.height
        }

    with tempfile.NamedTemporaryFile(suffix=".png") as temp_file:
        crop.save(temp_file.name, format="PNG", compress_level=1)
        encoded = base64.b64encode(Path(temp_file.name).read_bytes()).decode("ascii")

    return {
        "dataUrl": f"data:image/png;base64,{encoded}",
        "width": crop.width,
        "height": crop.height
    }


def find_top_boundary_before_question(image, previous_marker_y, current_marker_y, pixels_per_point):
    width, height = image.size
    left = int(10 * pixels_per_point)
    right = width - left
    start_y = clamp(previous_marker_y + int(55 * pixels_per_point), 0, height - 1)
    end_y = clamp(current_marker_y - int(3 * pixels_per_point), start_y + 1, height)
    if end_y <= start_y + int(16 * pixels_per_point):
        return current_marker_y - int(18 * pixels_per_point)

    minimum_ink_pixels = max(8, int((right - left) * 0.002))
    pixels = image.load()
    blank_run = 0
    required_blank = max(8, int(3 * pixels_per_point))
    seen_current_ink = False

    for y in range(end_y - 1, start_y - 1, -1):
        ink_pixels = 0
        for x in range(left, right):
            red, green, blue = pixels[x, y]
            if red < 170 or green < 170 or blue < 170:
                ink_pixels += 1
                if ink_pixels >= minimum_ink_pixels:
                    break

        if ink_pixels >= minimum_ink_pixels:
            seen_current_ink = True
            blank_run = 0
        elif seen_current_ink:
            blank_run += 1
            if blank_run >= required_blank:
                return clamp(
                    y + blank_run // 2,
                    previous_marker_y + int(55 * pixels_per_point),
                    current_marker_y - int(6 * pixels_per_point)
                )

    return clamp(
        current_marker_y - int(18 * pixels_per_point),
        previous_marker_y + int(55 * pixels_per_point),
        current_marker_y - int(6 * pixels_per_point)
    )


def find_bottom_boundary_after_question(image, current_marker_y, next_marker_y, pixels_per_point):
    width, height = image.size
    left = int(60 * pixels_per_point)
    right = width - left
    start_y = clamp(current_marker_y + int(55 * pixels_per_point), 0, height - 1)
    end_y = clamp(next_marker_y - int(3 * pixels_per_point), start_y + 1, height)
    if end_y <= start_y + int(16 * pixels_per_point):
        return next_marker_y - int(10 * pixels_per_point)

    minimum_ink_pixels = max(8, int((right - left) * 0.002))
    row_has_ink = []
    pixels = image.load()
    for y in range(start_y, end_y):
        ink_pixels = 0
        for x in range(left, right):
            red, green, blue = pixels[x, y]
            if red < 90 and green < 90 and blue < 90:
                ink_pixels += 1
                if ink_pixels >= minimum_ink_pixels:
                    break
        row_has_ink.append(ink_pixels >= minimum_ink_pixels)

    blank_run = 0
    required_blank = max(8, int(3 * pixels_per_point))
    seen_ink = False
    for row, has_ink in enumerate(row_has_ink):
        if has_ink:
            seen_ink = True
            blank_run = 0
        elif seen_ink:
            blank_run += 1
            if blank_run >= required_blank:
                return clamp(
                    start_y + row - blank_run // 2,
                    current_marker_y + int(55 * pixels_per_point),
                    next_marker_y - int(6 * pixels_per_point)
                )

    return clamp(
        next_marker_y - int(10 * pixels_per_point),
        current_marker_y + int(55 * pixels_per_point),
        next_marker_y - int(6 * pixels_per_point)
    )


def extract_key(pdf_path, questions):
    pdf_path = normalize_pdf_for_poppler(pdf_path)
    page_count = get_pdf_page_count(pdf_path)
    text = run_command([
        "pdftotext",
        "-layout",
        "-f",
        "1",
        "-l",
        str(min(2, page_count)),
        str(pdf_path),
        "-"
    ]).stdout.decode("utf-8", errors="ignore")
    key_section = isolate_key_section(text)
    tokens = re.findall(r"-?\d+(?:\.\d+)?", key_section)
    question_by_no = {int(question["questionNo"]): question for question in questions}
    answers = {}

    for index in range(0, len(tokens) - 1, 2):
        question_no = int(float(tokens[index]))
        answer = tokens[index + 1]
        question = question_by_no.get(question_no)
        if not question:
            continue
        answers[str(question_no)] = normalize_key_value(answer, question.get("answerType", "mcq"))

    return {
        "correctAnswers": answers,
        "answerKeyText": "\n".join(f"{key} {answers[key]}" for key in sorted(answers, key=lambda item: int(item)))
    }


def isolate_key_section(text):
    lower = text.lower()
    start_index = lower.find("key sheet")
    start = start_index + len("key sheet") if start_index >= 0 else 0
    end_candidates = []
    for needle in [" solution", "\n sec:", "\nsec:", " page 1"]:
        index = lower.find(needle, start)
        if index >= 0:
            end_candidates.append(index)
    end = min(end_candidates) if end_candidates else len(text)
    return text[start:end]


def normalize_key_value(value, answer_type):
    stripped = str(value).strip()
    if answer_type == "mcq":
        if re.match(r"^[1-4]$", stripped):
            return OPTION_LETTERS[int(stripped) - 1]
        return stripped.upper()
    return stripped


def subject_for_question(question_no, expected_questions):
    if expected_questions >= 75:
        if question_no <= 25:
            return "Mathematics"
        if question_no <= 50:
            return "Physics"
        return "Chemistry"
    return "Paper"


def answer_type_for_question(question_no, answer_mode):
    if answer_mode == "mcq":
        return "mcq"
    position = ((question_no - 1) % 25) + 1
    return "numeric" if position > 20 else "mcq"


def clamp(value, minimum, maximum):
    return min(max(int(value), int(minimum)), int(maximum))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("mode", choices=["questions", "key"])
    parser.add_argument("--pdf", required=True)
    parser.add_argument("--settings-json", default="{}")
    parser.add_argument("--questions-json", default="[]")
    parser.add_argument("--output-dir")
    args = parser.parse_args()

    try:
      if args.mode == "questions":
          print(json.dumps(extract_questions(Path(args.pdf), json.loads(args.settings_json), args.output_dir)))
      else:
          print(json.dumps(extract_key(Path(args.pdf), json.loads(args.questions_json))))
    except Exception as error:
        print(json.dumps({"error": str(error)}))
        sys.exit(1)


if __name__ == "__main__":
    main()