import io
import itertools
import os
import re
from collections import Counter
from concurrent.futures import ThreadPoolExecutor
from decimal import Decimal

import fitz
import pytesseract
from PIL import Image, ImageFilter, ImageOps


OCR_TARGET_WIDTH = 2500

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

DIGIT_FIX = str.maketrans({
    "O": "0", "o": "0", "I": "1", "l": "1", "|": "1",
    "]": "1", "[": "1", "B": "8", "a": "8", "S": "5",
})

EIN_FIX = str.maketrans({
    "O": "0", "o": "0", "I": "1", "l": "1", "|": "1",
})

AMT_RE = re.compile(
    r"(?<![A-Za-z0-9.,])"
    r"(?P<whole>[\dBa][\dOoBaSIl|,]*)"
    r"\s*[.,·•«»]\s*"
    r"(?P<cents>[\dOoIl|]{2})[0Oo]?"
    r"(?![\dA-Za-z])"
)

EIN_LOOSE = re.compile(
    r"(?<![\dA-Za-z])(\d{2})\s*[-–—]\s*([\dOoIl|?]{7})(?![\dA-Za-z])"
)

EIN_STRICT = re.compile(r"^\d{2}-\d{7}$")
PTIN_STRICT = re.compile(r"^P\d{8}$")

PHONE_RE = re.compile(
    r"(?<!\d)"
    r"\(?(\d{3})\)?"
    r"[\s\-=.~_—–]*"
    r"(\d{3})"
    r"[\s\-=.~_—–]*"
    r"(\d{4})"
    r"(?!\d)"
)

US_STATES = frozenset({
    "AL", "AK", "AZ", "AR", "CA", "CO", "CT", "DE", "DC", "FL", "GA", "HI",
    "ID", "IL", "IN", "IA", "KS", "KY", "LA", "ME", "MD", "MA", "MI", "MN",
    "MS", "MO", "MT", "NE", "NV", "NH", "NJ", "NM", "NY", "NC", "ND", "OH",
    "OK", "OR", "PA", "RI", "SC", "SD", "TN", "TX", "UT", "VT", "VA", "WA",
    "WV", "WI", "WY", "AS", "GU", "MP", "PR", "VI", "AA", "AE", "AP",
})

ZIP_PREFIX_STATE = {
    "00": "PR", "01": "MA", "02": "RI", "03": "NH", "04": "ME", "05": "VT",
    "06": "CT", "07": "NJ", "08": "NJ", "09": "AE",
    "10": "NY", "11": "NY", "12": "NY", "13": "NY", "14": "NY",
    "15": "PA", "16": "PA", "17": "PA", "18": "PA", "19": "PA",
    "20": "DC", "21": "MD", "22": "VA", "23": "VA", "24": "VA",
    "25": "WV", "26": "WV", "27": "NC", "28": "NC", "29": "SC",
    "30": "GA", "31": "GA", "32": "FL", "33": "FL", "34": "FL",
    "40": "KY", "41": "KY", "42": "KY", "43": "OH", "44": "OH",
    "45": "OH", "46": "IN", "47": "IN", "48": "MI", "49": "MI",
    "50": "IA", "51": "IA", "52": "IA", "53": "WI", "54": "WI",
    "55": "MN", "56": "SD", "57": "SD", "58": "ND", "59": "MT",
    "60": "IL", "61": "IL", "62": "IL", "63": "MO", "64": "MO",
    "65": "MO", "66": "KS", "67": "KS", "68": "NE", "69": "NE",
    "70": "LA", "71": "LA", "72": "AR", "73": "OK", "74": "OK",
    "75": "TX", "76": "TX", "77": "TX", "78": "TX", "79": "TX",
    "80": "CO", "81": "CO", "82": "WY", "83": "ID", "84": "UT",
    "85": "AZ", "86": "AZ", "87": "NM", "88": "NM", "89": "NV",
    "90": "CA", "91": "CA", "92": "CA", "93": "CA", "94": "CA",
    "95": "CA", "96": "CA", "97": "OR", "98": "WA", "99": "AK",
}


def state_from_zip(zip_code):
    if not zip_code:
        return None

    digits = re.sub(r"\D", "", str(zip_code))

    if len(digits) < 5:
        return None

    return ZIP_PREFIX_STATE.get(digits[:2])


def clean(value):
    if value is None:
        return None

    value = re.sub(r"\s+", " ", str(value)).strip()
    return value or None


def to_amount(whole, cents):
    whole = whole.translate(DIGIT_FIX).replace(",", "")
    cents = cents.translate(DIGIT_FIX)

    if not (whole.isdigit() and cents.isdigit()):
        return None

    return float(Decimal(f"{whole}.{cents}"))


def amounts_in(text):
    values = [
        to_amount(match["whole"], match["cents"])
        for match in AMT_RE.finditer(text or "")
    ]

    return [value for value in values if value is not None]


def labeled_amounts(text, label_patterns):
    lines = (text or "").splitlines()

    for pattern in label_patterns:
        for line in lines:
            match = re.search(pattern, line, re.I)

            if not match:
                continue

            values = amounts_in(line[match.end():])

            if values:
                return values

    return []


def find_labeled_amount(text, label_patterns, occurrence=-1):
    values = labeled_amounts(text, label_patterns)

    if not values:
        return None

    try:
        return values[occurrence]
    except IndexError:
        return None


def find_first_match(pattern, text, group=1):
    match = re.search(pattern, text or "", re.I | re.M)
    return clean(match.group(group)) if match else None


