engine/id_parser.py

import re
from datetime import datetime, date
from difflib import SequenceMatcher


# ── Country-specific ID number validation ──────────────────────────────
# MRZ/VLM nationality codes (ISO 3166-1 alpha-3) → validation rules.
# SA uses Luhn check; others use regex only.

COUNTRY_ID_FORMATS = {
    "ZAF": {"regex": r"^[0-9]{13}$", "name": "South Africa", "has_luhn": True},
    "NGA": {"regex": r"^[0-9]{11}$", "name": "Nigeria", "has_luhn": False},
    "KEN": {"regex": r"^[0-9]{1,9}$", "name": "Kenya", "has_luhn": False},
    "ZWE": {"regex": r"^[0-9]{8,9}[A-Za-z]\d{2}$", "name": "Zimbabwe", "has_luhn": False},
    "UGA": {"regex": r"^[A-Z0-9]{14}$", "name": "Uganda", "has_luhn": False},
    "ZMB": {"regex": r"^[0-9]{10}$", "name": "Zambia", "has_luhn": False},
    "CIV": {"regex": r"^[A-Z]{2}[0-9]+$", "name": "Côte d'Ivoire", "has_luhn": False},
}


def validate_id_for_country(id_number: str, country_code: str) -> dict:
    """Validate an ID number against country-specific format rules.

    Args:
        id_number: The ID number to validate.
        country_code: ISO 3166-1 alpha-3 country code (e.g., 'ZAF', 'NGA').

    Returns:
        Dict with 'valid' (bool), 'country' (str), 'method' (str).
    """
    if not id_number or not country_code:
        return {"valid": False, "country": country_code, "method": "unknown"}

    fmt = COUNTRY_ID_FORMATS.get(country_code.upper())
    if not fmt:
        # Unknown country — can't validate, not a failure
        return {"valid": True, "country": country_code, "method": "no_rules"}

    # Regex check
    if not re.match(fmt["regex"], id_number):
        return {"valid": False, "country": fmt["name"], "method": "regex"}

    # SA gets additional Luhn validation
    if fmt["has_luhn"]:
        luhn_ok = validate_luhn(id_number)
        return {"valid": luhn_ok, "country": fmt["name"], "method": "luhn" if luhn_ok else "luhn_failed"}

    return {"valid": True, "country": fmt["name"], "method": "regex"}


# OCR commonly confuses these characters with digits.
# Conservative set — only substitutions that are visually unambiguous.
# Aggressive subs like S→5, B→8, G→6, D→0 create false ID numbers from normal text.
OCR_DIGIT_SUBS = str.maketrans({
    "O": "0", "o": "0",
    "I": "1", "l": "1", "|": "1",
    "Q": "0",
})

FIELD_LABELS = {
    "id_number": ["IDENTITY NUMBER", "IDENTITY NO", "ID NUMBER", "ID NO", "I.D. NO", "I.D. NUMBER"],
    "surname": ["SURNAME", "LAST NAME"],
    "names": ["NAMES", "FIRST NAMES", "FORENAMES"],
    "date_of_birth": ["DATE OF BIRTH", "BIRTH DATE", "DOB"],
    "sex": ["SEX", "GENDER"],
    "country_of_birth": ["COUNTRY OF BIRTH", "PLACE OF BIRTH"],
    "nationality": ["NATIONALITY"],
    "citizenship_status": ["STATUS", "CITIZENSHIP", "S.A. CITIZEN", "SA CITIZEN"],
}


def validate_luhn(id_number: str) -> bool:
    """Validate a 13-digit SA ID number using the SA-specific Luhn algorithm.

    SA uses a non-standard Luhn variant:
    1. Sum digits at odd positions (1st, 3rd, 5th... 1-indexed).
    2. Concatenate digits at even positions into one number, multiply by 2,
       then sum the individual digits of the result.
    3. Total = sum_odd + sum_even_digits.
    4. Check digit = (10 - (total % 10)) % 10.
    """
    if not id_number or len(id_number) != 13 or not id_number.isdigit():
        return False

    digits = [int(d) for d in id_number]

    # Sum of odd-positioned digits (0, 2, 4, 6, 8, 10 in 0-indexed = 1st, 3rd, 5th...)
    odd_sum = sum(digits[i] for i in range(0, 12, 2))

    # Concatenate even-positioned digits (1, 3, 5, 7, 9, 11 in 0-indexed) and multiply by 2
    even_concat = int("".join(str(digits[i]) for i in range(1, 12, 2)))
    even_doubled = even_concat * 2
    even_sum = sum(int(d) for d in str(even_doubled))

    total = odd_sum + even_sum
    check_digit = (10 - (total % 10)) % 10

    return check_digit == digits[12]


