import re


W2_FIELDS = {
    "tax_year": None,
    "control_number": None,
    "employer_ein": None,
    "employer_name": None,
    "employer_address": None,
    "employee_ssn": None,
    "employee_first_name": None,
    "employee_middle_initial": None,
    "employee_last_name": None,
    "employee_suffix": None,
    "employee_address": None,
    "box1_wages": None,
    "box2_fed_tax": None,
    "box3_ss_wages": None,
    "box4_ss_tax": None,
    "box5_medicare_wages": None,
    "box6_medicare_tax": None,
    "box7_ss_tips": None,
    "box8_allocated_tips": None,
    "box9": None,
    "box10_dependent_care": None,
    "box11_nonqualified": None,
    "box12a": None,
    "box12b": None,
    "box12c": None,
    "box12d": None,
    "box13_statutory_employee": False,
    "box13_retirement_plan": False,
    "box13_third_party_sick_pay": False,
    "box14_other": None,
    "box14a_other": None,
    "box14b_tipped_occupation_codes": None,
    "box15_state": None,
    "box15_state_id": None,
    "box16_state_wages": None,
    "box17_state_tax": None,
    "box18_local_wages": None,
    "box19_local_tax": None,
    "box20_locality_name": None,
}

MONEY_RE = re.compile(
    r"(?<![A-Za-z0-9])"
    r"\$?\s*\(?-?"
    r"(?:\d{1,3}(?:,\d{3})+|\d+)"
    r"\.\d{2}"
    r"\)?"
    r"(?![A-Za-z0-9])"
)

EIN_RE = re.compile(r"(?<!\d)(\d{2})\s*-\s*(\d{7})(?!\d)")
SSN_RE = re.compile(
    r"(?<!\d)(\d{3})\s*-\s*(\d{2})\s*-\s*(\d{4})(?!\d)"
)


def _clean(value):
    if value is None:
        return None
    value = re.sub(r"\s+", " ", str(value)).strip(" \t:|,")
    return value or None


def _amount(value):
    if value is None:
        return None
    value = str(value).strip()
    value = value.replace("$", "").replace(",", "").replace(" ", "")
    negative = value.startswith("(") and value.endswith(")")
    value = value.strip("()")
    try:
        number = round(float(value), 2)
        return -number if negative else number
    except (ValueError, TypeError):
        return None


def _money_values(line):
    return [
        _amount(match.group(0))
        for match in MONEY_RE.finditer(line or "")
    ]


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


def _normalized_lines(text):
    return [line.rstrip() for line in (text or "").splitlines()]


def _line_index(lines, pattern):
    for index, line in enumerate(lines):
        if re.search(pattern, line, re.I):
            return index
    return None


def _first_column(line):
    """Keep the left-hand W-2 copy in side-by-side PDF layouts."""
    if not line:
        return None
    stripped = line.lstrip()
    if not stripped:
        return None
    parts = re.split(r"\s{2,}", stripped, maxsplit=1)
    return _clean(parts[0]) if parts and parts[0] else None


def _section_lines(lines, heading_pattern, stop_pattern, limit=10):
    index = _line_index(lines, heading_pattern)
    if index is None:
        return []
    output = []
    for raw_line in lines[index + 1:index + 1 + limit]:
        if not raw_line.strip():
            continue
        if re.search(stop_pattern, raw_line, re.I):
            break
        value = _first_column(raw_line)
        if value:
            output.append(value)
    return output


def _next_nonempty(lines, index, limit=5):
    if index is None:
        return []
    found = []
    for line in lines[index + 1:index + 1 + limit]:
        if line.strip():
            value = _first_column(line)
            if value:
                found.append(value)
    return found


def _values_below_label(lines, label_pattern, count=2, limit=4):
    index = _line_index(lines, label_pattern)
    if index is None:
        return []
    for line in lines[index + 1:index + 1 + limit]:
        values = _money_values(line)
        if len(values) >= count:
            return values[:count]
    return []