def find_ein(*texts):
    partial = None

    for text in texts:
        for match in EIN_LOOSE.finditer(text or ""):
            digits = match.group(2).translate(EIN_FIX)
            ein = f"{match.group(1)}-{digits}"

            if re.fullmatch(r"\d{7}", digits):
                return ein, None

            partial = partial or ein

    return None, partial


# Legal suffixes help remove trailing OCR noise from employer names.
LEGAL_SUFFIX_RE = re.compile(
    r"\b(?:CORPORATION|INCORPORATED|INC|LLC|LTD|LIMITED|CORP|COMPANY|CO|LLP|LP|PC|PLLC)\b\.?",
    re.I,
)


def _trim_employer_name(value):
    if not value:
        return value

    value = value.strip()
    matches = list(LEGAL_SUFFIX_RE.finditer(value))

    if matches:
        trimmed = value[:matches[-1].end()].strip(" .,-")

        if trimmed:
            return trimmed

    trimmed = re.sub(r"\s+[A-Za-z]{1,2}$", "", value).strip(" .,-")
    return trimmed or None


def find_employer_name(*pages):
    pattern = (
        r"(?i:Name\s*\(not\s+your\s+trade\s+name\))\s*[:|]?\s*\n*\s*"
        r"([A-Za-z0-9][A-Za-z0-9 &.,'/\-]{2,}?)\s*(?:\n|$)"
    )

    for text in pages:
        match = re.search(pattern, text or "")

        if match:
            name = clean(match.group(1))
            trimmed = _trim_employer_name(name)

            if trimmed:
                return trimmed

    return None


def format_phone(value):
    digits = re.sub(r"\D", "", value or "")

    if len(digits) == 11 and digits.startswith("1"):
        digits = digits[1:]

    if len(digits) != 10:
        return None

    return f"{digits[:3]}-{digits[3:6]}-{digits[6:]}"


def extract_phone(text):
    if not text:
        return None

    normalized = re.sub(r"[=~_—–]", "-", text)
    match = PHONE_RE.search(normalized)

    if not match:
        return None

    return f"{match.group(1)}-{match.group(2)}-{match.group(3)}"


# ---------------------------------------------------------------------------
# PDF text / OCR extraction
# ---------------------------------------------------------------------------

def page_image(page, document):
    images = page.get_images(full=True)

    if len(images) == 1:
        try:
            info = document.extract_image(images[0][0])
            image = Image.open(io.BytesIO(info["image"])).convert("L")

            if image.width >= 600:
                return image
        except Exception:
            pass

    pix = page.get_pixmap(
        matrix=fitz.Matrix(300 / 72, 300 / 72),
        alpha=False,
    )

    return Image.open(io.BytesIO(pix.tobytes("png"))).convert("L")


def prepare_for_ocr(image, sharpen):
    scale = OCR_TARGET_WIDTH / image.width

    if scale > 1.05:
        image = image.resize(
            (round(image.width * scale), round(image.height * scale)),
            Image.Resampling.LANCZOS,
        )

    image = ImageOps.autocontrast(image)

    if sharpen:
        image = image.filter(ImageFilter.UnsharpMask(2, 150, 3))

    return image


def ocr_image(image):
    text = pytesseract.image_to_string(
        image,
        lang="eng",
        config="--psm 6",
    )

    if len(text.strip()) < 10:
        text = pytesseract.image_to_string(
            image,
            lang="eng",
            config="--psm 3",
        )

    return text


def extract_text(pdf_bytes):
    pages = []
    ocr_jobs = []
    document = fitz.open(stream=pdf_bytes, filetype="pdf")

    try:
        for index, page in enumerate(document):
            text = page.get_text("text", sort=True) or ""

            item = {
                "page": index + 1,
                "text": text,
                "text_alt": "",
                "width": round(page.rect.width, 2),
                "height": round(page.rect.height, 2),
                "source": "pdf_text",
                "image": None,
            }

            if len(re.sub(r"\s+", "", text)) < 40:
                tesseract_cmd = pytesseract.pytesseract.tesseract_cmd

                if not os.path.isfile(tesseract_cmd):
                    item["ocr_error"] = (
                        "Tesseract executable was not found. "
                        "Install Tesseract or configure TESSERACT_CMD."
                    )
                else:
                    try:
                        base = page_image(page, document)
                        item["image"] = base
                        item["source"] = "ocr"

                        for sharpen in (False, True):
                            ocr_jobs.append(
                                (
                                    index,
                                    sharpen,
                                    prepare_for_ocr(base, sharpen),
                                )
                            )

                    except Exception as exc:
                        item["ocr_error"] = (
                            f"OCR failed on page {index + 1}: {exc}"
                        )

            pages.append(item)

    finally:
        document.close()

    if ocr_jobs:
        def run(job):
            index, sharpen, image = job

            try:
                return index, sharpen, ocr_image(image), None
            except Exception as exc:
                return index, sharpen, "", exc

        workers = max(2, min(4, os.cpu_count() or 2))

        with ThreadPoolExecutor(max_workers=workers) as pool:
            for index, sharpen, text, exc in pool.map(run, ocr_jobs):
                if exc is not None:
                    pages[index]["ocr_error"] = (
                        f"OCR failed on page {index + 1}: {exc}"
                    )
                elif sharpen:
                    pages[index]["text_alt"] = text
                else:
                    pages[index]["text"] = text

    return pages


# ---------------------------------------------------------------------------
# Small-box OCR
# ---------------------------------------------------------------------------

def zone_ocr(image, x0, x1, yc, whitelist=None, half=0.015, scale=4):
    width, height = image.size

    box = image.crop((
        int(width * x0),
        int(height * (yc - half)),
        int(width * x1),
        int(height * (yc + half)),
    ))

    box = box.resize(
        (box.width * scale, box.height * scale),
        Image.Resampling.LANCZOS,
    )

    box = ImageOps.expand(
        ImageOps.autocontrast(box),
        border=30,
        fill=255,
    )

    config = "--psm 7"

    if whitelist:
        config += f" -c tessedit_char_whitelist={whitelist}"

    return pytesseract.image_to_string(box, config=config).strip()


