from __future__ import annotations

"""Mutation gate for pdffill_tasks, forged and verified 5/5 near-misses rejected.

The checker in pdffill_tasks.py could grade an output but nothing showed it would
reject a broken one. These are the deliberate near-misses it refuses, each one
verified to change the file and to be rejected.
"""
# ---- inlined from pdffill_tasks.py so this file runs on its own ----
"""H1: fill REAL government forms. Inputs are the actual fillable PDFs (IRS W-9, CRA TD1,
USCIS I-9) in data/forms/. The executor gets the blank form at INPUT_DIR/<form>.pdf, a
list of human labels with values, and must write the filled PDF to OUTPUT. It has to
discover which AcroForm field each label means (W-9 fields are named f1_01..f1_15 with no
tooltips; TD1 fields carry tooltips; I-9 field names are the labels themselves). The
checker knows the true mapping (from widget geometry, see data/forms/inspect_fields.py)
and reads the values back with pypdf. No LLM judges anything.
"""
# (moved to the top of this file)
import os
import random
import re
import shutil
import subprocess
import sys
import tempfile

CONTEXTS = ["w9", "td1", "i9"]
HERE = os.path.dirname(os.path.abspath(__file__))
FORMS = os.path.join(HERE, "..", "data", "forms")

# label -> true field name (established from widget rectangles and printed labels)
LABELS = {
    "w9": {
        "Name": "topmostSubform[0].Page1[0].f1_01[0]",
        "Business name": "topmostSubform[0].Page1[0].f1_02[0]",
        "Address": "topmostSubform[0].Page1[0].Address_ReadOrder[0].f1_07[0]",
        "City, state, and ZIP code": "topmostSubform[0].Page1[0].Address_ReadOrder[0].f1_08[0]",
        "Requester's name and address": "topmostSubform[0].Page1[0].f1_09[0]",
        "Account numbers": "topmostSubform[0].Page1[0].f1_10[0]",
    },
    "td1": {
        "Last name": "form1[0].Page1[0].Subform1[0].Identification[0].Last_name[0]",
        "First name and initials": "form1[0].Page1[0].Subform1[0].Identification[0].First_name[0]",
        "Date of birth": "form1[0].Page1[0].Subform1[0].Identification[0].DOB[0]",
        "Employee number": "form1[0].Page1[0].Subform1[0].Identification[0].Employee_number[0]",
        "Address": "form1[0].Page1[0].Subform1[0].Identification[0].Address[0]",
    },
    "i9": {
        "Last name (from Section 1)": "Last Name Family Name from Section 1",
        "First name (from Section 1)": "First Name Given Name from Section 1",
        "Middle initial (from Section 1)": "Middle initial if any from Section 1",
    },
}
VALUES = {
    "Name": ["Maria Okonkwo", "Devin Halloran", "Priya Raman-Cole"], "Business name": ["Okonkwo Consulting LLC", "Halloran Farms", "Raman-Cole Design"],
    "Address": ["18 Birch Lane Apt 4", "902 Mill Road", "77 Harbour St Unit 12"], "City, state, and ZIP code": ["Albany, NY 12207", "Bend, OR 97701", "Toledo, OH 43604"],
    "Requester's name and address": ["Northwind Payables, 1 Market St, Reno NV", "Quill Notes Inc, 40 Elm Ave, Austin TX"], "Account numbers": ["AC-4471", "AC-9020, AC-9021"],
    "Last name": ["Tremblay", "Nguyen", "Okafor", "Bergeron", "Sidhu", "MacLeod", "Petrova"],
    "First name and initials": ["Louise M.", "Daniel", "Chidi A.", "Amrit K.", "Fiona", "Marc-Andre", "Yelena R."],
    "Date of birth": ["1988/04/12", "1975/11/30", "1994/07/03", "1969/02/21", "2001/09/15", "1983/12/08"],
    "Employee number": ["E-10442", "E-2210", "E-77", "E-3391", "E-508", "E-91020"], "Last name (from Section 1)": ["Fernandez", "Whitfield"], "First name (from Section 1)": ["Carla", "Jerome"],
    "Middle initial (from Section 1)": ["R", "T"],
}
FORM_TITLE = {"w9": "IRS Form W-9 (Request for Taxpayer Identification Number and Certification)",
              "td1": "CRA Form TD1 (Personal Tax Credits Return)", "i9": "USCIS Form I-9 (Employment Eligibility Verification)"}


def build(n_per: int = 8, seed: int = 61) -> list[dict]:
    rng = random.Random(seed)
    specs = []
    for ctx in CONTEXTS:
        labels = list(LABELS[ctx])
        for i in range(n_per):
            k = rng.randint(max(2, len(labels) - 2), len(labels))
            chosen = rng.sample(labels, k)
            values = {lab: rng.choice(VALUES[lab]) for lab in chosen}
            specs.append({"context": ctx, "id": f"{ctx}{i}", "form": ctx, "values": values,
                          "spec": (f"Fill the real {FORM_TITLE[ctx]}: the blank form is the file {ctx}.pdf inside the directory "
                                   f"named by the environment variable INPUT_DIR (os.environ['INPUT_DIR']). Fill EXACTLY these fields, "
                                   f"identified by their printed labels on the form, with these values, and leave every other "
                                   f"field empty:\n" + "\n".join(f"- {lab}: {val}" for lab, val in values.items()) +
                                   "\nThe filled values must be stored as the fields' values (AcroForm /V) so a PDF reader shows "
                                   "them, the form must keep all its pages, and the output is written to OUTPUT.")})
    rng.shuffle(specs)
    return specs


