"""
Scientific PDF parser — GROBID client with pymupdf4llm fallback.

Strategy:
1. Try GROBID (if server reachable at GROBID_URL env var or localhost:8070)
   - POST /api/processFulltextDocument → TEI XML
   - Parse TEI for: title, abstract, sections, paragraphs, figures, tables,
     equations, citations, references
2. Fall back to pymupdf4llm markdown parse (same as pdf_digital) + heuristic
   equation/citation detection

GROBID TEI namespace: http://www.tei-c.org/ns/1.0
"""

from __future__ import annotations

import ipaddress
import logging
import os
import re
import socket
import urllib.parse
from collections.abc import Iterator

_log = logging.getLogger(__name__)

from .base import BaseParser, FileRecord, RawElement  # noqa: E402
from .pdf_digital import _parse_markdown_page  # noqa: E402

_TEI_NS = "http://www.tei-c.org/ns/1.0"
_T = f"{{{_TEI_NS}}}"

# Citation heuristics: [1], [Author, 2020], (Author et al., 2020)
_CITATION_RE = re.compile(
    r"(\[\d+(?:,\s*\d+)*\]|\[[\w\s]+,\s*\d{4}\]|\([\w\s]+\s+et\s+al\.,?\s*\d{4}\))"
)


_PRIVATE_NETWORKS = [
    ipaddress.ip_network("10.0.0.0/8"),
    ipaddress.ip_network("172.16.0.0/12"),
    ipaddress.ip_network("192.168.0.0/16"),
    ipaddress.ip_network("169.254.0.0/16"),
    ipaddress.ip_network("127.0.0.0/8"),
    ipaddress.ip_network("::1/128"),
]


def _grobid_url() -> str | None:
    raw = os.environ.get("GROBID_URL", "http://localhost:8070")
    parsed = urllib.parse.urlparse(raw)

    if parsed.scheme not in ("http", "https"):
        _log.warning("GROBID_URL rejected: scheme %r is not http/https", parsed.scheme)
        return None

    hostname = parsed.hostname
    if not hostname:
        _log.warning("GROBID_URL rejected: missing hostname")
        return None

    try:
        resolved = socket.getaddrinfo(hostname, None)
    except socket.gaierror as exc:
        _log.warning("GROBID_URL rejected: hostname %r could not be resolved: %s", hostname, exc)
        return None

    for _family, _type, _proto, _canonname, sockaddr in resolved:
        addr_str = sockaddr[0]
        try:
            ip = ipaddress.ip_address(addr_str)
        except ValueError:
            continue
        for net in _PRIVATE_NETWORKS:
            if ip in net:
                _log.warning(
                    "GROBID_URL rejected: resolved address %s is in private/reserved range %s",
                    ip,
                    net,
                )
                return None

    return raw


def _try_grobid(path) -> str | None:  # type: ignore[no-untyped-def]
    """POST to GROBID and return TEI XML string, or None on any failure."""
    try:
        import requests  # type: ignore[import-untyped]

        base = _grobid_url()
        if base is None:
            return None
        url = f"{base}/api/processFulltextDocument"
        with open(str(path), "rb") as f:
            resp = requests.post(
                url,
                files={"input": f},
                params={"consolidateReferences": "0", "includeRawAffiliations": "0"},
                timeout=120,
            )
        if resp.status_code == 200:
            return resp.text
    except Exception:
        pass
    return None


def _tei_text(el) -> str:  # type: ignore[no-untyped-def]
    """Extract all text content from a TEI element."""
    return " ".join(el.itertext()).strip()


def _parse_tei(tei_xml: str) -> Iterator[RawElement]:
    """Parse GROBID TEI XML and yield RawElements."""
    try:
        from lxml import etree
    except ImportError as e:
        raise ImportError("lxml required: pip install lxml") from e

    _safe_parser = etree.XMLParser(resolve_entities=False, no_network=True)
    root = etree.fromstring(tei_xml.encode(), _safe_parser)

    # Title
    title_el = root.find(f".//{_T}titleStmt/{_T}title[@type='main']")
    if title_el is None:
        title_el = root.find(f".//{_T}title")
    if title_el is not None:
        title = _tei_text(title_el)
        if title:
            yield RawElement(
                element_type="structural:title",
                text=title,
                heading_level=1,
                detection_tier=1,
                extra={"extraction_method": "grobid:tei", "confidence": 0.99},
            )

    # Abstract
    abstract_el = root.find(f".//{_T}abstract")
    if abstract_el is not None:
        abstract_text = _tei_text(abstract_el)
        if abstract_text:
            yield RawElement(
                element_type="text:abstract",
                text=abstract_text,
                extra={"extraction_method": "grobid:tei", "confidence": 0.99},
            )

    # Body
    body = root.find(f".//{_T}body")
    if body is not None:
        yield from _walk_tei_body(body)

    # References
    ref_list = root.find(f".//{_T}listBibl")
    if ref_list is not None:
        for i, bibl in enumerate(ref_list.findall(f"{_T}biblStruct")):
            ref_text = _tei_text(bibl)
            if ref_text:
                yield RawElement(
                    element_type="scientific:reference_entry",
                    text=ref_text,
                    reference_data={"index": i + 1, "raw": ref_text},
                    extra={"extraction_method": "grobid:tei", "confidence": 0.97},
                )


