"""Form-independent preprocessing and OMRChecker wrapper for TDC forms."""

from __future__ import annotations

import argparse
import csv
import json
import os
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path
from typing import Any

import cv2
import fitz
import numpy as np
from PIL import Image, ImageOps

try:
    from pillow_heif import register_heif_opener

    register_heif_opener()
    HEIC_AVAILABLE = True
except ImportError:
    HEIC_AVAILABLE = False


FULL_SIZE = (1700, 2200)
PDF_SIZE = (612.0, 792.0)
BORDER_PDF = (27.0, 27.0, 585.0, 765.0)
MARKERS_PDF = (39.0, 39.0, 573.0, 753.0)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--form-id", required=True)
    parser.add_argument("--input", type=Path, required=True)
    parser.add_argument("--bundle-root", type=Path, required=True)
    parser.add_argument("--omrchecker-root", type=Path, required=True)
    parser.add_argument("--output-json", type=Path)
    parser.add_argument("--artifacts-dir", type=Path)
    return parser.parse_args()


def load_image(path: Path) -> np.ndarray:
    if path.suffix.lower() == ".pdf":
        document = fitz.open(path)
        if len(document) != 1:
            raise ValueError(f"Expected a one-page PDF; received {len(document)} pages")
        page = document[0]
        pixmap = page.get_pixmap(matrix=fitz.Matrix(2.5, 2.5), alpha=False)
        image = np.frombuffer(pixmap.samples, dtype=np.uint8).reshape(pixmap.height, pixmap.width, pixmap.n)
        return cv2.cvtColor(image, cv2.COLOR_RGB2BGR)

    if path.suffix.lower() in {".heic", ".heif"} and not HEIC_AVAILABLE:
        raise RuntimeError("HEIC input requires the pillow-heif package")
    with Image.open(path) as source:
        corrected = ImageOps.exif_transpose(source).convert("RGB")
        return cv2.cvtColor(np.asarray(corrected), cv2.COLOR_RGB2BGR)


def order_points(points: np.ndarray) -> np.ndarray:
    points = points.astype("float32")
    ordered = np.zeros((4, 2), dtype="float32")
    sums = points.sum(axis=1)
    differences = np.diff(points, axis=1).reshape(-1)
    ordered[0] = points[np.argmin(sums)]
    ordered[2] = points[np.argmax(sums)]
    ordered[1] = points[np.argmin(differences)]
    ordered[3] = points[np.argmax(differences)]
    return ordered


