"""Generate synthetic forms for classes the static web crawl under-represents (card payment, OTP, password
change/reset, checkout address, identifier-first logins, SPA div-soup markup), then extract them with the
standard extractor. They are labeled later by the same teacher pass as real forms.

Usage: python synth.py OUT_DIR N [--conc 32] [--seed 5]
"""
import argparse
import asyncio
import gzip
import json
import os
import random
import re
import sys

import aiohttp

sys.path.insert(0, "/job/work")
from zap.extract import extract_page  # noqa: E402

SCENARIOS = [
    (8, "checkout payment step with credit/debit card fields (card number, name on card, expiry, CVC)"),
    (4, "checkout page with shipping address AND card payment fields in one form"),
    (4, "card payment form where expiry month and expiry year are two separate <select> dropdowns"),
    (2, "donation form with amount, name, email and card details"),
    (3, "subscription billing form with card number, expiry (MM/YY single field), CVC and billing ZIP"),
    (2, "add a new payment card form in account settings (cardholder name, number, expiry, security code)"),
    (4, "two-factor authentication page asking for a 6-digit code from an authenticator app (one input)"),
    (4, "SMS verification page where the 6-digit code is split into six single-character inputs"),
    (3, "email verification page asking for a code sent by email, with a resend link"),
    (1, "2FA page asking for a backup/recovery code"),
    (4, "change password page for a logged-in user: current password, new password, confirm new password"),
    (3, "reset password page reached from an email link: new password and confirm password only"),
    (2, "forgot password page asking for the account email address"),
    (1, "forgot password page asking for username or email"),
    (3, "identifier-first sign-in step 1: only an email or phone field and a Next button"),
    (3, "sign-in step 2: the username is shown as text and only the password field is present"),
    (1, "login form with phone number and password"),
    (4, "shipping address form: full name, address line 1, address line 2, city, state/province, postal code, country, phone"),
    (2, "billing address form with first name, last name, company, street, house number, city, ZIP, country"),
    (2, "account profile / billing details edit form with company name and VAT field"),
    (1, "login form inside a modal dialog on a shop homepage"),
    (1, "registration form with username, email, password and repeat password"),
]
STYLES = [
    "plain semantic HTML with <form>, <label for> and ids",
    "Bootstrap 5 markup",
    "Tailwind CSS utility classes and no <form> element (a React SPA, fields inside divs, button type=button)",
    "Material UI style: floating labels placed after the input, generated class names like css-1x2y3z",
    "Angular reactive form with formcontrolname attributes and mat-form-field wrappers",
    "Vue app with data-v-xxxx attributes and v-model style names",
    "WooCommerce / WordPress markup",
    "Shopify-like checkout markup with data-testid attributes",
    "legacy table-based layout with <td> labels",
    "Stripe Elements-like structure where each card field is its own div with an input",
    "minimal markup where inputs have only placeholders and no labels",
    "inputs with aria-label only and obfuscated names like field_7f3a2 or input_12",
]
QUIRKS = [
    "no autocomplete attributes at all",
    "autocomplete=\"off\" on every input",
    "correct HTML autocomplete tokens on every input",
    "some inputs use type=\"text\" even for email/phone/number",
    "the password inputs have a show/hide toggle button next to them",
    "field names are in the page language (not English)",
    "a hidden honeypot text input styled off-screen",
    "a CAPTCHA text input near the submit button",
    "inputmode and maxlength attributes where appropriate",
    "",
]
LANGS = ["English"] * 8 + ["Spanish", "French", "German", "Portuguese (Brazil)", "Italian", "Dutch", "Polish",
                           "Russian", "Japanese", "Chinese (Simplified)", "Korean", "Turkish", "Arabic", "Hindi",
                           "Indonesian", "Vietnamese", "Swedish", "Czech", "Ukrainian", "Greek", "Hebrew", "Thai"]
