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
from fastapi import FastAPI, File, HTTPException, UploadFile
from fastapi.concurrency import run_in_threadpool
from fastapi.middleware.cors import CORSMiddleware


app = FastAPI(
    title="OLPMS Form 941 Python Parser",
    version="2.1.0",
    root_path="/olpms-parser-api",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        "http://localhost:5173",
        "http://127.0.0.1:5173",
    ],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

TESSERACT_PATH = os.getenv(
    "TESSERACT_CMD",
    r"C:\Program Files\Tesseract-OCR\tesseract.exe",
)

if os.path.isfile(TESSERACT_PATH):
    pytesseract.pytesseract.tesseract_cmd = TESSERACT_PATH

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": "MN", "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(m["whole"], m["cents"]) for m in AMT_RE.finditer(text or "")]
    return [v for v in values if v 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 m in EIN_LOOSE.finditer(text or ""):
            digits = m.group(2).translate(EIN_FIX)
            ein = f"{m.group(1)}-{digits}"
            if re.fullmatch(r"\d{7}", digits):
                return ein, None
            partial = partial or ein
    return None, partial


# Legal-suffix list used to trim trailing OCR noise off the employer name
# (e.g. "SANATECH PRINTING CORPORATION Oo" -> "SANATECH PRINTING CORPORATION").
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):
    """Remove trailing OCR noise from an employer name.

    If the name contains a legal suffix (CORPORATION, INC, LLC, …), cut
    everything after the LAST one. Otherwise, strip a trailing 1-2 letter
    token that is almost certainly OCR garbage.
    """
    if not value:
        return value

    v = value.strip()

    matches = list(LEGAL_SUFFIX_RE.finditer(v))
    if matches:
        last = matches[-1]
        trimmed = v[:last.end()].strip(" .,-")
        if trimmed:
            return trimmed

    trimmed = re.sub(r"\s+[A-Za-z]{1,2}$", "", v).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:
        m = re.search(pattern, text or "")
        if m:
            name = clean(m.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
    norm = re.sub(r"[=~_—–]", "-", text)
    m = PHONE_RE.search(norm)
    if not m:
        return None
    return f"{m.group(1)}-{m.group(2)}-{m.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.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.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":
            m = EIN_LOOSE.search(raw)
            if m:
                digits = m.group(2).translate(EIN_FIX)
                if digits.isdigit():
                    value = f"{m.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":
            m = re.search(r"\b(\d{5}(?:-\d{4})?)\b", raw)
            value = m.group(1) if m else None
        elif kind == "state":
            m = re.search(r"\b([A-Z]{2})\b", raw.upper())
            if m and m.group(1) in US_STATES:
                value = m.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


# Line 1's answer box on page 1.
PAGE1_LINE1_X = (0.725, 0.94)
PAGE1_LINE1_Y = 0.365


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

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

    for dy in (0.0, -0.010, 0.010, -0.020, 0.020):
        yc = PAGE1_LINE1_Y + dy
        crop = image.crop((
            max(0, int(width * x0)),
            max(0, int(height * (yc - 0.015))),
            min(width, int(width * x1)),
            min(height, int(height * (yc + 0.015))),
        ))
        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)
            for psm in (10, 13, 8):
                try:
                    raw = pytesseract.image_to_string(
                        enlarged,
                        config=f"--psm {psm} -c tessedit_char_whitelist=0123456789",
                    )
                except Exception:
                    continue
                digits = re.sub(r"\D", "", raw)
                if len(digits) == 1:
                    votes[int(digits)] += 1

    if not votes:
        return None
    ranked = votes.most_common()
    if len(ranked) > 1 and ranked[0][1] == ranked[1][1]:
        return None
    value, count = ranked[0]
    return value if count >= 3 else None


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(i):
        if i >= len(pages):
            return ""
        return (pages[i].get("text") or "") + "\n" + (pages[i].get("text_alt") or "")

    p1, p2, p3, p4 = (page_text(i) for i 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(i) for i in range(len(pages)))
    warnings = []
    auto_filled = []
    result = {}

    for key in ALL_FIELDS:
        result.setdefault(key, None)

    # ---- 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 (page 1)
    m = re.search(
        r"\bAddress\s*[:|]?\s*([^\n\[]+)",
        p1,
        re.I,
    )
    if m:
        addr = clean(m.group(1))
        addr = re.sub(r"énd\b", "2nd", addr, flags=re.I)
        addr = re.sub(r"\s+", " ", addr).strip(" -|,.")
        if addr:
            result["employer_address"] = addr

    m = 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 m:
        result["employer_city"] = clean(m.group(1))
        result["employer_state"] = m.group(2)
        result["employer_zip"] = m.group(3)
    else:
        m = 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 m:
            result["employer_city"] = clean(m.group(1))
            result["employer_zip"] = m.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):
        m = 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 m:
            return None

        block = m.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)

        nums = re.findall(r"(?<![\d.])(\d{1,4})(?![\d.])", block)
        if not nums:
            return None
        return int(nums[-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
    ):
        warnings.append(
            f"Line 1 OCR conflict: box={line1_from_box}, text={text_line1}. "
            "Verify the PDF."
        )
        line1 = text_line1
    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)",
    ):
        m = re.search(pattern, p1, re.I)
        if m:
            result["quarter"] = int(m.group(1))
            break

    if result["quarter"] is None and filename:
        m = re.search(
            r"(?:^|[^A-Za-z0-9])(?:([1-4])[qQ](?:\d{2})?|[qQ]([1-4]))(?![A-Za-z0-9])",
            filename,
        )
        if m:
            result["quarter"] = int(m.group(1) or m.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_vals = amounts_in(tail.group(1)) if tail else []
    liab_p4 = p4_vals[0] if p4_vals else None
    result["total_liability"] = liab_p2 if liab_p2 is not None else liab_p4

    month_candidates = []
    for n in (1, 2, 3):
        cands = []
        v = find_labeled_amount(p2, [rf"Month\s*{n}\b"], occurrence=0)
        if v is not None:
            cands.append(v)
        seg = re.search(
            rf"Month\s*{n}\b(.*?)(?=Month\s*[123]\b|Total\s+liability|$)",
            p4, re.I | re.S,
        )
        if seg:
            cands.extend(amounts_in(seg.group(1))[:2])
        month_candidates.append(list(dict.fromkeys(cands)))

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

    total = result["total_liability"]
    if total is not None and all(month_candidates):
        current_sum = sum(result[f"month{n}_liability"] for n 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 n, value in enumerate(combo, start=1):
                        result[f"month{n}_liability"] = value
                    break

    month_keys = [f"month{n}_liability" for n in (1, 2, 3)]
    missing_months = [k for k in month_keys if result[k] is None]
    if len(missing_months) == 1 and total is not None:
        known = sum(result[k] for k in month_keys if result[k] is not None)
        result[missing_months[0]] = round(total - known, 2)
        auto_filled.append(missing_months[0])

    # ---- Totals consensus
    no_adj = all(
        result.get(k) is None
        for k 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 = [
        v for v in (
            result["line6_total_before_adj"],
            result["line10_total_after_adj"],
            result["line12_total_taxes"],
            liab_p2, liab_p4,
        ) if v 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 c in candidates:
            if c and abs(round(c * 0.153, 2) - e5) <= 0.05:
                return c
        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"]:
        m = 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 m:
            title = clean(m.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"]:
        m = 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 m:
            result["designee_phone"] = f"{m.group(1)}-{m.group(2)}-{m.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"]:
        m = re.search(r"\bP\s?(\d[\d ]{7,9}\d)\b", prep_text)
        if m:
            digits = re.sub(r"\D", "", m.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"]:
        m = 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 m:
            addr = clean(m.group(1).strip(" -|."))
            addr = re.sub(r"\s*Phone\s*$", "", addr, flags=re.I)
            if addr:
                result["preparer_address"] = addr

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

    # ----- preparer_address street fallback
    if not result["preparer_address"]:
        m = 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 m:
            addr = clean(m.group(1).strip(" -|."))
            if addr:
                result["preparer_address"] = addr

    # ----- 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 m in re.finditer(r"State\s+([A-Za-z]{2})\b", prep_text):
            st = m.group(1).upper()
            if st in US_STATES:
                result["preparer_state"] = st
                break

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

    # =========================================================
    # FINAL CROSS-CHECK: infer state from ZIP
    # =========================================================
    zip_inferred = state_from_zip(result.get("preparer_zip"))
    if zip_inferred:
        if 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(k) for k in keys]
        expected = result.get(target)
        if expected is not None and all(v is not None for v 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(p.get("ocr_error") for p 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(p.get("source") == "ocr" for p 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


# ---------------------------------------------------------------------------
# Routes
# ---------------------------------------------------------------------------

@app.get("/health")
def health():
    return {"status": "ok", "service": "olpms-941-python-parser"}


@app.post("/parse-941")
async def parse_941_endpoint(file: UploadFile = File(...)):
    if not file.filename or not file.filename.lower().endswith(".pdf"):
        raise HTTPException(status_code=400, detail="Please upload a PDF file.")

    pdf_bytes = await file.read()
    if not pdf_bytes:
        raise HTTPException(status_code=400, detail="The uploaded PDF is empty.")
    if len(pdf_bytes) > 25 * 1024 * 1024:
        raise HTTPException(
            status_code=413,
            detail="PDF is too large. Maximum allowed size is 25 MB.",
        )

    try:
        pages = await run_in_threadpool(extract_text, pdf_bytes)
        if not pages:
            raise HTTPException(status_code=422, detail="The PDF contains no readable pages.")

        result = await run_in_threadpool(parse_941, pages, file.filename)
        result["filename"] = file.filename
        return result
    except HTTPException:
        raise
    except Exception as exc:
        raise HTTPException(
            status_code=422,
            detail=f"Unable to parse this Form 941 PDF: {exc}",
        ) from exc