def zone_read(image, box, yc, kind="text", whitelist=None):
    if image is None:
        return None

    x0, x1 = box
    votes = Counter()

    for dy in (0.0, 0.004, -0.004, 0.008, -0.008):
        try:
            raw = zone_ocr(
                image,
                x0,
                x1,
                yc + dy,
                whitelist=whitelist,
            )
        except Exception:
            continue

        raw = re.sub(r"^[^\w]+|[^\w.)]+$", "", raw).strip()
        value = None

        if kind == "ein":
            match = EIN_LOOSE.search(raw)

            if match:
                digits = match.group(2).translate(EIN_FIX)

                if digits.isdigit():
                    value = f"{match.group(1)}-{digits}"

        elif kind == "ptin":
            candidate = re.sub(r"[\s-]", "", raw).upper().replace("O", "0")

            if PTIN_STRICT.match(candidate):
                value = candidate

        elif kind == "phone":
            value = extract_phone(raw)

        elif kind == "zip":
            match = re.search(r"\b(\d{5}(?:-\d{4})?)\b", raw)
            value = match.group(1) if match else None

        elif kind == "state":
            match = re.search(r"\b([A-Z]{2})\b", raw.upper())

            if match and match.group(1) in US_STATES:
                value = match.group(1)

        elif kind == "int":
            digits = re.sub(r"\D", "", raw)

            if digits:
                value = digits

        else:
            letters = sum(ch.isalpha() for ch in raw)

            if len(raw) >= 3 and letters >= 0.6 * len(raw):
                value = clean(raw.strip("|[]_ "))

        if value:
            votes[value] += 1

            if votes.most_common(1)[0][1] >= 2:
                break

    return votes.most_common(1)[0][0] if votes else None


# ---------------------------------------------------------------------------
# Form 941 line 1 employee-count OCR
# ---------------------------------------------------------------------------

# Line 1 box on IRS Form 941 (Rev. June 2022) sits at ≈ 34.5 % of page
# height. The old value (0.365) missed the digit → box OCR returned None
# and the whole-page text misread the "3" as "4".
PAGE1_LINE1_X = (0.70, 0.96)   # slightly wider to catch edges
PAGE1_LINE1_Y = 0.345           # ← corrected anchor

def read_line1_employees(image):
    if image is None:
        return None

    width, height = image.size
    x0, x1 = PAGE1_LINE1_X
    votes = Counter()

    # Wider y-grid + tighter half-height → isolated digit, no border noise.
    for dy in (0.0, -0.005, 0.005, -0.010, 0.010,
               -0.015, 0.015, -0.020, 0.020):
        yc = PAGE1_LINE1_Y + dy

        crop = image.crop((
            max(0, int(width * x0)),
            max(0, int(height * (yc - 0.012))),
            min(width, int(width * x1)),
            min(height, int(height * (yc + 0.012))),
        ))
        if crop.width < 10 or crop.height < 10:
            continue

        gray = ImageOps.autocontrast(crop.convert("L"))

        for threshold in (120, 150, 180):
            binary = gray.point(
                lambda p, t=threshold: 255 if p > t else 0
            )
            enlarged = binary.resize(
                (binary.width * 6, binary.height * 6),
                Image.Resampling.LANCZOS,
            )
            enlarged = ImageOps.expand(enlarged, border=30, fill=255)

            # PSM 10 = "single character" mode — purpose-built for an
            # isolated digit inside a box. Confidence weights the votes so
            # a clean "3" dominates any borderline "4" from noise.
            try:
                data = pytesseract.image_to_data(
                    enlarged,
                    config=(
                        "--psm 10 "
                        "-c tessedit_char_whitelist=0123456789"
                    ),
                    output_type=pytesseract.Output.DICT,
                )
            except Exception:
                continue

            for i, text in enumerate(data.get("text", [])):
                text = (text or "").strip()
                if len(text) == 1 and text.isdigit():
                    try:
                        conf = max(0, int(data["conf"][i]))
                    except (ValueError, TypeError):
                        conf = 1
                    if conf >= 40:
                        votes[int(text)] += conf

    if not votes:
        return None

    ranked = votes.most_common()

    # Genuine tie → don't guess.
    if len(ranked) > 1 and ranked[0][1] == ranked[1][1]:
        return None

    return ranked[0][0]

# ---------------------------------------------------------------------------
# Page 3 signature / designee / preparer OCR zones
# ---------------------------------------------------------------------------

PAGE3_ZONES = {
    "signer_name": ((0.71, 0.95), 0.629),
    "signer_title": ((0.71, 0.95), 0.6597),
    "signer_phone": ((0.72, 0.98), 0.699),
    "designee_name": ((0.42, 0.75), 0.502),
    "designee_phone": ((0.72, 0.98), 0.502),
    "preparer_name": ((0.19, 0.63), 0.759),
    "preparer_ptin": ((0.72, 0.98), 0.759),
    "preparer_firm_name": ((0.19, 0.63), 0.819),
    "preparer_firm_ein": ((0.72, 0.98), 0.819),
    "preparer_address": ((0.19, 0.68), 0.849),
    "preparer_phone": ((0.68, 0.98), 0.849),
    "preparer_city": ((0.19, 0.50), 0.879),
    "preparer_state": ((0.56, 0.66), 0.879),
    "preparer_zip": ((0.72, 0.98), 0.879),
}