SITES = ["online store", "bank", "airline", "streaming service", "SaaS dashboard", "forum", "news site", "university portal",
         "government service", "hotel booking site", "food delivery app", "crypto exchange", "telecom operator",
         "insurance company", "game platform", "email provider", "charity", "gym", "car rental", "marketplace",
         "developer platform", "social network", "utility company", "ticketing site", "healthcare portal"]

PROMPT = """Write the HTML of a realistic web page section for this form.

Form: {scenario}
Website: a {site} (invent a plausible brand name)
Language of all visible text: {lang}
Markup style: {style}
Detail: {quirk}

Output exactly:
TITLE: <the page's <title> text>
URL: <a plausible https URL of the page>
<the HTML of the page section containing the form: include a heading, the fields, buttons and links; no <html>, <head>, <script> or <style>>"""


async def gen(session, url, prompt, sem):
    body = {"model": "qwen", "messages": [{"role": "user", "content": prompt}], "temperature": 1.0, "top_p": 0.95,
            "max_tokens": 1800, "chat_template_kwargs": {"enable_thinking": False}}
    async with sem:
        for _ in range(3):
            try:
                async with session.post(url, json=body, timeout=aiohttp.ClientTimeout(total=900)) as r:
                    j = await r.json()
                return j["choices"][0]["message"]["content"] or ""
            except Exception:  # noqa: BLE001
                await asyncio.sleep(3)
    return ""


def parse(text):
    t = re.search(r"TITLE:\s*(.+)", text)
    u = re.search(r"URL:\s*(\S+)", text)
    m = re.search(r"<[a-zA-Z]", text[u.end():] if u else text)
    if not (t and u and m):
        return None
    html = (text[u.end():] if u else text)[m.start():]
    html = re.sub(r"```\s*$", "", html.strip()).strip()
    return t.group(1).strip()[:200], u.group(1).strip(), html


async def main(a):
    os.makedirs(a.out, exist_ok=True)
    rng = random.Random(a.seed)
    weights = [w for w, _ in SCENARIOS]
    jobs = []
    for i in range(a.n):
        sc = rng.choices(SCENARIOS, weights)[0][1]
        jobs.append({"i": i, "scenario": sc, "site": rng.choice(SITES), "lang": rng.choice(LANGS),
                     "style": rng.choice(STYLES), "quirk": rng.choice(QUIRKS)})
    url = f"http://127.0.0.1:{a.port}/v1/chat/completions"
    sem = asyncio.Semaphore(a.conc)
    out = gzip.open(os.path.join(a.out, f"part-synth-{a.seed}.jsonl.gz.tmp"), "wt")
    stats = {"ok": 0, "bad": 0, "forms": 0}
    async with aiohttp.ClientSession(connector=aiohttp.TCPConnector(limit=a.conc)) as s:
        async def one(job):
            text = await gen(s, url, PROMPT.format(**job), sem)
            p = parse(text)
            if not p:
                stats["bad"] += 1
                return
            title, page_url, html = p
            doc = f"<html><head><title>{title}</title></head><body>{html}</body></html>"
            forms = extract_page(doc, page_url)
            stats["ok"] += 1
            for k, f in enumerate(forms):
                f.update(url=page_url + f"#synth{job['i']}", host=f"synth,{a.seed},{job['i']}", lang=job["lang"],
                         src="synth", fidx=k, synth=job)
                out.write(json.dumps(f, ensure_ascii=False) + "\n")
                stats["forms"] += 1
            if (stats["ok"] + stats["bad"]) % 200 == 0:
                print(stats, flush=True)
        await asyncio.gather(*[one(j) for j in jobs])
    out.close()
    os.rename(os.path.join(a.out, f"part-synth-{a.seed}.jsonl.gz.tmp"), os.path.join(a.out, f"part-synth-{a.seed}.jsonl.gz"))
    print("DONE", stats, flush=True)


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("out")
    ap.add_argument("n", type=int)
    ap.add_argument("--conc", type=int, default=32)
    ap.add_argument("--seed", type=int, default=5)
    ap.add_argument("--port", type=int, default=30020)
    asyncio.run(main(ap.parse_args()))