def parse_id_number(id_number: str) -> dict:
    """Extract encoded fields from a valid SA ID number.

    SA ID format: YYMMDD SSSS C A Z
    - YYMMDD: Date of birth
    - SSSS: Gender (0000-4999=Female, 5000-9999=Male)
    - C: Citizenship (0=SA citizen, 1=permanent resident)
    - A: Usually 8 or 9
    - Z: Luhn check digit
    """
    result = {
        "date_of_birth": None,
        "sex": None,
        "citizenship": None,
        "is_valid": False,
    }

    if not id_number or len(id_number) != 13 or not id_number.isdigit():
        return result

    result["is_valid"] = validate_luhn(id_number)

    # Parse date of birth
    yy = int(id_number[0:2])
    mm = int(id_number[2:4])
    dd = int(id_number[4:6])

    # Century heuristic: if YY > current 2-digit year, assume 1900s
    current_yy = datetime.now().year % 100
    year = 1900 + yy if yy > current_yy else 2000 + yy

    try:
        dob = date(year, mm, dd)
        result["date_of_birth"] = dob.isoformat()
    except ValueError:
        pass  # Invalid date encoded in ID

    # Parse gender
    gender_seq = int(id_number[6:10])
    result["sex"] = "Male" if gender_seq >= 5000 else "Female"

    # Parse citizenship
    citizenship_digit = int(id_number[10])
    result["citizenship"] = "SA Citizen" if citizenship_digit == 0 else "Permanent Resident"

    return result


def extract_id_number(text_blocks: list) -> tuple:
    """Find a 13-digit ID number in OCR text blocks.

    Tries raw text first, then applies OCR error corrections.
    Returns (id_number, confidence) or (None, 0).
    """
    candidates = []

    for item in text_blocks:
        if len(item) >= 3:
            _, text, confidence = item[0], item[1], item[2]
        else:
            text = item if isinstance(item, str) else str(item)
            confidence = 0.0

        # Clean the text
        cleaned = re.sub(r"[\s\-\.]", "", text)

        # Try finding 13 consecutive digits directly
        matches = re.findall(r"\d{13}", cleaned)
        for match in matches:
            if validate_luhn(match):
                candidates.append((match, confidence))

        # Apply OCR substitutions and try again
        substituted = cleaned.translate(OCR_DIGIT_SUBS)
        # Extract only digit characters after substitution
        digits_only = re.sub(r"[^0-9]", "", substituted)

        # Slide through looking for 13-digit sequences
        for i in range(len(digits_only) - 12):
            seq = digits_only[i:i + 13]
            if validate_luhn(seq) and seq not in [c[0] for c in candidates]:
                candidates.append((seq, confidence * 0.9))  # Slightly lower confidence for substituted

    if candidates:
        # Return the highest confidence match
        candidates.sort(key=lambda x: x[1], reverse=True)
        return candidates[0]

    return None, 0.0


def _bbox_center_y(bbox):
    """Get vertical center of a bounding box."""
    return (bbox[0][1] + bbox[2][1]) / 2


def _bbox_center_x(bbox):
    """Get horizontal center of a bounding box."""
    return (bbox[0][0] + bbox[2][0]) / 2


def _bbox_right(bbox):
    """Get right edge of a bounding box."""
    return bbox[1][0]


def _bbox_left(bbox):
    """Get left edge of a bounding box."""
    return bbox[0][0]