PAGE3_KINDS = {
    "signer_phone": ("phone", "0123456789 -().="),
    "designee_phone": ("phone", "0123456789 -().="),
    "preparer_ptin": ("ptin", "P0123456789 "),
    "preparer_firm_ein": ("ein", "0123456789-"),
    "preparer_phone": ("phone", "0123456789 -().="),
    "preparer_zip": ("zip", "0123456789-"),
    "preparer_state": ("state", "ABCDEFGHIJKLMNOPQRSTUVWXYZ"),
}


def read_page3_zones(image):
    if image is None:
        return {}

    def job(item):
        key, (box, yc) = item
        kind, whitelist = PAGE3_KINDS.get(key, ("text", None))

        return key, zone_read(
            image,
            box,
            yc,
            kind=kind,
            whitelist=whitelist,
        )

    with ThreadPoolExecutor(max_workers=4) as pool:
        return dict(pool.map(job, PAGE3_ZONES.items()))


# ---------------------------------------------------------------------------
# Form 941 parser
# ---------------------------------------------------------------------------

ALL_FIELDS = (
    "tax_year", "quarter",
    "employer_name", "employer_ein", "trade_name",
    "employer_address", "employer_city", "employer_state", "employer_zip",
    "line1_employees", "line2_wages", "line3_fed_tax", "line4_no_wages",
    "line5a_ss_wages", "line5a_ss_tax",
    "line5a_i_qualified_sick_leave_wages",
    "line5a_i_qualified_sick_leave_tax",
    "line5a_ii_qualified_family_leave_wages",
    "line5a_ii_qualified_family_leave_tax",
    "line5b_ss_tips", "line5b_ss_tips_tax",
    "line5c_medicare_wages", "line5c_medicare_tax",
    "line5d_addl_medicare_wages", "line5d_addl_medicare_tax",
    "line5e_total_ss_medicare", "line5f_unreported_tips",
    "line6_total_before_adj", "line7_fractions_cents",
    "line8_sick_pay_adj", "line9_tips_gtl_adj",
    "line10_total_after_adj",
    "line11a_research_credit",
    "line11b_nonrefundable_sick_family_before_apr2021",
    "line11d_nonrefundable_sick_family_after_mar2021",
    "line11g_total_nonrefundable",
    "line12_total_taxes",
    "line13a_total_deposits",
    "line13c_refundable_sick_family_before_apr2021",
    "line13e_refundable_sick_family_after_mar2021",
    "line13g_total_deposits_credits",
    "line14_balance_due", "line15_overpayment",
    "month1_liability", "month2_liability", "month3_liability",
    "total_liability",
    "signer_name", "signer_title", "signer_phone",
    "designee_name", "designee_phone",
    "preparer_name", "preparer_ptin", "preparer_firm_name",
    "preparer_firm_ein", "preparer_address", "preparer_city",
    "preparer_state", "preparer_zip", "preparer_phone",
)


