"""Scientific literature loaders for bounded, provenance-preserving snapshots."""

from __future__ import annotations

import re
from collections.abc import Iterable
from pathlib import Path
from xml.etree import ElementTree

from graph.graph import EdgeRecord, GraphStore
from graph.real.corpus import RealGraph, SourceMetadata, relation_phrase


def _text(element: ElementTree.Element | None) -> str:
    return " ".join("".join(element.itertext()).split()) if element is not None else ""


def _slug(value: str, *, limit: int = 120) -> str:
    return re.sub(r"[^A-Za-z0-9]+", "_", value.strip()).strip("_")[:limit] or "unknown"


def _make_graph(
    source: SourceMetadata,
    nodes: dict[str, str],
    edges: Iterable[EdgeRecord],
    labels: dict[str, str],
    attributes: dict[str, dict[str, str]],
) -> RealGraph:
    edge_list = tuple(edges)
    relations = sorted({edge.relation_id for edge in edge_list})
    return RealGraph(
        graph=GraphStore(
            nodes,
            edge_list,
            graph_version=f"{source.source_id}@{source.version}",
            node_attributes=attributes,
            relation_attributes={
                relation: {
                    "functional": "false",
                    "description": relation_phrase(relation),
                }
                for relation in relations
            },
        ),
        node_labels=labels,
        relation_phrases={
            relation: relation_phrase(relation) for relation in relations
        },
        source=source,
    )


def load_pubmed_xml(
    path: str | Path,
    source: SourceMetadata,
    *,
    max_records: int | None = None,
) -> RealGraph:
    """Load PubMed titles/abstracts and citation structure as a graph."""
    root = ElementTree.parse(path).getroot()
    nodes: dict[str, str] = {}
    labels: dict[str, str] = {}
    attributes: dict[str, dict[str, str]] = {}
    edges: list[EdgeRecord] = []
    records = root.findall(".//PubmedArticle")
    if max_records is not None:
        records = records[:max_records]
    for record in records:
        pmid = _text(record.find(".//PMID"))
        article = record.find(".//Article")
        if not pmid or article is None:
            continue
        article_id = f"pubmed:{pmid}"
        title = _text(article.find("ArticleTitle"))
        abstract = " ".join(
            _text(item) for item in article.findall("Abstract/AbstractText")
        )
        journal = _text(article.find("Journal/Title"))
        year = _text(article.find("Journal/JournalIssue/PubDate/Year"))
        if not year:
            year = _text(article.find("Journal/JournalIssue/PubDate/MedlineDate"))[:4]
        nodes[article_id] = "scientific_article"
        labels[article_id] = title[:240] or pmid
        attributes[article_id] = {
            "pmid": pmid,
            "title": title[:4000],
            "abstract": abstract[:16000],
            "publication_year": year,
            "evidence_kind": "scientific_abstract",
        }
        if journal:
            journal_id = f"pubmed:journal:{_slug(journal)}"
            nodes[journal_id] = "journal"
            labels[journal_id] = journal
            edges.append(
                EdgeRecord(
                    article_id,
                    "published_in",
                    journal_id,
                    "scientific_article",
                    "journal",
                    relation_phrase("published_in"),
                    {"evidence_kind": "scientific_bibliographic_observation"},
                )
            )
        for author in article.findall("AuthorList/Author"):
            last = _text(author.find("LastName"))
            initials = _text(author.find("Initials"))
            name = " ".join(part for part in (last, initials) if part)
            if not name:
                continue
            author_id = f"pubmed:author:{_slug(name)}"
            nodes[author_id] = "author"
            labels[author_id] = name
            edges.append(
                EdgeRecord(
                    article_id,
                    "written_by",
                    author_id,
                    "scientific_article",
                    "author",
                    relation_phrase("written_by"),
                    {"evidence_kind": "scientific_bibliographic_observation"},
                )
            )
        for mesh in record.findall(".//MeshHeading/DescriptorName"):
            name = _text(mesh)
            if not name:
                continue
            mesh_id = f"pubmed:mesh:{_slug(name)}"
            nodes[mesh_id] = "mesh_descriptor"
            labels[mesh_id] = name
            edges.append(
                EdgeRecord(
                    article_id,
                    "has_subject",
                    mesh_id,
                    "scientific_article",
                    "mesh_descriptor",
                    relation_phrase("has_subject"),
                    {"evidence_kind": "scientific_subject_annotation"},
                )
            )
    return _make_graph(source, nodes, edges, labels, attributes)