def _find_value_for_label(label_idx, results, y_tolerance=25, x_tolerance=50, skip_indices=None):
    """Find the OCR text value associated with a label.

    Looks to the right of the label first, then below it.
    Skips indices in skip_indices (other known labels) to avoid picking labels as values.
    """
    if skip_indices is None:
        skip_indices = set()

    label_bbox = results[label_idx][0]
    label_y = _bbox_center_y(label_bbox)
    label_right = _bbox_right(label_bbox)
    label_left = _bbox_left(label_bbox)
    label_width = label_right - label_left
    label_height = label_bbox[2][1] - label_bbox[0][1]

    right_candidates = []
    below_candidates = []

    # Max horizontal gap for "right of label" — prevent picking text far away
    max_right_gap = max(label_width * 0.5, 100)
    # Scale below-label tolerances with label dimensions
    below_y_max = max(80, label_height * 5)
    effective_x_tolerance = max(x_tolerance, label_width * 0.4)

    for i, (bbox, text, conf) in enumerate(results):
        if i == label_idx or i in skip_indices:
            continue

        text_y = _bbox_center_y(bbox)
        text_left = _bbox_left(bbox)

        # Right of label: same vertical level, starts after label ends, not too far
        gap = text_left - label_right
        if abs(text_y - label_y) < y_tolerance and gap > -10 and gap < max_right_gap:
            right_candidates.append((gap, text, conf))

        # Below label: lower vertical position, roughly same horizontal area
        elif 5 < (text_y - label_y) < below_y_max and abs(text_left - label_left) < effective_x_tolerance:
            distance = text_y - label_y
            below_candidates.append((distance, text, conf))

    # Prefer right-of-label, fall back to below
    if right_candidates:
        right_candidates.sort(key=lambda x: x[0])
        return right_candidates[0][1], right_candidates[0][2]

    if below_candidates:
        below_candidates.sort(key=lambda x: x[0])
        return below_candidates[0][1], below_candidates[0][2]

    return None, 0.0


def _match_label(text, label_variants):
    """Check if text matches any label variant."""
    text_upper = text.upper().strip()
    if len(text_upper) < 2:
        return False
    for variant in label_variants:
        # Exact match
        if text_upper == variant:
            return True
        # Substring match — require at least 50% overlap to avoid false positives
        if variant in text_upper or text_upper in variant:
            min_len = min(len(text_upper), len(variant))
            max_len = max(len(text_upper), len(variant))
            if min_len >= 3 and min_len / max_len > 0.5:
                return True
        # Fuzzy match for OCR errors in labels (>= 5 chars, >= 80% similarity)
        if len(text_upper) >= 5 and len(variant) >= 5:
            ratio = SequenceMatcher(None, text_upper, variant).ratio()
            if ratio >= 0.8:
                return True
    return False


def extract_fields(ocr_results: list) -> dict:
    """Extract structured SA ID fields from EasyOCR output.

    Two-pass approach:
    1. Identify all label positions in the OCR results.
    2. For each label, find the nearest non-label text as the value.
       This prevents picking other labels (e.g. "Nationality") as values.

    Args:
        ocr_results: List of (bounding_box, text, confidence) from EasyOCR.

    Returns:
        Dict with extracted fields and confidence scores.
    """
    fields = {}
    confidence = {}

    # Pass 1: Identify which OCR results are labels
    label_indices = {}  # field_name -> index in ocr_results
    for field_name, label_variants in FIELD_LABELS.items():
        for i, (bbox, text, conf) in enumerate(ocr_results):
            if _match_label(text, label_variants):
                label_indices[field_name] = i
                break

    # Set of all label indices — these should be skipped when looking for values
    all_label_idxs = set(label_indices.values())

    # Pass 2: Find values for each identified label
    for field_name, label_idx in label_indices.items():
        if field_name == "id_number":
            continue  # Handled separately with OCR error correction
        value, value_conf = _find_value_for_label(
            label_idx, ocr_results, skip_indices=all_label_idxs
        )
        if value:
            fields[field_name] = value.strip()
            confidence[field_name] = round(value_conf, 2)

    # Extract ID number using specialized logic (handles OCR errors)
    id_number, id_conf = extract_id_number(ocr_results)
    if id_number:
        fields["id_number"] = id_number
        confidence["id_number"] = round(id_conf, 2)

    # Normalize sex field
    if "sex" in fields:
        sex_upper = fields["sex"].upper().strip()
        if sex_upper in ("M", "MALE"):
            fields["sex"] = "Male"
        elif sex_upper in ("F", "FEMALE"):
            fields["sex"] = "Female"

    return {"fields": fields, "confidence": confidence}