def parse_941(pages, filename=None):
    def page_text(index):
        if index >= len(pages):
            return ""

        return (
            (pages[index].get("text") or "")
            + "\n"
            + (pages[index].get("text_alt") or "")
        )

    p1, p2, p3, p4 = (page_text(index) for index in range(4))

    img1 = pages[0].get("image") if len(pages) > 0 else None
    img2 = pages[1].get("image") if len(pages) > 1 else None
    img3 = pages[2].get("image") if len(pages) > 2 else None

    full_text = "\n".join(page_text(index) for index in range(len(pages)))
    warnings = []
    auto_filled = []
    result = {key: None for key in ALL_FIELDS}

    # ---- Tax year
    year_match = (
        re.search(r"941\s+for\s+(20\d{2})", p1, re.I)
        or re.search(r"Calendar year\s*[\[\(]?\s*(20\d{2})", full_text, re.I)
        or re.search(r"\b(20\d{2})\b", p1)
    )

    result["tax_year"] = int(year_match.group(1)) if year_match else None

    # ---- Employer EIN
    ein = (
        zone_read(
            img2,
            (0.66, 0.95),
            0.081,
            kind="ein",
            whitelist="0123456789-",
        )
        if img2 is not None
        else None
    )

    partial = None

    if not ein:
        ein, partial = find_ein(p2, p4, p1)

    result["employer_ein"] = ein
    result["employer_ein_partial"] = partial

    if partial and not ein:
        warnings.append(
            f"Employer EIN only partly readable ({partial}); enter it manually."
        )
    elif not ein:
        warnings.append("Employer EIN could not be confidently extracted.")

    result["employer_name"] = find_employer_name(p2, p4)

    result["trade_name"] = find_first_match(
        r"Trade\s+name\s*\(if\s+any\)\s*[:|]?\s*\n*\s*"
        r"([A-Za-z0-9][A-Za-z0-9 &.,'/\-]{2,}?)\s*(?:\n|$)",
        p2,
    )

    # ---- Employer address / city / state / ZIP
    match = re.search(r"\bAddress\s*[:|]?\s*([^\n\[]+)", p1, re.I)

    if match:
        address = clean(match.group(1))
        address = re.sub(r"énd\b", "2nd", address, flags=re.I)
        address = re.sub(r"\s+", " ", address).strip(" -|,.")

        if address:
            result["employer_address"] = address

    match = re.search(
        r"\n\s*([A-Z][A-Za-z]+(?:\s+[A-Z][A-Za-z]+){0,2})\s*\|\s*"
        r"([A-Z]{2})\s*\|\s*"
        r"(\d{5}(?:-\d{4})?)\b",
        p1,
    )

    if match:
        result["employer_city"] = clean(match.group(1))
        result["employer_state"] = match.group(2)
        result["employer_zip"] = match.group(3)
    else:
        match = re.search(
            r"\n\s*([A-Z][A-Za-z]+(?:\s+[A-Z][A-Za-z]+){0,2})\s+"
            r"(\d{5}(?:-\d{4})?)\b",
            p1,
        )

        if match:
            result["employer_city"] = clean(match.group(1))
            result["employer_zip"] = match.group(2)

    if not result["employer_state"]:
        inferred = state_from_zip(result.get("employer_zip"))

        if inferred:
            result["employer_state"] = inferred
            auto_filled.append("employer_state")

    # ---- Source text
    result["source_page_count"] = len(pages)
    result["source_pages"] = [
        {
            "page": page["page"],
            "source": page.get("source", "unknown"),
            "text": page.get("text", ""),
            "text_alt": page.get("text_alt", ""),
            **(
                {"ocr_error": page["ocr_error"]}
                if page.get("ocr_error")
                else {}
            ),
        }
        for page in pages
    ]

    # ---- Single-amount lines
    searchable_text = p1 + "\n" + p2

    field_patterns = {
        "line2_wages": [
            r"Wages\s*,?\s*tips\s*,?\s*and\s*other\s*compensation",
        ],
        "line3_fed_tax": [
            r"Federal\s+income\s+tax\s+withheld",
        ],
        "line5e_total_ss_medicare": [
            r"Total\s+social\s+security\s+and\s+Medicare\s+taxes",
        ],
        "line5f_unreported_tips": [
            r"Section\s+3121\s*\(?q\)?\s*Notice\s+and\s+Demand",
        ],
        "line6_total_before_adj": [
            r"Total\s+taxes\s+before\s+adjustments",
            r"^\s*6\b",
        ],
        "line7_fractions_cents": [
            r"adjustment\s+for\s+fractions\s+of\s+cents",
        ],
        "line8_sick_pay_adj": [
            r"adjustment\s+for\s+sick\s+pay",
        ],
        "line9_tips_gtl_adj": [
            r"group[-\s]*term\s*life",
        ],
        "line10_total_after_adj": [
            r"Total\s+taxes\s+after\s+adjustments\.\s*Combine",
            r"^\s*10\b",
        ],
        "line11a_research_credit": [
            r"Qualified\s+small\s+business\s+payroll\s+tax\s+credit",
            r"research\s+activities",
        ],
        "line11b_nonrefundable_sick_family_before_apr2021": [
            r"Nonrefundable\s+portion\s+of\s+credit\s+for\s+qualified\s+sick"
            r"\s+and\s+family\s+leave\s+wages\s+for\s+leave\s+taken\s+before"
            r"\s+April\s+1,\s*2021",
        ],
        "line11d_nonrefundable_sick_family_after_mar2021": [
            r"Nonrefundable\s+portion\s+of\s+credit\s+for\s+qualified\s+sick"
            r"\s+and\s+family\s+leave\s+wages\s+for\s+leave\s+taken\s+after"
            r"\s+March\s+31,\s*2021",
        ],
        "line11g_total_nonrefundable": [
            r"Total\s+nonrefundable\s+credits",
        ],
        "line12_total_taxes": [
            r"^\s*12\b",
        ],
        "line13a_total_deposits": [
            r"Total\s+deposits\s+for\s+this\s+quarter",
            r"overpayments\s+applied\s+from\s+Form\s+941-X",
        ],
        "line13c_refundable_sick_family_before_apr2021": [
            r"Refundable\s+portion\s+of\s+credit\s+for\s+qualified\s+sick"
            r"\s+and\s+family\s+leave\s+wages\s+for\s+leave\s+taken\s+before"
            r"\s+April\s+1,\s*2021",
        ],
        "line13e_refundable_sick_family_after_mar2021": [
            r"Refundable\s+portion\s+of\s+credit\s+for\s+qualified\s+sick"
            r"\s+and\s+family\s+leave\s+wages\s+for\s+leave\s+taken\s+after"
            r"\s+March\s+31,\s*2021",
        ],
        "line13g_total_deposits_credits": [
            r"Total\s+deposits\s+and\s+refundable",
        ],
        "line14_balance_due": [
            r"Balance\s+due",
        ],
        "line15_overpayment": [
            r"Overpayment",
        ],
    }

    for key, patterns in field_patterns.items():
        value = find_labeled_amount(searchable_text, patterns)

        if value is not None:
            result[key] = value

    # ---- Line 1: number of employees
    line1_from_box = read_line1_employees(img1)

    def extract_line1_employees(text):
        match = re.search(
            r"Number\s+of\s+employees\s+who\s+received\s+wages"
            r"(.*?)"
            r"(?=\n\s*2\s+Wages|\Z)",
            text,
            re.I | re.S,
        )

        if not match:
            return None

        block = match.group(1)

        block = re.sub(
            r"\b(?:Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Sept|Oct|Nov|Dec)"
            r"[a-z]*\.?\s*\d{1,2}\b",
            " ",
            block,
            flags=re.I,
        )

        block = re.sub(r"Quarter\s*[1-4]?", " ", block, flags=re.I)
        numbers = re.findall(r"(?<![\d.])(\d{1,4})(?![\d.])", block)

        if not numbers:
            return None

        return int(numbers[-1])

    text_line1 = extract_line1_employees(p1)

    if (
        line1_from_box is not None
        and text_line1 is not None
        and line1_from_box != text_line1
    ):
        # The box reading comes from an isolated crop of the digit, so it
        # is more reliable than the whole-page text OCR (which misreads
        # 3 ↔ 4 / 5 ↔ 6 on these forms). Prefer it.
        warnings.append(
            f"Line 1 OCR conflict: box={line1_from_box}, "
            f"text={text_line1}. Using image (box) reading — "
            "verify against the original PDF."
        )
        line1 = line1_from_box
    else:
        line1 = (
            line1_from_box
            if line1_from_box is not None
            else text_line1
        )

    if line1 is not None:
        result["line1_employees"] = line1

    # ---- Line 4 checkbox
    if re.search(
        r"Check\s+and\s+go\s+to\s+line\s+6[^\n]*[Xx✓✔☒■]",
        p1,
        re.I,
    ):
        result["line4_no_wages"] = True

    # ---- Lines 5a-5d
    vals5a = labeled_amounts(
        p1,
        [r"Taxable\s+social\s+security\s+wages"],
    )
    vals5a_i = labeled_amounts(
        p1,
        [r"Qualified\s+sick\s+leave\s+wages"],
    )
    vals5a_ii = labeled_amounts(
        p1,
        [r"Qualified\s+family\s+leave\s+wages"],
    )
    vals5b = labeled_amounts(
        p1,
        [r"Taxable\s+social\s+security\s+tips"],
    )
    vals5c = labeled_amounts(
        p1,
        [r"Taxable\s+Medicare\s+wages"],
    )
    vals5d = labeled_amounts(
        p1,
        [r"Additional\s+Medicare\s+Tax\s+withholding"],
    )

    def split_cols(values):
        if len(values) >= 2:
            return values[0], values[-1]

        return None, (values[0] if values else None)

    (
        result["line5a_ss_wages"],
        result["line5a_ss_tax"],
    ) = split_cols(vals5a)

    (
        result["line5a_i_qualified_sick_leave_wages"],
        result["line5a_i_qualified_sick_leave_tax"],
    ) = split_cols(vals5a_i)

    (
        result["line5a_ii_qualified_family_leave_wages"],
        result["line5a_ii_qualified_family_leave_tax"],
    ) = split_cols(vals5a_ii)

    (
        result["line5b_ss_tips"],
        result["line5b_ss_tips_tax"],
    ) = split_cols(vals5b)

    (
        result["line5c_medicare_wages"],
        result["line5c_medicare_tax"],
    ) = split_cols(vals5c)

    (
        result["line5d_addl_medicare_wages"],
        result["line5d_addl_medicare_tax"],
    ) = split_cols(vals5d)

    # ---- Quarter
    for pattern in (
        r"(?<![A-Za-z])(?:[Xx✓✔☒■])\s*([1-4])\s*:\s*(?:January|April|July|October)",
        r"\[\s*[Xx✓✔☒]\s*\]\s*([1-4])\s*:\s*(?:January|April|July|October)",
    ):
        match = re.search(pattern, p1, re.I)

        if match:
            result["quarter"] = int(match.group(1))
            break

    if result["quarter"] is None and filename:
        match = re.search(
            r"(?:^|[^A-Za-z0-9])(?:([1-4])[qQ](?:\d{2})?|[qQ]([1-4]))(?![A-Za-z0-9])",
            filename,
        )

        if match:
            result["quarter"] = int(match.group(1) or match.group(2))
            auto_filled.append("quarter")

    # ---- Liability
    liab_p2 = find_labeled_amount(
        p2,
        [r"Total\s+liability\s+for\s+quarter"],
    )

    tail = re.search(
        r"Total\s+liability\s+for\s+the\s+quarter(.*)",
        p4,
        re.I | re.S,
    )

    p4_values = amounts_in(tail.group(1)) if tail else []
    liab_p4 = p4_values[0] if p4_values else None

    result["total_liability"] = (
        liab_p2 if liab_p2 is not None else liab_p4
    )

    month_candidates = []

    for number in (1, 2, 3):
        candidates = []

        value = find_labeled_amount(
            p2,
            [rf"Month\s*{number}\b"],
            occurrence=0,
        )

        if value is not None:
            candidates.append(value)

        segment = re.search(
            rf"Month\s*{number}\b(.*?)(?=Month\s*[123]\b|Total\s+liability|$)",
            p4,
            re.I | re.S,
        )

        if segment:
            candidates.extend(amounts_in(segment.group(1))[:2])

        month_candidates.append(list(dict.fromkeys(candidates)))

    for number, candidates in enumerate(month_candidates, start=1):
        result[f"month{number}_liability"] = (
            candidates[0] if candidates else None
        )

    total = result["total_liability"]

    if total is not None and all(month_candidates):
        current_sum = sum(
            result[f"month{number}_liability"]
            for number in (1, 2, 3)
        )

        if abs(current_sum - total) > 0.02:
            for combo in itertools.product(*month_candidates):
                if abs(sum(combo) - total) <= 0.02:
                    for number, value in enumerate(combo, start=1):
                        result[f"month{number}_liability"] = value
                    break

    month_keys = [
        f"month{number}_liability"
        for number in (1, 2, 3)
    ]

    missing_months = [
        key for key in month_keys
        if result[key] is None
    ]

    if len(missing_months) == 1 and total is not None:
        known = sum(
            result[key]
            for key in month_keys
            if result[key] is not None
        )

        result[missing_months[0]] = round(total - known, 2)
        auto_filled.append(missing_months[0])

    # ---- Totals consensus
    no_adj = all(
        result.get(key) is None
        for key in (
            "line7_fractions_cents",
            "line8_sick_pay_adj",
            "line9_tips_gtl_adj",
            "line11a_research_credit",
            "line11b_nonrefundable_sick_family_before_apr2021",
            "line11d_nonrefundable_sick_family_after_mar2021",
            "line11g_total_nonrefundable",
        )
    )

    votes = [
        value
        for value in (
            result["line6_total_before_adj"],
            result["line10_total_after_adj"],
            result["line12_total_taxes"],
            liab_p2,
            liab_p4,
        )
        if value is not None
    ]

    if no_adj and votes:
        value, count = Counter(votes).most_common(1)[0]

        if count >= 2:
            for key in (
                "line6_total_before_adj",
                "line10_total_after_adj",
                "line12_total_taxes",
                "total_liability",
            ):
                if result[key] != value:
                    result[key] = value
                    auto_filled.append(key)

    # ---- Line 3
    e5 = result["line5e_total_ss_medicare"]
    l6 = result["line6_total_before_adj"]

    if l6 is not None and e5 is not None:
        derived3 = round(
            l6 - e5 - (result["line5f_unreported_tips"] or 0),
            2,
        )

        read3 = result["line3_fed_tax"]

        if derived3 >= 0 and (
            read3 is None or abs(read3 - derived3) > 0.02
        ):
            result["line3_fed_tax"] = derived3
            auto_filled.append("line3_fed_tax")

    # ---- Reconcile wages
    def reconcile_wages():
        if e5 is None:
            return None

        tax5a = result["line5a_ss_tax"]

        candidates = [
            result["line2_wages"],
            result["line5a_ss_wages"],
            result["line5c_medicare_wages"],
            round(tax5a / 0.124, 2) if tax5a else None,
            round(e5 / 0.153, 2),
        ]

        for candidate in candidates:
            if candidate and abs(round(candidate * 0.153, 2) - e5) <= 0.05:
                return candidate

        return None

    wages = reconcile_wages()

    if wages is not None:
        if result["line5a_ss_wages"] != wages:
            result["line5a_ss_wages"] = wages
            auto_filled.append("line5a_ss_wages")

        tax5a = round(wages * 0.124, 2)

        if result["line5a_ss_tax"] != tax5a:
            result["line5a_ss_tax"] = tax5a
            auto_filled.append("line5a_ss_tax")

        if result["line5c_medicare_wages"] != wages:
            result["line5c_medicare_wages"] = wages
            auto_filled.append("line5c_medicare_wages")

        tax5c = round(wages * 0.029, 2)

        if result["line5c_medicare_tax"] != tax5c:
            result["line5c_medicare_tax"] = tax5c
            auto_filled.append("line5c_medicare_tax")

        line2 = result["line2_wages"]

        if line2 is None or line2 < wages - 0.01:
            result["line2_wages"] = wages
            auto_filled.append("line2_wages")

    elif e5 is not None and (vals5a or vals5c):
        warnings.append(
            "Lines 5a / 5c could not be matched to line 5e; "
            "check them against the PDF."
        )

    if result["line2_wages"] is None:
        warnings.append("Line 2 wages could not be read; enter it manually.")

    # ---- Signer / designee / preparer from image zones
    zones = read_page3_zones(img3)

    for key, value in zones.items():
        if value is not None:
            result[key] = value

    # ---- Text-based fallbacks (page 3)
    signer_match = re.search(
        r"Sign\s+here(.*?)(?=Paid\s+Preparer\s+Use\s+Only|$)",
        p3,
        re.I | re.S,
    )

    signer_text = signer_match.group(1) if signer_match else p3

    prep_match = re.search(
        r"Paid\s+Preparer\s+Use\s+Only(.*)",
        p3,
        re.I | re.S,
    )

    prep_text = prep_match.group(1) if prep_match else p3

    designee_match = re.search(
        r"May\s+we\s+speak\s+with\s+your\s+third-party\s+designee(.*?)(?=Sign\s+here|$)",
        p3,
        re.I | re.S,
    )

    designee_text = designee_match.group(1) if designee_match else ""

    # ----- signer_name
    if not result["signer_name"]:
        result["signer_name"] = find_first_match(
            r"Print\s+your\s*\n\s*([A-Z][A-Za-z.'\- ]{4,})\s*\n",
            signer_text,
        )

    # ----- signer_title
    if not result["signer_title"]:
        match = re.search(
            r"Print\s+your\s+name\s+and\s+title\s*\n*\s*"
            r"([A-Z][A-Za-z.'\- ]{4,})\s*\n*\s*([A-Za-z][A-Za-z .'\-]{2,})",
            p3,
            re.I,
        )

        if match:
            title = clean(match.group(2))

            if title and title.lower() not in ("sign", "date"):
                result["signer_title"] = title

    # ----- signer_phone
    if not result["signer_phone"]:
        result["signer_phone"] = (
            extract_phone(signer_text)
            or extract_phone(p3)
        )

    # ----- designee_name
    if not result["designee_name"]:
        result["designee_name"] = find_first_match(
            r"Designee.?s\s+name\s+and\s+phone\s+number\s+"
            r"([A-Z][A-Za-z.'\- ]+?)\s+\d",
            designee_text or p3,
        )

    # ----- designee_phone
    if not result["designee_phone"]:
        match = re.search(
            r"Designee.?s\s+name\s+and\s+phone\s+number"
            r".*?(\d{3})\D{0,3}(\d{3})\D{0,3}(\d{4})",
            designee_text or p3,
            re.I | re.S,
        )

        if match:
            result["designee_phone"] = (
                f"{match.group(1)}-{match.group(2)}-{match.group(3)}"
            )
        else:
            result["designee_phone"] = (
                extract_phone(designee_text)
                or extract_phone(p3)
            )

    # ----- preparer_name
    if not result["preparer_name"]:
        result["preparer_name"] = find_first_match(
            r"Preparer.?s\s*name\s*\|?\s*([A-Z][A-Za-z.'\- ]+?)\s+PTIN",
            prep_text,
        )

    # ----- preparer_ptin
    if not result["preparer_ptin"]:
        match = re.search(r"\bP\s?(\d[\d ]{7,9}\d)\b", prep_text)

        if match:
            digits = re.sub(r"\D", "", match.group(1))

            if len(digits) == 8:
                result["preparer_ptin"] = f"P{digits}"

    # ----- preparer_firm_name
    if not result["preparer_firm_name"]:
        result["preparer_firm_name"] = find_first_match(
            r"([A-Z][A-Z&.,' ]{5,}?(?:INC|LLC|CORP|CPA|SERVICES)\.?)\s+EIN",
            prep_text,
        )

    # ----- preparer_firm_ein
    if not result["preparer_firm_ein"]:
        ein2, _ = find_ein(prep_text)

        if ein2:
            result["preparer_firm_ein"] = ein2

    # ----- preparer_address
    if not result["preparer_address"]:
        match = re.search(
            r"(?:Address|waaress|nwaaress|nwaress|nadd?ress|nadie|ndie|ddress|Addr)"
            r"\s*[:|.]?\s*"
            r"([A-Za-z0-9][^\n|]*?)"
            r"(?=\s*(?:Phone|City|State|ZIP|Page\s*3|Form\s*941|$))",
            prep_text,
            re.I,
        )

        if match:
            address = clean(match.group(1).strip(" -|."))
            address = re.sub(r"\s*Phone\s*$", "", address, flags=re.I)

            if address:
                result["preparer_address"] = address

    # ----- preparer_address PO Box fallback
    if not result["preparer_address"]:
        match = re.search(
            r"\b((?:P\.?\s*O\.?|PO)\s*Box\s+\d+[A-Za-z0-9 \-]*)",
            prep_text,
            re.I,
        )

        if match:
            address = clean(match.group(1).strip(" -|."))

            if address:
                result["preparer_address"] = address

    # ----- preparer_address street fallback
    if not result["preparer_address"]:
        match = re.search(
            r"\n\s*(\d+\s+[A-Z][A-Za-z0-9.'\- ]{3,}(?:"
            r"Street|St|Avenue|Ave|Road|Rd|Boulevard|Blvd|Lane|Ln|Drive|Dr|"
            r"Court|Ct|Place|Pl|Way|Terrace|Ter|Circle|Cir|Suite|Ste|Floor|Fl)"
            r"\.?[^\n|]*)",
            prep_text,
        )

        if match:
            address = clean(match.group(1).strip(" -|."))

            if address:
                result["preparer_address"] = address

    # ----- preparer_phone
    if not result["preparer_phone"]:
        result["preparer_phone"] = (
            extract_phone(prep_text)
            or extract_phone(p3)
        )

    # ----- preparer_city
    if not result["preparer_city"]:
        result["preparer_city"] = find_first_match(
            r"City\s+([A-Z][A-Za-z.'\- ]+?)\s+State",
            prep_text,
        )

    # ----- preparer_state
    if not result["preparer_state"]:
        for match in re.finditer(r"State\s+([A-Za-z]{2})\b", prep_text):
            state = match.group(1).upper()

            if state in US_STATES:
                result["preparer_state"] = state
                break

    # ----- preparer_zip
    if not result["preparer_zip"]:
        match = re.search(
            r"ZIP\s*(?:code)?\s*[:|]?\s*(\d{5}(?:-\d{4})?)",
            prep_text,
            re.I,
        )

        if match:
            result["preparer_zip"] = match.group(1)

    # ---- Infer preparer state from ZIP
    zip_inferred = state_from_zip(result.get("preparer_zip"))

    if zip_inferred and result["preparer_state"] != zip_inferred:
        result["preparer_state"] = zip_inferred
        auto_filled.append("preparer_state")

    # ----- signer_name fallback
    if not result["signer_name"]:
        fallback = result["preparer_name"] or result["designee_name"]

        if fallback:
            result["signer_name"] = fallback
            auto_filled.append("signer_name")

    # ----- signer_phone fallback
    if not result["signer_phone"]:
        result["signer_phone"] = (
            result["preparer_phone"]
            or result["designee_phone"]
        )

        if result["signer_phone"]:
            auto_filled.append("signer_phone")

    # ---- Consistency checks
    def check_sum(keys, target, message):
        values = [result.get(key) for key in keys]
        expected = result.get(target)

        if expected is not None and all(value is not None for value in values):
            if abs(sum(values) - expected) > 0.02:
                warnings.append(message)

    check_sum(
        [
            "line3_fed_tax",
            "line5e_total_ss_medicare",
            "line5f_unreported_tips",
        ],
        "line6_total_before_adj",
        "Line 6 does not match lines 3 + 5e + 5f.",
    )

    check_sum(
        [
            "month1_liability",
            "month2_liability",
            "month3_liability",
        ],
        "total_liability",
        "Monthly liabilities do not match total liability.",
    )

    if not result["employer_name"]:
        warnings.append("Employer name could not be confidently extracted.")

    if result["quarter"] is None:
        warnings.append("Quarter checkbox needs verification.")

    if result["tax_year"] is None:
        warnings.append("Tax year could not be confidently extracted.")

    if any(page.get("ocr_error") for page in pages):
        warnings.append("One or more pages could not be OCR-processed.")

    if not full_text.strip():
        warnings.append("No readable PDF text was extracted.")

    if any(page.get("source") == "ocr" for page in pages):
        warnings.append(
            "OCR-derived values require review against the original PDF."
        )

    result["parse_warnings"] = warnings
    result["auto_filled"] = sorted(set(auto_filled))
    result["review_required"] = bool(warnings)

    return result


# Keep FastAPI routes in app.py.
# This module contains Form 941 parsing logic only.