def _set_pair(result, key1, key2, values):
    if len(values) >= 1:
        result[key1] = values[0]
    if len(values) >= 2:
        result[key2] = values[1]


def _find_text_after_heading(lines, heading_pattern, limit=5):
    index = _line_index(lines, heading_pattern)
    values = _next_nonempty(lines, index, limit)
    return values[0] if values else None


def _find_address_after_heading(lines, heading_pattern, limit=6):
    index = _line_index(lines, heading_pattern)
    if index is None:
        return None
    parts = []
    stop_pattern = (
        r"^(?:[a-e]\s+)?"
        r"(?:Employee(?:'s)? social security|"
        r"Employer ID|Employer identification|"
        r"7\s+Social security tips|"
        r"10\s+Dependent care|"
        r"13\s+Statutory employee|"
        r"15\s+State|"
        r"Form W-2)\b"
    )
    for raw_line in lines[index + 1:index + 1 + limit]:
        if not raw_line.strip():
            continue
        if re.search(stop_pattern, raw_line, re.I):
            break
        value = _first_column(raw_line)
        if value:
            parts.append(value)
    return _clean(" ".join(parts)) if parts else None


def _first_ein(text):
    match = EIN_RE.search(text or "")
    if not match:
        return None
    return f"{match.group(1)}-{match.group(2)}"


def _first_ssn(text):
    match = SSN_RE.search(text or "")
    if not match:
        return None
    return f"{match.group(1)}-{match.group(2)}-{match.group(3)}"


def _control_number(text):
    match = re.search(r"\b([A-Z]{2}\d{7})\b", text or "", re.I)
    if match:
        return match.group(1)
    lines = _normalized_lines(text)
    index = _line_index(lines, r"\bControl\s+number\b")
    if index is not None:
        for line in lines[index:index + 4]:
            candidates = re.findall(r"\b[A-Z0-9][A-Z0-9-]{3,}\b", line, re.I)
            for value in candidates:
                if (
                    re.search(r"[A-Z]", value, re.I)
                    and re.search(r"\d", value)
                    and value.lower() not in {
                        "control", "number", "wages", "federal",
                        "income", "tax", "withheld",
                    }
                ):
                    return value
    return None


def _employee_name(text):
    lines = _normalized_lines(text)
    name = _find_text_after_heading(
        lines,
        r"\be\s+Employee(?:'s)?\s+name",
        limit=4,
    )
    if not name:
        return None, None, None, None
    parts = name.split()
    if not parts:
        return None, None, None, None
    first = parts[0]
    middle = None
    suffix = None
    if parts[-1].rstrip(".").upper() in {"JR", "SR", "II", "III", "IV", "V"}:
        suffix = parts.pop().rstrip(".")
    if len(parts) >= 3 and len(parts[1].rstrip(".")) == 1:
        middle = parts[1].rstrip(".")
        last = " ".join(parts[2:])
    elif len(parts) >= 2:
        last = " ".join(parts[1:])
    else:
        last = None
    return first, middle, last, suffix


# ★ FIX: accept BOTH state-ID shapes:
#     NY style:  XX-XXXXXXX   (e.g. 81-4376648)
#     NJ style:  XXXXXXXXXXXX (e.g. 842395406000)
#   4–21 chars of [A-Z0-9-] starting with an alphanumeric.
def _state_row(text):
    lines = _normalized_lines(text)
    for line in lines:
        match = re.match(
            r"^\s*([A-Z]{2})\s+"
            r"([A-Z0-9][A-Z0-9\-]{3,20})"
            r"\s+(.*)$",
            line,
            re.I,
        )
        if not match:
            continue
        state = match.group(1).upper()
        state_id = re.sub(r"\s+", "", match.group(2))
        values = _money_values(match.group(3))
        if len(values) >= 2:
            return state, state_id, values[0], values[1]
    return None, None, None, None


