"""Zap form extractor (canonical Python implementation).

Finds fillable fields in an HTML page, groups them into forms, and produces
  * the serialized text the Zap model reads (``serialize``), and
  * a cleaned, index-annotated HTML snippet for the teacher LLM (``teacher_html``).

The JavaScript port in ``js/zap-extract.js`` implements the same rules on a live DOM.
Every rule here is expressible with standard DOM APIs, so keep it that way.
"""
from __future__ import annotations

import copy
import hashlib
import re
import unicodedata
from urllib.parse import unquote, urlsplit

import lxml.html
from lxml import etree

SKIP_INPUT_TYPES = {"hidden", "submit", "button", "reset", "image", "file", "checkbox", "radio", "range", "color"}
FIELD_TAGS = {"input", "select", "textarea"}
NO_TEXT_TAGS = {"script", "style", "noscript", "template", "svg", "select", "textarea", "option", "head"}
HEADING_TAGS = {"h1", "h2", "h3", "h4", "h5", "h6", "legend"}
DATA_ATTRS = ["data-testid", "data-test", "data-qa", "data-cy", "formcontrolname", "data-field", "data-name", "data-type"]
MAX_FIELDS = 40
MAX_FORMS = 10

# Key tokens (special tokens of the Zap tokenizer).
CTX_KEYS = ["[TITLE]", "[URL]", "[FORM]", "[HEAD]", "[BTN]", "[LINK]", "[TXT]", "[HID]"]
FIELD_KEYS = ["[FLD]", "[TYPE]", "[NAME]", "[ID]", "[CLASS]", "[AC]", "[PH]", "[LBL]", "[ARIA]", "[TIP]",
              "[NEAR]", "[AFTER]", "[MODE]", "[MAX]", "[REQ]", "[OPT]", "[DATA]"]
SPECIAL_TOKENS = CTX_KEYS + FIELD_KEYS

# Character budgets (Unicode code points).
CTX_LIMITS = {"title": 60, "url": 60, "form": 60, "head": 60, "btn": 60, "link": 80, "txt": 120, "hid": 40}
FIELD_LIMITS = {"name": 40, "id": 40, "class": 30, "ac": 30, "ph": 50, "lbl": 60, "aria": 50, "tip": 40,
                "near": 50, "after": 30, "mode": 12, "max": 6, "opt": 60, "data": 30}
CTX_ORDER = [("title", "[TITLE]"), ("url", "[URL]"), ("form", "[FORM]"), ("head", "[HEAD]"), ("btn", "[BTN]"),
             ("link", "[LINK]"), ("txt", "[TXT]"), ("hid", "[HID]")]
FIELD_ORDER = [("type", "[TYPE]"), ("name", "[NAME]"), ("id", "[ID]"), ("class", "[CLASS]"), ("ac", "[AC]"),
               ("ph", "[PH]"), ("lbl", "[LBL]"), ("aria", "[ARIA]"), ("tip", "[TIP]"), ("near", "[NEAR]"),
               ("after", "[AFTER]"), ("mode", "[MODE]"), ("max", "[MAX]"), ("req", "[REQ]"), ("opt", "[OPT]"),
               ("data", "[DATA]")]

_WS = re.compile(r"\s+")


def norm(s: str | None) -> str:
    """NFKC + collapse whitespace + strip."""
    if not s:
        return ""
    return _WS.sub(" ", unicodedata.normalize("NFKC", s)).strip()


def cut(s: str, n: int) -> str:
    return s[:n].rstrip()


def words(s: str | None) -> str:
    """Identifier -> lowercase words: 'loginForm_Email2' -> 'login form email 2'."""
    if not s:
        return ""
    s = unicodedata.normalize("NFKC", s)
    s = re.sub(r"([a-z])([A-Z])", r"\1 \2", s)
    s = re.sub(r"([A-Z]+)([A-Z][a-z])", r"\1 \2", s)
    s = re.sub(r"([A-Za-z])([0-9])", r"\1 \2", s)
    s = re.sub(r"([0-9])([A-Za-z])", r"\1 \2", s)
    out = "".join(c if unicodedata.category(c)[0] in "LNM" else " " for c in s)
    return _WS.sub(" ", out).strip().lower()