def find_borders(image: np.ndarray, maximum: int = 8) -> list[np.ndarray]:
    height, width = image.shape[:2]
    scale = min(1.0, 1600.0 / max(height, width))
    small = cv2.resize(image, None, fx=scale, fy=scale, interpolation=cv2.INTER_AREA)
    gray = cv2.cvtColor(small, cv2.COLOR_BGR2GRAY)
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
    edges = cv2.Canny(blurred, 30, 110)
    edges = cv2.dilate(edges, np.ones((3, 3), np.uint8), iterations=1)
    contours, _ = cv2.findContours(edges, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
    area_total = small.shape[0] * small.shape[1]
    target_ratio = (BORDER_PDF[2] - BORDER_PDF[0]) / (BORDER_PDF[3] - BORDER_PDF[1])
    candidates: list[tuple[float, np.ndarray]] = []
    for contour in contours:
        area = cv2.contourArea(contour)
        if area < area_total * 0.18:
            continue
        perimeter = cv2.arcLength(contour, True)
        for epsilon in (0.008, 0.012, 0.018, 0.025, 0.035):
            approximation = cv2.approxPolyDP(contour, epsilon * perimeter, True)
            if len(approximation) != 4 or not cv2.isContourConvex(approximation):
                continue
            points = order_points(approximation.reshape(4, 2))
            top = np.linalg.norm(points[1] - points[0])
            bottom = np.linalg.norm(points[2] - points[3])
            left = np.linalg.norm(points[3] - points[0])
            right = np.linalg.norm(points[2] - points[1])
            ratio = ((top + bottom) / 2) / max((left + right) / 2, 1)
            ratio_error = abs(ratio - target_ratio) / target_ratio
            # A phone photo may contain both the white paper edge and the
            # slightly smaller printed black form border. Prefer the latter by
            # measuring darkness directly along the quadrilateral, rather than
            # automatically choosing the largest page-shaped contour.
            line_mask = np.zeros(gray.shape, dtype=np.uint8)
            cv2.polylines(
                line_mask,
                [points.astype(np.int32)],
                True,
                255,
                thickness=max(3, round(min(small.shape[:2]) / 350)),
            )
            line_mean = float(cv2.mean(gray, mask=line_mask)[0])
            line_darkness = 1.0 - line_mean / 255.0
            score = 1.3 * line_darkness - 0.6 * ratio_error + 0.1 * (area / area_total)
            candidates.append((score, points / scale))
            break
    selected: list[np.ndarray] = []
    for _score, points in sorted(candidates, key=lambda item: item[0], reverse=True):
        if any(float(np.mean(np.linalg.norm(order_points(points) - order_points(existing), axis=1))) < 12 for existing in selected):
            continue
        selected.append(points)
        if len(selected) == maximum:
            break
    return selected


def rectify_to_pdf_quad(image: np.ndarray, points: np.ndarray, pdf_quad: tuple[float, float, float, float]) -> np.ndarray:
    width, height = FULL_SIZE
    x0 = pdf_quad[0] / PDF_SIZE[0] * width
    y0 = pdf_quad[1] / PDF_SIZE[1] * height
    x1 = pdf_quad[2] / PDF_SIZE[0] * width
    y1 = pdf_quad[3] / PDF_SIZE[1] * height
    destination = np.array([[x0, y0], [x1, y0], [x1, y1], [x0, y1]], dtype="float32")
    transform = cv2.getPerspectiveTransform(order_points(points), destination)
    return cv2.warpPerspective(image, transform, FULL_SIZE, borderValue=(255, 255, 255))


def rectify(image: np.ndarray, border: np.ndarray) -> np.ndarray:
    return rectify_to_pdf_quad(image, border, BORDER_PDF)


def find_marker_quad(image: np.ndarray, marker: np.ndarray) -> tuple[np.ndarray, dict[str, float]] | None:
    """Locate the four bullseye fiducials, including partially clipped ones."""
    height, width = image.shape[:2]
    scale = min(1.0, 1600.0 / max(height, width))
    small = cv2.resize(image, None, fx=scale, fy=scale, interpolation=cv2.INTER_AREA)
    gray = cv2.cvtColor(small, cv2.COLOR_BGR2GRAY)
    sh, sw = gray.shape
    padding = round(max(sh, sw) * 0.055)
    padded = cv2.copyMakeBorder(gray, padding, padding, padding, padding, cv2.BORDER_CONSTANT, value=255)
    marker_gray = marker if marker.ndim == 2 else cv2.cvtColor(marker, cv2.COLOR_BGR2GRAY)
    regions = {
        "upper_left": (-padding, sw * 0.42, -padding, sh * 0.36),
        "upper_right": (sw * 0.58, sw + padding, -padding, sh * 0.36),
        "lower_right": (sw * 0.58, sw + padding, sh * 0.64, sh + padding),
        "lower_left": (-padding, sw * 0.42, sh * 0.64, sh + padding),
    }
    centers = []
    scores: dict[str, float] = {}
    minimum = max(24, round(min(sh, sw) * 0.020))
    maximum = max(minimum + 4, round(min(sh, sw) * 0.085))
    for name, (x0f, x1f, y0f, y1f) in regions.items():
        x0, x1 = round(x0f + padding), round(x1f + padding)
        y0, y1 = round(y0f + padding), round(y1f + padding)
        roi = padded[y0:y1, x0:x1]
        best_score = -1.0
        best_center = None
        for size in range(minimum, maximum + 1, 4):
            resized_marker = cv2.resize(marker_gray, (size, size), interpolation=cv2.INTER_AREA)
            response = cv2.matchTemplate(roi, resized_marker, cv2.TM_CCOEFF_NORMED)
            _, score, _, location = cv2.minMaxLoc(response)
            if score > best_score:
                best_score = float(score)
                best_center = (
                    x0 + location[0] + size / 2 - padding,
                    y0 + location[1] + size / 2 - padding,
                )
        if best_center is None or best_score < 0.46:
            return None
        centers.append(best_center)
        scores[name] = round(best_score, 3)
    points = np.asarray(centers, dtype="float32") / scale
    if cv2.contourArea(order_points(points)) < height * width * 0.28:
        return None
    return points, scores


def disk_mean(gray: np.ndarray, center: tuple[int, int], radius: int) -> float:
    mask = np.zeros(gray.shape, dtype=np.uint8)
    cv2.circle(mask, center, radius, 255, -1)
    return float(cv2.mean(gray, mask=mask)[0])


def control_measurements(rectified: np.ndarray, geometry: dict) -> dict[str, Any]:
    gray = cv2.cvtColor(rectified, cv2.COLOR_BGR2GRAY) if rectified.ndim == 3 else rectified
    scale_x = FULL_SIZE[0] / PDF_SIZE[0]
    scale_y = FULL_SIZE[1] / PDF_SIZE[1]
    samples = []
    expected = []
    names = []
    for station_name, station in geometry["control_marks"]["stations"].items():
        for index, ((x, y), bit) in enumerate(zip(station["pdf_top_origin_centers"], station["pattern"])):
            center = (round(x * scale_x), round(y * scale_y))
            mean = disk_mean(gray, center, 11)
            samples.append(mean)
            expected.append(int(bit))
            names.append(f"{station_name}_{index + 1}")

    darkest_indices = set(np.argsort(samples)[: sum(expected)].tolist())
    predicted = [1 if index in darkest_indices else 0 for index in range(len(samples))]
    matches = sum(a == b for a, b in zip(predicted, expected))
    black_means = [value for value, bit in zip(samples, expected) if bit]
    white_means = [value for value, bit in zip(samples, expected) if not bit]
    separation = float(np.mean(white_means) - np.mean(black_means))
    darkest_expected_black = float(max(black_means))
    darkest_expected_white = float(min(white_means))
    worst_case_gap = darkest_expected_white - darkest_expected_black
    contrast_valid = (
        darkest_expected_black <= 180
        and darkest_expected_white >= 100
        and worst_case_gap >= 25
    )
    exact = matches == len(expected) and contrast_valid
    return {
        "matches": matches,
        "total": len(expected),
        "rank_pattern_exact": matches == len(expected),
        "contrast_valid": contrast_valid,
        "exact": exact,
        "separation_gray_levels": round(separation, 2),
        "worst_expected_black_mean": round(darkest_expected_black, 2),
        "worst_expected_white_mean": round(darkest_expected_white, 2),
        "worst_case_gap": round(worst_case_gap, 2),
        "samples": {name: round(value, 2) for name, value in zip(names, samples)},
        "predicted_bits": predicted,
        "expected_bits": expected,
    }


def orientation_candidates(image: np.ndarray, geometry: dict, marker: np.ndarray) -> list[dict[str, Any]]:
    rotations = [
        (0, image),
        (90, cv2.rotate(image, cv2.ROTATE_90_CLOCKWISE)),
        (180, cv2.rotate(image, cv2.ROTATE_180)),
        (270, cv2.rotate(image, cv2.ROTATE_90_COUNTERCLOCKWISE)),
    ]
    candidates = []
    for degrees, rotated in rotations:
        if rotated.shape[0] <= rotated.shape[1]:
            continue
        for border_index, border in enumerate(find_borders(rotated)):
            normalized = rectify(rotated, border)
            controls = control_measurements(normalized, geometry)
            candidates.append(
                {
                    "rotation_degrees_clockwise": degrees,
                    "border_candidate": border_index,
                    "image": normalized,
                    "controls": controls,
                    "alignment_method": "printed_border",
                }
            )
        marker_match = find_marker_quad(rotated, marker)
        if marker_match:
            marker_points, marker_scores = marker_match
            normalized = rectify_to_pdf_quad(rotated, marker_points, MARKERS_PDF)
            controls = control_measurements(normalized, geometry)
            candidates.append(
                {
                    "rotation_degrees_clockwise": degrees,
                    "border_candidate": None,
                    "image": normalized,
                    "controls": controls,
                    "alignment_method": "corner_fiducials",
                    "marker_match_scores": marker_scores,
                }
            )
    return candidates


def illumination_metrics(image: np.ndarray) -> dict[str, float]:
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    background = cv2.GaussianBlur(gray, (0, 0), 75)
    central = background[100:-100, 100:-100]
    p10, p90 = np.percentile(central, [10, 90])
    return {
        "background_p10": round(float(p10), 2),
        "background_p90": round(float(p90), 2),
        "background_range": round(float(p90 - p10), 2),
        "laplacian_variance": round(float(cv2.Laplacian(gray, cv2.CV_64F).var()), 2),
    }


def normalize_illumination(image: np.ndarray, metrics: dict[str, float]) -> tuple[np.ndarray, str]:
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    if metrics["background_range"] < 28:
        return gray, "grayscale"
    background = cv2.GaussianBlur(gray, (0, 0), 70)
    return cv2.divide(gray, background, scale=255), "gaussian_background_divide"


def run_omr(image: np.ndarray, form_dir: Path, omrchecker_root: Path) -> dict[str, str]:
    with tempfile.TemporaryDirectory(prefix="tdc-omr-wrapper-") as temporary:
        root = Path(temporary)
        inputs = root / "inputs"
        outputs = root / "outputs"
        inputs.mkdir()
        shutil.copy2(form_dir / "template.json", inputs / "template.json")
        shutil.copy2(form_dir / "omr_marker.jpg", inputs / "omr_marker.jpg")
        cv2.imwrite(str(inputs / "normalized.png"), image)
        environment = os.environ.copy()
        environment["MPLCONFIGDIR"] = str(root / "matplotlib")
        completed = subprocess.run(
            [sys.executable, str(omrchecker_root / "main.py"), "-i", str(inputs), "-o", str(outputs)],
            cwd=omrchecker_root,
            env=environment,
            text=True,
            capture_output=True,
        )
        if completed.returncode:
            raise RuntimeError(f"OMRChecker exited {completed.returncode}: {completed.stderr[-2000:]}")
        csv_files = list((outputs / "Results").glob("Results_*.csv"))
        if len(csv_files) != 1:
            raise RuntimeError(f"Expected one OMR results CSV; found {len(csv_files)}")
        with csv_files[0].open(newline="", encoding="utf-8-sig") as stream:
            rows = list(csv.DictReader(stream))
        if len(rows) != 1:
            raise RuntimeError(f"Expected one OMR result; found {len(rows)}")
        return rows[0]


def decode_driver_id(raw: str, mapping: dict) -> str:
    if len(raw) != 3 or not raw.isdigit():
        raise ValueError(f"Driver ID row output must contain three digits; received {raw!r}")
    characters = []
    for position, row_character in enumerate(raw):
        value = mapping["positions"][position][int(row_character)]
        if value is None:
            raise ValueError(f"Driver ID position {position + 1} used an unavailable row")
        characters.append(value)
    return "".join(characters)


def marker_crop(image: np.ndarray) -> np.ndarray:
    x0 = round(MARKERS_PDF[0] / PDF_SIZE[0] * FULL_SIZE[0])
    y0 = round(MARKERS_PDF[1] / PDF_SIZE[1] * FULL_SIZE[1])
    x1 = round(MARKERS_PDF[2] / PDF_SIZE[0] * FULL_SIZE[0])
    y1 = round(MARKERS_PDF[3] / PDF_SIZE[1] * FULL_SIZE[1])
    return cv2.resize(image[y0:y1, x0:x1], FULL_SIZE, interpolation=cv2.INTER_LINEAR)


def field_measurements(image: np.ndarray, template: dict) -> dict[str, dict[str, Any]]:
    aligned = marker_crop(image)
    gray = aligned if aligned.ndim == 2 else cv2.cvtColor(aligned, cv2.COLOR_BGR2GRAY)
    default_width, default_height = template["bubbleDimensions"]
    measurements: dict[str, dict[str, Any]] = {}
    for block in template["fieldBlocks"].values():
        values = block["bubbleValues"]
        vertical = block["direction"] == "vertical"
        width, height = block.get("bubbleDimensions", [default_width, default_height])
        radius = max(6, round(min(width, height) * 0.24))
        for field_index, field in enumerate(block["fieldLabels"]):
            if not field.startswith("q"):
                continue
            scores = []
            for value_index, value in enumerate(values):
                x = block["origin"][0] + width / 2
                y = block["origin"][1] + height / 2
                if vertical:
                    x += field_index * block["labelsGap"]
                    y += value_index * block["bubblesGap"]
                else:
                    x += value_index * block["bubblesGap"]
                    y += field_index * block["labelsGap"]
                darkness = 1.0 - disk_mean(gray, (round(x), round(y)), radius) / 255.0
                scores.append((value, darkness))
            ranked = sorted(scores, key=lambda item: item[1], reverse=True)
            top_gap = ranked[0][1] - ranked[1][1]
            second_cluster_gap = ranked[1][1] - ranked[2][1] if len(ranked) > 2 else 0.0
            # A genuine double mark has two bubbles separated from all remaining
            # bubbles. Broad shadow/bleed-through darkens several alternatives
            # together and therefore does not produce this two-bubble cluster.
            if top_gap <= 0.16 and second_cluster_gap >= 0.13:
                selected = {ranked[0][0], ranked[1][0]}
                local_result = "".join(value for value, _score in scores if value in selected)
                local_class = "multiple"
            elif top_gap >= 0.055:
                local_result = ranked[0][0]
                local_class = "single"
            else:
                local_result = ""
                local_class = "indeterminate"
            measurements[field] = {
                "local_result": local_result,
                "local_class": local_class,
                "top_darkness": round(float(ranked[0][1]), 3),
                "top_gap": round(float(top_gap), 3),
                "second_cluster_gap": round(float(second_cluster_gap), 3),
                "scores": {value: round(float(score), 3) for value, score in scores},
            }
    return measurements


def reconcile_answers(
    omr_row: dict[str, str],
    measurements: dict[str, dict[str, Any]],
    profile: dict,
) -> tuple[dict[str, str], list[dict[str, Any]]]:
    answers: dict[str, str] = {}
    issues: list[dict[str, Any]] = []
    optional = {int(profile["bonus_question"])} if profile.get("bonus_question") else set()
    for number in range(1, profile["questions"] + 1):
        field = f"q{number}"
        omr_result = omr_row.get(field, "")
        local = measurements.get(field, {})
        local_result = local.get("local_result", "")
        local_class = local.get("local_class", "indeterminate")
        final_result = omr_result
        reason = None

        if number in optional and local_class == "indeterminate":
            final_result = ""
            if omr_result:
                reason = "optional question treated as blank because all bubbles are equally dark"
        elif local_class == "multiple" and omr_result != local_result:
            final_result = local_result
            reason = "secondary measurement detected multiple marks"
        elif local_class == "single" and omr_result != local_result:
            final_result = local_result
            if not omr_result:
                reason = "low-contrast answer recovered by secondary measurement"
            elif len(omr_result) > 1:
                reason = "false multiple resolved by secondary measurement"
            else:
                reason = "OMR and secondary measurement disagree"
        elif len(omr_result) > 1:
            reason = "multiple marks"
        elif omr_result and local_class == "indeterminate":
            reason = "low local darkness separation"

        if not final_result and number not in optional:
            reason = reason or "required question is blank or unreadable"

        answers[field] = final_result
        if reason:
            issues.append(
                {
                    "field": field,
                    "omr_result": omr_result,
                    "result": final_result,
                    "reason": reason,
                    **{key: value for key, value in local.items() if key != "scores"},
                }
            )
    return answers, issues


def process(
    input_path: Path,
    form_id: str,
    bundle_root: Path,
    omrchecker_root: Path,
    artifacts_dir: Path | None = None,
) -> dict[str, Any]:
    form_dir = bundle_root / "forms" / form_id
    if not form_dir.is_dir():
        raise ValueError(f"Unknown form ID: {form_id}")
    profile = json.loads((form_dir / "form_profile.json").read_text())
    geometry = json.loads((bundle_root / "shared/geometry_and_controls.json").read_text())
    mappings = json.loads((bundle_root / "shared/id_mappings.json").read_text())
    template = json.loads((form_dir / "template.json").read_text())
    marker = cv2.imread(str(form_dir / "omr_marker.jpg"), cv2.IMREAD_GRAYSCALE)
    if marker is None:
        raise RuntimeError(f"Unable to load corner marker for {form_id}")

    source = load_image(input_path)
    candidates = orientation_candidates(source, geometry, marker)
    if not candidates:
        return {"status": "rescan", "form_id": form_id, "flags": ["form border not found"]}
    candidates.sort(
        key=lambda item: (
            item["controls"]["exact"],
            item["controls"]["matches"],
            item["controls"]["worst_case_gap"],
            item["controls"]["separation_gray_levels"],
        ),
        reverse=True,
    )
    winner = candidates[0]
    metrics = illumination_metrics(winner["image"])
    normalized, normalization = normalize_illumination(winner["image"], metrics)
    normalized_controls = control_measurements(normalized, geometry)
    flags = []
    if not normalized_controls["exact"]:
        flags.append("control pattern mismatch")
    if normalized_controls["separation_gray_levels"] < 60:
        flags.append("low control-mark contrast")
    if metrics["laplacian_variance"] < 45:
        flags.append("image may be blurred")

    result: dict[str, Any] = {
        "status": "rescan" if any("control" in flag for flag in flags) else "accepted",
        "form_id": form_id,
        "input_file": input_path.name,
        "rotation_degrees_clockwise": winner["rotation_degrees_clockwise"],
        "alignment_method": winner["alignment_method"],
        "normalization": normalization,
        "quality": {
            **metrics,
            "orientation_controls_before_normalization": winner["controls"],
            "controls": normalized_controls,
        },
        "flags": flags,
    }
    if result["status"] == "rescan":
        return result

    omr_row = run_omr(normalized, form_dir, omrchecker_root)
    raw_id = omr_row.get(profile["raw_id_output"], "")
    try:
        driver_id = decode_driver_id(raw_id, mappings[profile["id_mapping"]])
    except ValueError as error:
        driver_id = ""
        flags.append(str(error))

    measurements = field_measurements(normalized, template)
    answers, confidence_issues = reconcile_answers(omr_row, measurements, profile)
    if confidence_issues:
        flags.append("one or more marked fields require review")
    result.update(
        {
            "status": "review" if flags else "accepted",
            "driver_id": driver_id,
            "driver_id_rows": raw_id,
            "answers": answers,
            "confidence_issues": confidence_issues,
            "secondary_measurements": measurements,
        }
    )

    if artifacts_dir:
        artifacts_dir.mkdir(parents=True, exist_ok=True)
        cv2.imwrite(str(artifacts_dir / "normalized.png"), normalized)
        (artifacts_dir / "result.json").write_text(json.dumps(result, indent=2) + "\n")
    return result


def main() -> int:
    args = parse_args()
    result = process(
        args.input.resolve(),
        args.form_id,
        args.bundle_root.resolve(),
        args.omrchecker_root.resolve(),
        args.artifacts_dir.resolve() if args.artifacts_dir else None,
    )
    text = json.dumps(result, indent=2) + "\n"
    if args.output_json:
        args.output_json.parent.mkdir(parents=True, exist_ok=True)
        args.output_json.write_text(text)
    else:
        print(text, end="")
    return 0 if result["status"] != "rescan" else 2


if __name__ == "__main__":
    raise SystemExit(main())