def _local_row(lines):
    values = _values_below_label(
        lines,
        r"\b18\s+Local\s+wages",
        count=2,
        limit=4,
    )
    locality = None
    index = _line_index(lines, r"\b18\s+Local\s+wages")
    if index is not None:
        non_empty_count = 0
        for line in lines[index + 1:index + 4]:
            if not line.strip():
                continue
            non_empty_count += 1
            if non_empty_count > 2:
                break
            segments = [s for s in re.split(r"\s{2,}", line.strip()) if s]
            for seg in segments:
                if not re.fullmatch(
                    r"[A-Z][A-Za-z]+(?:\s+[A-Z][A-Za-z]+){0,2}", seg
                ):
                    continue
                if re.search(
                    r"Copy|Federal|Form|Wage|Statement|Internal|Revenue|Service|"
                    r"Furnish|This|Information|Being|Local|Income|Tax|Locality|Name",
                    seg, re.I,
                ):
                    continue
                locality = seg
                break
            if locality:
                break
    return values, locality


def parse_w2(pages, filename=None):
    result = dict(W2_FIELDS)
    page_texts = [
        (page.get("text") or "") + "\n" + (page.get("text_alt") or "")
        for page in pages
    ]
    text = "\n".join(page_texts)
    lines = _normalized_lines(text)
    warnings = []

    if not text.strip():
        warnings.append("No readable W-2 text was extracted.")

    # Tax year
    year = _find(
        r"\b(?:Tax\s+Year|Calendar\s+Year)\s*[: ]*\s*(20\d{2})\b",
        text,
    )
    if not year:
        year = _find(
            r"\bForm\s+W-2\b[^\n]{0,100}?\b(20\d{2})\b",
            text,
        )
    if not year and filename:
        year = _find(r"\b(20\d{2})\b", filename)
    result["tax_year"] = int(year) if year else None

    # Control number, EIN, SSN
    result["control_number"] = _control_number(text)
    result["employer_ein"] = _first_ein(text)
    result["employee_ssn"] = _first_ssn(text)

    # Employer name and address: left-hand copy only
    employer_lines = _section_lines(
        lines,
        r"\bc\s+Employer(?:'s)?\s+name",
        r"^\s*d\s+Employee(?:'s)?\s+social security",
        limit=8,
    )
    result["employer_name"] = employer_lines[0] if employer_lines else None
    result["employer_address"] = (
        _clean(" ".join(employer_lines[1:]))
        if len(employer_lines) > 1
        else None
    )

    if not result["employer_name"]:
        result["employer_name"] = _find_text_after_heading(
            lines,
            r"\bc\s+Employer(?:'s)?\s+name",
            limit=4,
        )
    if not result["employer_address"]:
        result["employer_address"] = _find_address_after_heading(
            lines,
            r"\bc\s+Employer(?:'s)?\s+name",
            limit=6,
        )

    # Employee name and address
    (
        result["employee_first_name"],
        result["employee_middle_initial"],
        result["employee_last_name"],
        result["employee_suffix"],
    ) = _employee_name(text)

    employee_lines = _section_lines(
        lines,
        r"\be\s+Employee(?:'s)?\s+name",
        r"^\s*7\s+Social security tips",
        limit=8,
    )
    result["employee_address"] = (
        _clean(" ".join(employee_lines[1:]))
        if len(employee_lines) > 1
        else None
    )
    if not result["employee_address"]:
        result["employee_address"] = _find_address_after_heading(
            lines,
            r"\be\s+Employee(?:'s)?\s+name",
            limit=6,
        )

    # Federal boxes
    federal_pairs = [
        ("box1_wages", "box2_fed_tax", r"\b1\s+Wages,\s*tips,\s*other\s+comp"),
        ("box3_ss_wages", "box4_ss_tax", r"\b3\s+Social\s+security\s+wages"),
        ("box5_medicare_wages", "box6_medicare_tax", r"\b5\s+Medicare\s+wages\s+and\s+tips"),
        ("box7_ss_tips", "box8_allocated_tips", r"\b7\s+Social\s+security\s+tips"),
        ("box10_dependent_care", "box11_nonqualified", r"\b10\s+Dependent\s+care\s+benefits"),
    ]
    for first_key, second_key, label_pattern in federal_pairs:
        values = _values_below_label(lines, label_pattern, count=2, limit=4)
        _set_pair(result, first_key, second_key, values)

    result["box9"] = None

    # Box 12 codes
    for letter in ("a", "b", "c", "d"):
        match = re.search(
            rf"\b12\s*{letter}\s+Code\s+([A-Z])\s*"
            rf"\$?\s*([\d,]+\.\d{{2}})",
            text,
            re.I,
        )
        if match:
            result[f"box12{letter}"] = (
                f"{match.group(1).upper()} "
                f"{match.group(2).replace(',', '')}"
            )

    # Box 13 checkboxes
    checkbox_labels = [
        ("box13_statutory_employee", r"Statutory\s+employee"),
        ("box13_retirement_plan", r"Retirement\s+plan"),
        ("box13_third_party_sick_pay", r"Third[- ]party\s+sick\s+pay"),
    ]
    for key, label_pattern in checkbox_labels:
        result[key] = bool(
            re.search(
                label_pattern
                + r"[^\n]{0,35}?"
                + r"(?:\[\s*[xX]\s*\]|☒|✓|✔)",
                text,
                re.I,
            )
        )

    # Box 14 — accept state-specific labels in addition to SDI / PFL
    other_values = []
    for label in ("SDI", "PFL", "FLI", "UI/WF/HC"):
        match = re.search(
            rf"(?<![A-Za-z]){re.escape(label)}"
            rf"\s*:\s*\$?\s*([\d,]+(?:\.\d{{1,2}})?)",
            text,
            re.I,
        )
        if match:
            amount = match.group(1).replace(",", "")
            if "." not in amount:
                amount = f"{amount}.00"
            else:
                whole, frac = amount.split(".", 1)
                amount = f"{whole}.{frac.ljust(2, '0')[:2]}"
            other_values.append(f"{label}: {amount}")
    result["box14_other"] = "; ".join(other_values) if other_values else None
    result["box14a_other"] = result["box14_other"]
    result["box14b_tipped_occupation_codes"] = None

    # State boxes 15-17 (accepts both EIN-style and NJ-style state IDs)
    (
        result["box15_state"],
        result["box15_state_id"],
        result["box16_state_wages"],
        result["box17_state_tax"],
    ) = _state_row(text)

    # Local boxes 18-20
    local_values, locality = _local_row(lines)
    if len(local_values) >= 1:
        result["box18_local_wages"] = local_values[0]
    if len(local_values) >= 2:
        result["box19_local_tax"] = local_values[1]
    if not locality:
        index = _line_index(lines, r"\b18\s+Local\s+wages")
        if index is not None:
            for line in lines[index + 1:index + 4]:
                if re.search(r"\bNYC\b", line, re.I):
                    locality = "NYC"
                    break
    result["box20_locality_name"] = locality

    # Metadata
    result["form_type"] = "W-2"
    result["form_code"] = "W-2"
    result["filename"] = filename
    result["source_page_count"] = len(pages)
    result["source_pages"] = [
        {
            "page": page.get("page", index + 1),
            "source": page.get("source", "unknown"),
            **({"ocr_error": page["ocr_error"]} if page.get("ocr_error") else {}),
        }
        for index, page in enumerate(pages)
    ]
    result["extraction_method"] = (
        "ocr"
        if any(page.get("source") == "ocr" for page in pages)
        else "pdf_text"
    )

    # Warnings
    required_fields = (
        "tax_year",
        "employer_ein",
        "employer_name",
        "employee_ssn",
        "box1_wages",
        "box2_fed_tax",
    )
    missing = [key for key in required_fields if result.get(key) in (None, "")]
    if missing:
        warnings.append(
            "Could not confidently extract: "
            + ", ".join(missing)
            + ". Verify these fields against the original W-2."
        )
    if result["extraction_method"] == "ocr":
        warnings.append("OCR-derived values must be checked against the original PDF.")

    result["parse_warnings"] = warnings
    result["raw_text"] = text
    result["review_required"] = bool(warnings)
    return result