def tag(el) -> str:
    t = el.tag
    return t.lower() if isinstance(t, str) else ""


def attr(el, name: str) -> str:
    v = el.get(name)
    return v if v is not None else ""


def is_element(node) -> bool:
    return isinstance(node.tag, str)


def text(el) -> str:
    """textContent without script/style/select/textarea/option/svg/template/noscript, normalized."""
    parts: list[str] = []

    def walk(e):
        if tag(e) in NO_TEXT_TAGS:
            return
        if e.text and is_element(e):
            parts.append(e.text)
        for c in e:
            if is_element(c):
                walk(c)
            if c.tail:
                parts.append(c.tail)

    if is_element(el):
        walk(el)
    return norm(" ".join(parts))


def input_type(el) -> str:
    t = tag(el)
    if t == "input":
        return (attr(el, "type").strip().lower() or "text")
    return t


def hidden_attr(el) -> bool:
    if el.get("hidden") is not None or attr(el, "aria-hidden").strip().lower() == "true":
        return True
    style = attr(el, "style").replace(" ", "").lower()
    return "display:none" in style or "visibility:hidden" in style


def is_field(el) -> bool:
    t = tag(el)
    if t not in FIELD_TAGS:
        return False
    if t == "input" and input_type(el) in SKIP_INPUT_TYPES:
        return False
    return el.get("disabled") is None


def visible_within(el, root) -> bool:
    node = el
    while node is not None:
        if hidden_attr(node):
            return False
        if node is root:
            return True
        node = node.getparent()
    return True


def contains_field(el) -> bool:
    if tag(el) in FIELD_TAGS:
        return True
    return any(tag(d) in FIELD_TAGS for d in el.iter() if d is not el)


# ---------------------------------------------------------------- sibling walks (DOM node order)

def prev_nodes(el):
    """Previous siblings of el in DOM order, nearest first: ('t', str) or ('e', element)."""
    parent = el.getparent()
    if parent is None:
        return
    kids = list(parent)
    i = kids.index(el)
    for j in range(i - 1, -1, -1):
        k = kids[j]
        if k.tail:
            yield ("t", k.tail)
        if is_element(k):
            yield ("e", k)
    if parent.text:
        yield ("t", parent.text)


def next_nodes(el):
    parent = el.getparent()
    if parent is None:
        return
    if el.tail:
        yield ("t", el.tail)
    kids = list(parent)
    i = kids.index(el)
    for k in kids[i + 1:]:
        if is_element(k):
            yield ("e", k)
        if k.tail:
            yield ("t", k.tail)


def near_text(el, root) -> str:
    node = el
    for _level in range(3):
        seen = 0
        for kind, v in prev_nodes(node):
            if kind == "t":
                t = norm(v)
            else:
                if contains_field(v):
                    return ""
                t = text(v)
            if t:
                return t
            if kind == "e":
                seen += 1
                if seen >= 3:
                    break
        if node is root:
            break
        node = node.getparent()
        if node is None or node is root:
            break
    return ""


def after_text(el) -> str:
    seen = 0
    for kind, v in next_nodes(el):
        if kind == "t":
            t = norm(v)
        else:
            if contains_field(v):
                return ""
            t = text(v)
        if t:
            return t
        if kind == "e":
            seen += 1
            if seen >= 2:
                break
    return ""


# ---------------------------------------------------------------- labels

class Page:
    def __init__(self, doc, url: str):
        self.doc = doc
        self.url = url
        self.by_id: dict[str, object] = {}
        self.label_for: dict[str, list] = {}
        for e in doc.iter():
            if not is_element(e):
                continue
            i = e.get("id")
            if i and i not in self.by_id:
                self.by_id[i] = e
            if tag(e) == "label" and e.get("for"):
                self.label_for.setdefault(e.get("for"), []).append(e)
        t = doc.find(".//title")
        self.title = norm(t.text_content()) if t is not None else ""