def _walk_tei_body(body, depth: int = 0) -> Iterator[RawElement]:  # type: ignore[no-untyped-def]
    """Recursively walk TEI body/div elements."""
    for el in body:
        local = el.tag.replace(_T, "") if el.tag.startswith(_T) else el.tag

        if local == "div":
            head = el.find(f"{_T}head")
            if head is not None:
                head_text = _tei_text(head)
                n = head.get("n", "")
                level = int(n.split(".")[0]) + 1 if n and n[0].isdigit() else depth + 2
                level = min(max(level, 2), 6)
                if head_text:
                    yield RawElement(
                        element_type="structural:section_header",
                        text=head_text,
                        heading_level=level,
                        detection_tier=1,
                        extra={"extraction_method": "grobid:tei", "confidence": 0.98},
                    )
            yield from _walk_tei_body(el, depth + 1)

        elif local == "p":
            text = _tei_text(el)
            if not text:
                continue
            citations = _CITATION_RE.findall(text)
            extra: dict = {"extraction_method": "grobid:tei", "confidence": 0.95}
            if citations:
                extra["inline_citations"] = citations
            yield RawElement(element_type="text:narrative", text=text, extra=extra)

        elif local == "formula":
            formula_text = _tei_text(el)
            label = el.get("{http://www.w3.org/XML/1998/namespace}id") or None
            if formula_text:
                yield RawElement(
                    element_type="scientific:equation_display",
                    text=formula_text,
                    equation_data={"latex": formula_text, "label": label},
                    extra={"extraction_method": "grobid:tei", "confidence": 0.93},
                )

        elif local == "figure":
            desc = el.find(f"{_T}figDesc")
            head = el.find(f"{_T}head")
            caption = (
                _tei_text(desc)
                if desc is not None
                else (_tei_text(head) if head is not None else "")
            )
            label = _tei_text(head) if head is not None else ""

            tbl = el.find(f".//{_T}table")
            if tbl is not None:
                yield from _parse_tei_table(tbl, caption)
            elif caption:
                yield RawElement(
                    element_type="media:figure",
                    text=caption,
                    figure_data={
                        "caption": caption,
                        "label": label or None,
                        "image_ref": None,
                        "image_sha256": None,
                        "ocr_text": None,
                    },
                    extra={"extraction_method": "grobid:tei", "confidence": 0.90},
                )


def _parse_tei_table(tbl, caption: str = "") -> Iterator[RawElement]:  # type: ignore[no-untyped-def]
    """Convert a TEI table element to a dual-format RawElement."""
    headers: list[str] = []
    rows: list[list[str]] = []

    for row in tbl.findall(f".//{_T}row"):
        cells = [_tei_text(c) for c in row.findall(f"{_T}cell")]
        if not cells:
            continue
        role = row.get("role", "")
        if role == "label" and not headers:
            headers = cells
        else:
            rows.append(cells)

    if not headers and rows:
        headers = rows.pop(0)

    md_lines: list[str] = []
    if headers:
        md_lines.append("| " + " | ".join(headers) + " |")
        md_lines.append("|" + "|".join(" --- " for _ in headers) + "|")
    for row in rows:
        padded = list(row) + [""] * max(0, len(headers) - len(row))
        md_lines.append("| " + " | ".join(padded) + " |")
    markdown = "\n".join(md_lines)

    yield RawElement(
        element_type="table:simple",
        text=markdown or caption,
        table_data={
            "markdown": markdown,
            "structured": {"headers": headers, "rows": rows},
            "has_header_row": bool(headers),
            "row_count": len(rows),
            "header_row_count": 1 if headers else 0,
            "col_count": len(headers) if headers else (len(rows[0]) if rows else 0),
            "caption": caption or None,
        },
        extra={"extraction_method": "grobid:tei", "confidence": 0.92},
    )


def _fallback_scientific(path) -> Iterator[RawElement]:  # type: ignore[no-untyped-def]
    """pymupdf4llm fallback with citation annotation heuristics."""
    try:
        import pymupdf  # type: ignore[import-untyped]
        import pymupdf4llm  # type: ignore[import-untyped]
    except ImportError as e:
        raise ImportError(f"pymupdf4llm required: pip install pymupdf4llm — {e}") from e

    doc = pymupdf.open(str(path))
    try:
        for page_num in range(doc.page_count):
            md_text = pymupdf4llm.to_markdown(doc, pages=[page_num])
            if not md_text or not md_text.strip():
                continue
            for el in _parse_markdown_page(
                md_text, page_num + 1, "pymupdf4llm:scientific-fallback"
            ):
                if el.element_type == "text:narrative":
                    citations = _CITATION_RE.findall(el.text)
                    if citations:
                        el.extra["inline_citations"] = citations
                yield el
    finally:
        doc.close()


class PDFScientificParser(BaseParser):
    handles = [("application/pdf", "scientific")]

    def parse(self, file_record: FileRecord) -> Iterator[RawElement]:
        tei_xml = _try_grobid(file_record.path)
        if tei_xml:
            try:
                yield from _parse_tei(tei_xml)
                return
            except Exception:
                pass  # fall through to pymupdf fallback

        yield from _fallback_scientific(file_record.path)