FENCE = re.compile(r"```(?:python|py)?\s*\n(.*?)```", re.S)


def extract_script(reply: str) -> str | None:
    m = FENCE.findall(reply or "")
    if m:
        return max(m, key=len)
    if "pypdf" in (reply or "") or "PdfWriter" in (reply or ""):
        return reply
    return None


def run_script(script: str, spec: dict | None = None, timeout_s: int = 120) -> tuple[str | None, str]:
    d = tempfile.mkdtemp(prefix="sg_fill_")
    inp = os.path.join(d, "inputs")
    os.makedirs(inp, exist_ok=True)
    form = (spec or {}).get("form", "w9")
    shutil.copy(os.path.join(FORMS, f"{form}.pdf"), os.path.join(inp, f"{form}.pdf"))
    out = os.path.join(d, "filled.pdf")
    path = os.path.join(d, "build.py")
    with open(path, "w", encoding="utf-8") as fh:
        fh.write(script)
    try:
        p = subprocess.run([sys.executable, path], cwd=d, env={**os.environ, "OUTPUT": out, "INPUT_DIR": inp, "PYTHONIOENCODING": "utf-8"},
                           capture_output=True, text=True, timeout=timeout_s)
        if p.returncode != 0:
            err = next((ln for ln in reversed(p.stderr.splitlines()) if "Error" in ln or "error" in ln), "")
            return None, f"exit {p.returncode}: {(err or p.stderr)[-300:]}"
    except subprocess.TimeoutExpired:
        return None, "timeout"
    return (out if os.path.exists(out) else None), (p.stdout + p.stderr)[-500:]


def _val(v) -> str:
    if v is None:
        return ""
    return str(v).strip()


def grade(spec: dict, path: str | None) -> tuple[int, str]:
    if not path:
        return 0, "no file"
    try:
        from pypdf import PdfReader
        blank = PdfReader(os.path.join(FORMS, f"{spec['form']}.pdf"))
        r = PdfReader(path)
    except Exception as e:  # noqa: BLE001
        return 0, f"unreadable: {str(e)[:80]}"
    try:
        if len(r.pages) != len(blank.pages):
            return 0, f"{len(r.pages)} pages, blank has {len(blank.pages)}"
        fields = r.get_fields() or {}
        mapping = LABELS[spec["form"]]
        for lab, val in spec["values"].items():
            fname = mapping[lab]
            if fname not in fields:
                return 0, f"field for {lab!r} ({fname[-24:]}) missing from output"
            got = _val(fields[fname].get("/V"))
            if got != val:
                return 0, f"{lab!r}: got {got[:40]!r}, want {val!r}"
        # every other mapped field must stay empty (no garbage fill)
        for lab, fname in mapping.items():
            if lab not in spec["values"] and fname in fields and _val(fields[fname].get("/V")):
                return 0, f"{lab!r} should be empty, got {_val(fields[fname].get('/V'))[:30]!r}"
        # a reader must be able to show the values: NeedAppearances or an appearance stream on a filled widget
        need = False
        try:
            af = r.trailer["/Root"].get("/AcroForm")
            need = bool(af and af.get("/NeedAppearances"))
        except Exception:  # noqa: BLE001
            pass
        if not need:
            has_ap = False
            for page in r.pages:
                for a in page.get("/Annots") or []:
                    a = a.get_object()
                    if a.get("/Subtype") == "/Widget" and a.get("/AP") and _val(a.get("/V")):
                        has_ap = True
                        break
                if has_ap:
                    break
            if not has_ap:
                return 0, "no NeedAppearances flag and no appearance stream on filled widgets: a reader would show blanks"
        return 1, "ok"
    except Exception as e:  # noqa: BLE001
        return 0, f"grader exception: {str(e)[:100]}"


def user_prompt(spec: dict) -> str:
    return (spec["spec"] + "\n\nWrite ONE complete Python script (pypdf and pdfplumber are installed) that opens the blank "
            "form from the path in INPUT_DIR, discovers the right AcroForm field for each printed label (inspect the field "
            "names, tooltips and widget positions), fills them, and writes the filled PDF to the path in OUTPUT. Output "
            "only the script in a single ```python fence.")
# ---- end of pdffill_tasks.py ----

import os
import shutil
from pypdf import PdfReader, PdfWriter