def label_text(el, page: Page) -> str:
    out: list[str] = []
    for ref in attr(el, "aria-labelledby").split():
        e = page.by_id.get(ref)
        if e is not None:
            out.append(text(e))
    i = el.get("id")
    if i:
        for lab in page.label_for.get(i, []):
            out.append(text(lab))
    node = el.getparent()
    while node is not None:
        if tag(node) == "label":
            out.append(text(node))
            break
        node = node.getparent()
    seen, uniq = set(), []
    for t in out:
        if t and t not in seen:
            seen.add(t)
            uniq.append(t)
    return " | ".join(uniq)


# ---------------------------------------------------------------- form context

def is_buttonish(e) -> bool:
    t = tag(e)
    if t == "button":
        return True
    if t == "input" and input_type(e) in ("submit", "button", "image"):
        return True
    return attr(e, "role").lower() == "button"


def heading(root) -> str:
    for e in root.iter():
        if is_element(e) and (tag(e) in HEADING_TAGS or attr(e, "role").lower() == "heading"):
            t = text(e)
            if t:
                return t
    node = root
    for _level in range(4):
        n = 0
        for kind, sib in prev_nodes(node):
            if kind != "e":
                continue
            n += 1
            if n > 5:
                break
            if tag(sib) in HEADING_TAGS or attr(sib, "role").lower() == "heading":
                t = text(sib)
                if t:
                    return t
            hs = [d for d in sib.iter() if is_element(d) and d is not sib and
                  (tag(d) in HEADING_TAGS or attr(d, "role").lower() == "heading")]
            for d in reversed(hs):
                t = text(d)
                if t:
                    return t
        node = node.getparent()
        if node is None or tag(node) in ("body", "html"):
            break
    return ""


def button_texts(root) -> str:
    out: list[str] = []
    for e in root.iter():
        if not is_element(e) or not is_buttonish(e):
            continue
        t = text(e) or norm(attr(e, "value")) or norm(attr(e, "aria-label")) or norm(attr(e, "alt")) or norm(attr(e, "title"))
        if t and t not in out:
            out.append(cut(t, 30))
        if len(out) >= 3:
            break
    return " | ".join(out)


def link_texts(root) -> str:
    out: list[str] = []
    for e in root.iter():
        if not is_element(e) or tag(e) != "a":
            continue
        t = text(e) or norm(attr(e, "aria-label")) or norm(attr(e, "title"))
        if t and t not in out:
            out.append(cut(t, 30))
        if len(out) >= 4:
            break
    return " | ".join(out)


def hidden_names(root) -> str:
    out: list[str] = []
    for e in root.iter():
        if is_element(e) and tag(e) == "input" and input_type(e) == "hidden":
            w = words(attr(e, "name"))
            if w and w not in out:
                out.append(w)
            if len(out) >= 4:
                break
    return " ".join(out)


def url_words(url: str) -> str:
    try:
        u = urlsplit(url)
    except ValueError:
        return ""
    host = (u.hostname or "").split(".")
    lead = host[0] if len(host) >= 3 and host[0] != "www" else ""
    return words(lead + " " + unquote(u.path))


def form_words(root) -> str:
    parts = [attr(root, "id"), attr(root, "name"), attr(root, "class")]
    if tag(root) == "form":
        try:
            parts.append(unquote(urlsplit(attr(root, "action")).path))
        except ValueError:
            pass
    else:
        parts.append(attr(root, "role"))
    return (words(" ".join(parts)) + " " + norm(attr(root, "aria-label")).lower()).strip()


# ---------------------------------------------------------------- fields