def _check_date_transposition(ocr_date: str, id_date: str) -> dict:
    """Check if day and month are swapped between OCR date and ID-encoded date.

    Common in SA because DD/MM and MM/DD formats are both used.
    Inspired by Smile Identity's verifyDOB approach.

    Returns dict with match status and details.
    """
    result = {"exact_match": False, "transposed_match": False, "corrected_date": None}

    if not ocr_date or not id_date:
        return result

    normalized = _normalize_date(ocr_date)
    if not normalized:
        return result

    if normalized == id_date:
        result["exact_match"] = True
        return result

    # Try swapping day and month in the OCR date
    try:
        ocr_parts = normalized.split("-")
        if len(ocr_parts) == 3:
            year, month, day = ocr_parts
            swapped = f"{year}-{day}-{month}"
            # Only valid if both day and month are <= 12 (ambiguous dates)
            if int(day) <= 12 and int(month) <= 12:
                swapped_date = date(int(year), int(day), int(month))
                swapped_str = swapped_date.isoformat()
                if swapped_str == id_date:
                    result["transposed_match"] = True
                    result["corrected_date"] = id_date
    except (ValueError, IndexError):
        pass

    return result


def cross_validate(fields: dict, parsed_id: dict) -> dict:
    """Compare OCR-extracted fields against ID-number-encoded values.

    Returns validation results dict.
    """
    validation = {
        "luhn_valid": parsed_id.get("is_valid", False),
        "dob_cross_check": None,
        "dob_transposed": False,
        "gender_cross_check": None,
        "citizenship_cross_check": None,
    }

    # DOB cross-check with transposition detection
    if parsed_id.get("date_of_birth") and fields.get("date_of_birth"):
        ocr_dob = fields["date_of_birth"].strip()
        id_dob = parsed_id["date_of_birth"]

        transposition = _check_date_transposition(ocr_dob, id_dob)

        if transposition["exact_match"]:
            validation["dob_cross_check"] = True
        elif transposition["transposed_match"]:
            # Day/month were swapped — auto-correct to the ID-encoded date
            validation["dob_cross_check"] = True
            validation["dob_transposed"] = True
            fields["date_of_birth"] = transposition["corrected_date"]
        else:
            normalized = _normalize_date(ocr_dob)
            validation["dob_cross_check"] = normalized == id_dob if normalized else None

    # Gender cross-check
    if parsed_id.get("sex") and fields.get("sex"):
        validation["gender_cross_check"] = (
            fields["sex"].upper().strip() == parsed_id["sex"].upper()
        )

    # Citizenship cross-check
    if parsed_id.get("citizenship") and fields.get("citizenship_status"):
        ocr_cit = fields["citizenship_status"].upper().strip()
        id_cit = parsed_id["citizenship"].upper()
        # Flexible matching: "SA CITIZEN" matches "CITIZEN", "RSA" etc.
        validation["citizenship_cross_check"] = (
            "CITIZEN" in ocr_cit and "CITIZEN" in id_cit
        ) or (
            "RESIDENT" in ocr_cit and "RESIDENT" in id_cit
        )

    return validation


def _normalize_date(date_str: str) -> str | None:
    """Try to parse various date formats into YYYY-MM-DD."""
    date_str = date_str.strip()
    formats = [
        "%Y-%m-%d",
        "%d-%m-%Y",
        "%d/%m/%Y",
        "%Y/%m/%d",
        "%d %b %Y",
        "%d %B %Y",
        "%Y %m %d",
        "%d %m %Y",
    ]
    for fmt in formats:
        try:
            return datetime.strptime(date_str, fmt).date().isoformat()
        except ValueError:
            continue
    return None