def build_reference(spec, out_path):
    form = spec['form']
    here = os.path.dirname(os.path.abspath(__file__)) if '__file__' in globals() else HERE
    src = os.path.join(FORMS, f"{form}.pdf")
    reader = PdfReader(src)
    writer = PdfWriter()
    writer.append(reader)

    mapping = LABELS[form]
    fill = {}
    for lab, val in spec['values'].items():
        fname = mapping[lab]
        fill[fname] = val

    for page in writer.pages:
        writer.update_page_form_field_values(page, fill)

    # Ensure NeedAppearances so a reader displays the values
    try:
        writer.set_need_appearances_writer(True)
    except Exception:
        try:
            acro = writer._root_object['/AcroForm']
            from pypdf.generic import BooleanObject, NameObject
            acro[NameObject('/NeedAppearances')] = BooleanObject(True)
        except Exception:
            pass

    with open(out_path, 'wb') as fh:
        writer.write(fh)
    return out_path


def _load(path):
    return PdfReader(path)


def spoil_extra_field(path):
    out = path + '.spoil_extra_field'
    shutil.copy(path, out)
    reader = PdfReader(out)
    writer = PdfWriter()
    writer.append(reader)

    mapping = LABELS['td1']
    # find a mapped field name NOT used in this spec's values, fill it with garbage
    used_names = set()
    fields = reader.get_fields() or {}
    # We don't have direct access to spec here, so pick any mapped field with empty /V
    target = None
    for lab, fname in mapping.items():
        cur = fields.get(fname)
        val = cur.get('/V') if cur else None
        if not val:
            target = fname
            break
    if target is None:
        # fallback: just pick the first mapped field
        target = next(iter(mapping.values()))

    for page in writer.pages:
        writer.update_page_form_field_values(page, {target: 'UNEXPECTED_VALUE'})

    try:
        writer.set_need_appearances_writer(True)
    except Exception:
        pass

    with open(out, 'wb') as fh:
        writer.write(fh)
    return out


def spoil_wrong_value(path):
    out = path + '.spoil_wrong_value'
    shutil.copy(path, out)
    reader = PdfReader(out)
    writer = PdfWriter()
    writer.append(reader)

    mapping = LABELS['td1']
    fields = reader.get_fields() or {}
    target = None
    for lab, fname in mapping.items():
        cur = fields.get(fname)
        val = cur.get('/V') if cur else None
        if val:
            target = fname
            break
    if target is None:
        target = next(iter(mapping.values()))

    for page in writer.pages:
        writer.update_page_form_field_values(page, {target: 'WRONG_VALUE_XYZ'})

    try:
        writer.set_need_appearances_writer(True)
    except Exception:
        pass

    with open(out, 'wb') as fh:
        writer.write(fh)
    return out


def spoil_missing_page(path):
    out = path + '.spoil_missing_page'
    reader = PdfReader(path)
    writer = PdfWriter()
    for i, page in enumerate(reader.pages):
        if i == len(reader.pages) - 1 and len(reader.pages) > 1:
            continue
        writer.add_page(page)
    # preserve acroform
    try:
        writer._root_object['/AcroForm'] = reader.trailer['/Root']['/AcroForm']
    except Exception:
        pass
    with open(out, 'wb') as fh:
        writer.write(fh)
    return out


def spoil_no_appearance(path):
    out = path + '.spoil_no_appearance'
    shutil.copy(path, out)
    reader = PdfReader(out)
    writer = PdfWriter()
    writer.append(reader)

    # Remove NeedAppearances flag
    try:
        acro = writer._root_object.get('/AcroForm')
        if acro is not None:
            acro_obj = acro.get_object()
            if '/NeedAppearances' in acro_obj:
                del acro_obj['/NeedAppearances']
    except Exception:
        pass

    # Remove /AP from every widget annotation so no appearance stream exists
    for page in writer.pages:
        annots = page.get('/Annots')
        if not annots:
            continue
        annots_obj = annots.get_object() if hasattr(annots, 'get_object') else annots
        for a in annots_obj:
            a_obj = a.get_object()
            if a_obj.get('/Subtype') == '/Widget' and '/AP' in a_obj:
                del a_obj['/AP']

    with open(out, 'wb') as fh:
        writer.write(fh)
    return out


def spoil_stringify_date(path):
    out = path + '.spoil_stringify_date'
    shutil.copy(path, out)
    reader = PdfReader(out)
    writer = PdfWriter()
    writer.append(reader)

    mapping = LABELS['td1']
    dob_field = mapping.get('Date of birth')
    fields = reader.get_fields() or {}
    cur = fields.get(dob_field)
    val = cur.get('/V') if cur else None
    if val:
        # subtly alter formatting (e.g. change separators), still "looks right" but fails exact match
        newval = str(val).replace('/', '-')
        for page in writer.pages:
            writer.update_page_form_field_values(page, {dob_field: newval})

    try:
        writer.set_need_appearances_writer(True)
    except Exception:
        pass

    with open(out, 'wb') as fh:
        writer.write(fh)
    return out


SPOILERS = [
    ('extra_field', spoil_extra_field),
    ('wrong_value', spoil_wrong_value),
    ('missing_page', spoil_missing_page),
    ('no_appearance', spoil_no_appearance),
    ('stringify_date', spoil_stringify_date),
]