def field_features(el, root, page: Page) -> dict:
    t = tag(el)
    f: dict = {"type": t if t != "input" else "input " + input_type(el)}
    f["name"] = words(attr(el, "name"))
    f["id"] = words(attr(el, "id"))
    f["class"] = words(attr(el, "class"))
    f["ac"] = norm(attr(el, "autocomplete")).lower()
    f["ph"] = norm(attr(el, "placeholder"))
    f["lbl"] = label_text(el, page)
    f["aria"] = norm(attr(el, "aria-label"))
    f["tip"] = norm(attr(el, "title"))
    near = near_text(el, root)
    f["near"] = "" if near in (f["lbl"], f["ph"], f["aria"]) else near
    f["after"] = after_text(el)
    f["mode"] = norm(attr(el, "inputmode")).lower()
    f["max"] = norm(attr(el, "maxlength"))
    f["req"] = el.get("required") is not None or attr(el, "aria-required").lower() == "true"
    if t == "select":
        opts = [text(o) or norm(o.text_content()) for o in el.iter() if is_element(o) and tag(o) == "option"]
        opts = [o for o in opts if o]
        f["opt"] = (str(len(opts)) + ": " + " | ".join(opts[:6])) if opts else ""
    else:
        f["opt"] = ""
    data = [norm(attr(el, a)) for a in DATA_ATTRS if attr(el, a).strip()]
    f["data"] = words(" ".join(data[:2]))
    for k, n in FIELD_LIMITS.items():
        if isinstance(f.get(k), str):
            f[k] = cut(f[k], n)
    return f


def form_context(root, page: Page) -> dict:
    c = {"title": page.title, "url": url_words(page.url), "form": form_words(root), "head": heading(root),
         "btn": button_texts(root), "link": link_texts(root), "txt": text(root), "hid": hidden_names(root)}
    return {k: cut(v, CTX_LIMITS[k]) for k, v in c.items()}


def group_root(el, body, in_form):
    node = el.getparent()
    for _level in range(6):
        if node is None or node is body or tag(node) in ("body", "html"):
            break
        if any(is_element(d) and is_buttonish(d) and id(d) not in in_form for d in node.iter()):
            return node
        node = node.getparent()
    node = el
    for _level in range(3):
        p = node.getparent()
        if p is None or p is body or tag(p) in ("body", "html"):
            break
        node = p
    return node


def find_forms(doc) -> list[tuple[object, list]]:
    """[(root, [field elements])], document order, <form>s first-seen order."""
    body = doc.find(".//body")
    if body is None:
        body = doc
    groups: dict[int, tuple[object, list]] = {}
    order: list[int] = []
    # keep the proxies alive: lxml recycles element proxies, so id() is only stable while referenced
    form_nodes = [d for f in body.iter("form") for d in f.iter()]
    in_form = {id(d) for d in form_nodes}
    for el in body.iter():
        if not is_element(el) or not is_field(el):
            continue
        root = None
        node = el.getparent()
        while node is not None:
            if tag(node) == "form":
                root = node
                break
            node = node.getparent()
        if root is None:
            root = group_root(el, body, in_form)
        if not visible_within(el, root):
            continue
        key = id(root)
        if key not in groups:
            groups[key] = (root, [])
            order.append(key)
        if len(groups[key][1]) < MAX_FIELDS:
            groups[key][1].append(el)
    return [groups[k] for k in order][:MAX_FORMS]


# ---------------------------------------------------------------- serialization

def serialize_ctx(ctx: dict) -> str:
    return " ".join(f"{key} {ctx[k]}" for k, key in CTX_ORDER if ctx.get(k))


def serialize_field(f: dict) -> str:
    out = ["[FLD]"]
    for k, key in FIELD_ORDER:
        v = f.get(k)
        if k == "req":
            if v:
                out.append("[REQ]")
        elif v:
            out.append(f"{key} {v}")
    return " ".join(out)


def serialize(ctx: dict, fields: list[dict]) -> str:
    return " ".join([serialize_ctx(ctx)] + [serialize_field(f) for f in fields])


# ---------------------------------------------------------------- teacher view

KEEP_ATTRS = {"id", "name", "type", "class", "placeholder", "autocomplete", "title", "for", "role", "action", "method",
              "inputmode", "maxlength", "minlength", "required", "value", "href", "data-zap", "pattern", "alt", "label",
              "multiple", "readonly"} | set(DATA_ATTRS)
DROP_TAGS = {"script", "style", "noscript", "template", "svg", "iframe", "img", "picture", "video", "audio", "canvas",
             "link", "meta", "object", "embed", "path", "source", "use", "head"}


def teacher_html(root, field_els, max_chars=9000) -> str | None:
    for i, el in enumerate(field_els, 1):
        el.set("data-zap", str(i))
    try:
        r = copy.deepcopy(root)
    finally:
        for el in field_els:
            el.attrib.pop("data-zap", None)
    r.tail = None
    for e in list(r.iter()):
        if not is_element(e):  # comments / PIs
            p = e.getparent()
            if p is not None:
                e.drop_tree() if hasattr(e, "drop_tree") else p.remove(e)
            continue
        if tag(e) in DROP_TAGS and e is not r:
            e.drop_tree()
    for e in r.iter():
        if not is_element(e):
            continue
        for a in list(e.attrib):
            if a.startswith("aria-") or a in KEEP_ATTRS:
                v = e.attrib[a]
                if tag(e) == "input" and a == "value" and input_type(e) not in ("submit", "button", "hidden"):
                    del e.attrib[a]
                    continue
                lim = 60 if a == "class" else 80
                if a == "href":
                    try:
                        v = urlsplit(v).path
                    except ValueError:
                        v = ""
                if len(v) > lim:
                    e.attrib[a] = v[:lim]
                elif a == "href":
                    e.attrib[a] = v
            else:
                del e.attrib[a]
        if tag(e) == "select":
            opts = [o for o in e if is_element(o) and tag(o) in ("option", "optgroup")]
            if len(opts) > 6:
                for o in opts[6:]:
                    e.remove(o)
                e.append(etree.Comment(f" +{len(opts) - 6} more options "))
        if e.text:
            e.text = cut(norm(e.text), 200) + (" " if e.text[-1:].isspace() else "")
        if e.tail:
            e.tail = " " + cut(norm(e.tail), 200) + " "
    html = lxml.html.tostring(r, encoding="unicode", method="html")
    html = re.sub(r">\s+<", "> <", html)
    html = _WS.sub(" ", html).strip()
    return html if len(html) <= max_chars else None


def signature(ctx: dict, fields: list[dict]) -> str:
    key = "|".join([ctx.get("form", ""), ctx.get("btn", "")] +
                   [f"{f['type']}/{f['name']}/{f['id']}/{f['ac']}/{f['ph']}/{f['lbl']}" for f in fields])
    return hashlib.sha1(key.encode()).hexdigest()[:16]


def html5_normalize(html: str) -> str:
    """Re-serialize through an HTML5-spec parser (lexbor) so the tree matches what browsers build
    (e.g. <form> inside <table>); lxml's libxml2 parser nests malformed markup differently."""
    try:
        from selectolax.lexbor import LexborHTMLParser
    except ImportError:
        return html
    try:
        return LexborHTMLParser(html).html or html
    except Exception:  # noqa: BLE001
        return html


def extract_page(html: str | bytes, url: str, want_teacher: bool = True, html5: bool = True) -> list[dict]:
    if html5 and isinstance(html, str):
        html = html5_normalize(html)
    try:
        doc = lxml.html.document_fromstring(html)
    except (etree.ParserError, ValueError):
        return []
    page = Page(doc, url)
    out = []
    for root, els in find_forms(doc):
        fields = [field_features(el, root, page) for el in els]
        ctx = form_context(root, page)
        rec = {"kind": "form" if tag(root) == "form" else "group", "ctx": ctx, "fields": fields,
               "sig": signature(ctx, fields)}
        if want_teacher:
            rec["html"] = teacher_html(root, els)
        out.append(rec)
    return out