"""Elasticsearch storage, source/extraction preservation, and full-match aggregations."""

import asyncio
import calendar
import contextlib
import math
import re
import unicodedata
from collections.abc import AsyncIterator
from datetime import UTC, date, datetime
from typing import Any

from elasticsearch import AsyncElasticsearch, AuthorizationException, BadRequestError, NotFoundError
from elasticsearch.helpers import async_bulk

from app.config import Settings, settings
from app.statistics import (
    DERIVED_SCALES,
    INCONCLUSIVE_REASONS,
    _analysis_type,
    _completion_date,
    _margin,
    _unreported,
    assign_bucket,
    derived_field,
    stored_analysis,
)

BUCKETS = ("effect", "credible_null", "reported_null", "inconclusive", "failed", "unreported")
# Direction is a property of an effect, not a bucket: these partition the effect bucket.
EFFECT_DIRECTIONS = {
    "favours_intervention": "favoursIntervention",
    "favours_comparator": "favoursComparator",
    "unclear": "unclear",
}
_CACHE_FIELDS = set(
    (
        "extracted_at extraction_version extraction_status extraction_evidence evidence_span "
        "extraction_source fulltext_status "
        "population intervention comparator outcome outcome_unit n n_intervention n_comparator "
        "mean_intervention mean_comparator sd_intervention sd_comparator "
        "events_intervention events_comparator "
        "estimate ci_low ci_high p_value effect_type ci_level ci_sides p_value_operator has_control "
        "design primary_outcome_met reported_result result_direction effect_direction "
        "outcome_direction"
    ).split()
)
_LINK_FIELDS = set(
    (
        "linked_papers canonical_id canonical_ids registry_title registry_abstract registry_url "
        "paper_result_label possible_abstract_spin"
    ).split()
)
_PAPER_FIELDS = set(
    (
        "title abstract authors venue embedding embedding_model result_label null_score evidence_span "
        "classification_method abstract_available work_type snapshot_provenance pmcid"
    ).split()
)


# Embedded studies, the corpus the gap map is drawn from. Reviews stay in: they
# occupy the space too, and `gapmap` excludes them from a region's attempts.
_EMBEDDED_QUERY = {
    "bool": {"filter": [{"term": {"record_kind": "study"}}, {"exists": {"field": "embedding"}}]}
}
# Everything the map needs per document and nothing else: abstracts would
# multiply the bytes on the wire without changing a single position.
_MAP_FIELDS = [
    "embedding",
    "title",
    "year",
    "bucket",
    "record_kind",
    "is_review",
    "cited_by_count",
    "source",
]
_PIT_KEEP_ALIVE = "10m"


def _split_budget(limit: int, slices: int) -> list[int | None]:
    """Share a document budget between concurrent slices; no limit means no cap."""
    if limit <= 0:
        return [None] * slices
    base, extra = divmod(limit, slices)
    return [base + (1 if index < extra else 0) for index in range(slices)]


def index_mapping(dimensions: int = 384) -> dict:
    properties: dict[str, Any] = {}
    for name in (
        "record_kind source result_label bucket evidence_tier effect_type overall_status "
        "pmids nct_ids referenced_works result_pmids canonical_id canonical_ids outcome_unit "
        "p_value_operator ci_sides extraction_version extraction_status embedding_model "
        "analysis_effect_type effect_direction outcome_direction work_type "
        "pmcid extraction_source fulltext_status reported_result result_direction"
    ).split():
        properties[name] = {"type": "keyword"}
    for name in "title abstract population intervention comparator outcome".split():
        properties[name] = {"type": "text"}
    # Separate retrieval text allows normalization without changing source quotes.
    properties["search_text"] = {"type": "text", "analyzer": "english"}
    for name in "evidence_span why_stopped rationale".split():
        properties[name] = {"type": "text", "index": False}
    for name in (
        "estimate ci_low ci_high p_value null_score ci_level n enrollment_actual enrollment_planned "
        "analysis_estimate analysis_se analysis_ci_low analysis_ci_high mde "
        "n_intervention n_comparator mean_intervention mean_comparator "
        "sd_intervention sd_comparator events_intervention events_comparator"
    ).split():
        properties[name] = {"type": "double"}
    # Arm-level effects on every derivable scale, so a query on any scale can be
    # bucketed over the full match set without reading documents.
    for scale in DERIVED_SCALES:
        for name in ("estimate", "ci_low", "ci_high"):
            properties[derived_field(scale, name)] = {"type": "double"}
    for name in "year cited_by_count".split():
        properties[name] = {"type": "integer"}
    # Citation-authority signals. rank_feature rejects non-positive values, so
    # writers log-scale and leave the field absent instead of storing 0.
    for name in ("authority", "pagerank"):
        properties[name] = {"type": "rank_feature"}
    for name in (
        "is_retracted is_review has_results has_linked_publication has_control spin_flag "
        "possible_abstract_spin reporting_missing significant_but_trivial abstract_available"
    ).split():
        properties[name] = {"type": "boolean"}
    for name in (
        "extracted_at primary_completion_date publication_date received_at reporting_due_date"
    ).split():
        properties[name] = {"type": "date"}
    properties["embedding"] = {
        "type": "dense_vector",
        "dims": dimensions,
        "index": True,
        "similarity": "cosine",
        "index_options": {"type": "int8_hnsw"},
    }
    for name in (
        "registry_analyses",
        "extraction_evidence",
        "metadata",
        "linked_papers",
        "snapshot_provenance",
    ):
        properties[name] = {"type": "object", "enabled": False}
    return {"dynamic": False, "properties": properties}


_QUERY_STOPWORDS = set(
    """
a an the and or of in on at to for from with without by versus vs as is are was were be been
being do does did can could should would will may might whether that this these those their
its our your among into over under than compared comparing comparison effect effects efficacy
impact influence intervention interventions treatment treatments therapy therapies supplementation
supplement supplements supplemental administration administered oral daily receiving receive
change changes changing increase increased increasing decrease decreased decreasing reduce reduced
reducing reduction improvement improved improving improve decline severity symptoms symptom
outcome outcomes measure measures measured score scores level levels rate rates global function
functional study studies trial trials controlled randomized randomised placebo control group groups
expression expressed elevated elevation higher lower associated association status presence
""".split()
)
_QUERY_STEMS = ("depress", "cognit", "arthroscop", "osteoarthrit", "hypertens")
# Two or more letters then digits (kim1, il6, covid19). A single letter is left alone
# so the vitamin d3/d2 handling keeps its own tokens.
_JOINED_DESIGNATOR = re.compile(r"([a-z]{2,})(\d+)")


def _search_normalize(value: str) -> str:
    folded = "".join(
        char for char in unicodedata.normalize("NFKD", value) if not unicodedata.combining(char)
    ).casefold()
    # NFKD turns subscript ₃ into 3; also align VitaminD3 and vitamin D(3).
    return re.sub(
        r"\bvitamin[\s-]*d(?:[\s-]*\(?([23])\)?)?(?![a-z0-9])",
        lambda match: "vitamin d" + (match[1] or ""),
        folded,
    )


def _concept_terms(value: str) -> list[str]:
    terms = []
    for word in re.findall(r"[a-z0-9]+", _search_normalize(value)):
        if word in _QUERY_STOPWORDS:
            continue
        token = next((stem for stem in _QUERY_STEMS if word.startswith(stem)), word)
        if token not in terms:
            terms.append(token)
    return terms


def _concept_query(terms: list[str], fields: list[str], *, require_all: bool) -> dict:
    clauses = []
    for term in terms:
        if term in _QUERY_STEMS:
            # The stored standard analyzer does not stem depression/depressive.
            # A small, explicit morphology set avoids changing an existing index.
            clauses.append(
                {
                    "bool": {
                        "should": [
                            {"prefix": {field.split("^")[0]: {"value": term}}} for field in fields
                        ],
                        "minimum_should_match": 1,
                    }
                }
            )
        elif joined := _JOINED_DESIGNATOR.fullmatch(term):
            # The standard analyzer indexes "KIM-1" as kim + 1, so a typed "kim1" would
            # otherwise match only the few documents that spell it without the hyphen.
            spaced = " ".join(joined.groups())
            clauses.append(
                {
                    "bool": {
                        "should": [
                            {"multi_match": {"query": term, "fields": fields}},
                            {"multi_match": {"query": spaced, "fields": fields, "type": "phrase"}},
                        ],
                        "minimum_should_match": 1,
                    }
                }
            )
        else:
            clauses.append({"multi_match": {"query": term, "fields": fields}})
    minimum = len(clauses) if require_all else max(1, math.ceil(len(clauses) * 0.6))
    return {"bool": {"should": clauses, "minimum_should_match": minimum}}


_POPULATION_GENERIC = set(
    """
adult adults patient patients participant participants people person persons individual individuals
population populations subject subjects volunteer volunteers men women male males female females
child children pediatric paediatric adolescent adolescents young younger old older elderly aged age
ages aging ageing year years month months day days healthy general community dwelling communitydwelling
post menopausal postmenopausal premenopausal middle middleaged senior seniors who whom have has had
having diagnosed diagnosis confirmed suspected established symptomatic asymptomatic mild moderate severe
major minor history risk risks high higher low lower at living suffering experiencing presenting
scheduled elective undergoing eligible enrolled treated untreated baseline underlying previously
currently type mellitus chronic acute recent newly stage stages all any approximately least more less
""".split()
)
_ANATOMICAL_TERMS = set(
    """
knee knees hip hips elbow elbows shoulder shoulders ankle ankles foot feet wrist wrists hand hands
spine spinal lumbar cervical thoracic lung lungs breast prostate kidney kidneys renal heart cardiac
brain colon colorectal rectal liver hepatic skin eye eyes ear ears dental tooth teeth jaw
""".split()
)


# An alias made only of these words would admit nearly every clinical record.
_ALIAS_TOO_BROAD = set(
    "disease diseases disorder disorders syndrome syndromes condition conditions illness "
    "illnesses complication complications health".split()
)
_MAX_POPULATION_ALIASES = 4


def _specific_terms(value: str) -> list[str]:
    return [
        term
        for term in _concept_terms(value)
        if term not in _POPULATION_GENERIC and not term.isdecimal()
    ]


def _population_terms(pico: dict) -> list[str]:
    return _specific_terms(pico.get("population") or "")


def _population_alias_terms(pico: dict) -> list[list[str]]:
    """Model-proposed equivalent condition names, reduced to specific terms.

    Aliases widen the full match set and therefore every bucket count, so an
    alias that adds nothing specific, or merely repeats the population, is dropped.
    """
    primary = set(_population_terms(pico))
    result: list[list[str]] = []
    for alias in list(pico.get("populationAliases") or [])[:_MAX_POPULATION_ALIASES]:
        terms = _specific_terms(str(alias))
        if not terms or terms in result:
            continue
        if all(term in _ALIAS_TOO_BROAD or term in _ANATOMICAL_TERMS for term in terms):
            continue
        specific = {term for term in terms if term not in _ALIAS_TOO_BROAD}
        if specific <= primary or primary <= specific:
            # A subset is broader than the population; a superset is already matched by it.
            continue
        result.append(terms)
    return result


def population_query(pico: dict) -> dict | None:
    """Constrain explicit conditions/sites; generic demographic descriptions only rank.

    This guard also applies to reference expansion. It deliberately does not try
    to infer diagnostic synonyms itself: absent textual support is a retrieval limitation.
    The parser may supply ``populationAliases`` (equivalent names of the same condition);
    each is a full alternative to the condition clause and is ignored without a population.
    Multiple anatomical sites are alternatives under the same condition.
    """
    terms = _population_terms(pico)
    if not terms:
        return None
    fields = ["population^3", "title^3", "abstract", "search_text"]
    anatomy = [term for term in terms if term in _ANATOMICAL_TERMS]
    condition = [term for term in terms if term not in _ANATOMICAL_TERMS]
    clauses = []
    if anatomy:
        clauses.append(
            {
                "bool": {
                    "should": [
                        _concept_query([term], fields, require_all=True) for term in anatomy
                    ],
                    "minimum_should_match": 1,
                }
            }
        )
    if condition:
        clauses.append(_concept_query(condition, fields, require_all=True))
    primary = {"bool": {"must": clauses}}
    alternatives = []
    for alias in _population_alias_terms(pico):
        must = [_concept_query(alias, fields, require_all=True)]
        # An alias that drops the anatomical site must not escape it: "OA" still needs "knee".
        if anatomy and not any(term in _ANATOMICAL_TERMS for term in alias):
            must.append(clauses[0])
        alternatives.append({"bool": {"must": must}})
    if not alternatives:
        return primary
    return {"bool": {"should": [primary, *alternatives], "minimum_should_match": 1}}


# ClinicalTrials.gov records have no citation count: ingestion writes 0 for every
# registry row. A citation bound applied to them would delete exactly the
# terminated and never-reported trials this tool exists to surface, so registry
# rows are exempt from citation bounds and the pipeline warns that they are.
# Publication dates do exist for registry rows (the primary completion date), so
# a date bound applies to papers and trials alike.
CITATION_EXEMPT_SOURCES = ("ctgov", "merged")


def citation_range(filters: dict | None) -> dict | None:
    """The citation bound on its own, or None when the request set neither end."""
    if not filters:
        return None
    bounds = {}
    if filters.get("minCitations") is not None:
        bounds["gte"] = int(filters["minCitations"])
    if filters.get("maxCitations") is not None:
        bounds["lte"] = int(filters["maxCitations"])
    return {"range": {"cited_by_count": bounds}} if bounds else None


def registry_citation_exemption_query(query: dict, filters: dict | None) -> dict | None:
    """Rows in the match set only because citation bounds skip the registry.

    Counting these makes the exemption visible: it is the number of trial records
    the citation requirement would have removed had it applied to them.
    """
    bound = citation_range(filters)
    if bound is None:
        return None
    return {
        "bool": {
            "must": [query],
            "filter": [{"terms": {"source": list(CITATION_EXEMPT_SOURCES)}}],
            "must_not": [bound],
        }
    }


def filter_clauses(filters: dict | None) -> list[dict]:
    """Pre-search corpus restrictions, as filter clauses over the full match set.

    These belong in the query's ``filter``, never in ``should``: the caller's
    expanded review references and every aggregation share that clause list, so a
    filtered-out record cannot re-enter the study list while the counts exclude it.
    A record with no ``publication_date`` cannot satisfy a date bound and drops out.
    """
    if not filters:
        return []
    clauses: list[dict] = []
    dates = {}
    if filters.get("yearFrom") is not None:
        dates["gte"] = f"{int(filters['yearFrom']):04d}-01-01"
    if filters.get("yearTo") is not None:
        dates["lte"] = f"{int(filters['yearTo']):04d}-12-31"
    if dates:
        clauses.append({"range": {"publication_date": dates}})
    citations = citation_range(filters)
    if citations is not None:
        clauses.append(
            {
                "bool": {
                    "should": [
                        citations,
                        {"terms": {"source": list(CITATION_EXEMPT_SOURCES)}},
                    ],
                    "minimum_should_match": 1,
                }
            }
        )
    return clauses


def lexical_query(
    pico: dict,
    idea: str,
    expanded_ids: list[str] | None = None,
    filters: dict | None = None,
) -> dict:
    """Require intervention/outcome plus explicit disease or anatomical scope.

    Untyped model synonyms are boosts, not alternative mandatory concepts: a mood
    synonym must not let a depression-only paper satisfy the vitamin D clause.
    Verified review references remain an explicit extension of this match universe.
    """
    intervention = _concept_terms(pico.get("intervention") or "")
    outcome = _concept_terms(pico.get("outcome") or "")
    intervention_fields = ["intervention^5", "title^3", "abstract", "search_text"]
    outcome_fields = ["outcome^5", "title^3", "abstract", "search_text"]
    must = []
    if intervention and outcome:
        if "arthroscop" in intervention:
            # "Arthroscopic" already supplies the surgical anchor. Preserve
            # an explicit choice of lavage OR debridement instead of requiring both.
            intervention = [term for term in intervention if term not in {"surgery", "surgical"}]
        primary_intervention = _concept_query(intervention, intervention_fields, require_all=True)
        if (
            "arthroscop" in intervention
            and {"lavage", "debridement"}.issubset(intervention)
            and re.search(r"\bor\b|/", _search_normalize(pico.get("intervention") or ""))
        ):
            anchor = [term for term in intervention if term not in {"lavage", "debridement"}]
            primary_intervention = {
                "bool": {
                    "must": [
                        _concept_query(anchor, intervention_fields, require_all=True),
                        {
                            "bool": {
                                "should": [
                                    _concept_query(
                                        [procedure], intervention_fields, require_all=True
                                    )
                                    for procedure in ("lavage", "debridement")
                                ],
                                "minimum_should_match": 1,
                            }
                        },
                    ]
                }
            }
        intervention_options = [primary_intervention]
        # Established intervention aliases for the three seeded clinical topics.
        aliases = list(pico.get("interventionAliases", []))[:6]
        if "vitamin" in intervention and any(term in intervention for term in ("d", "d3", "d2")):
            aliases += [
                "cholecalciferol",
                "ergocalciferol",
                "calciferol",
                "vitamin d3",
                "vitamin d2",
            ]
        elif "omega" in intervention and "3" in intervention:
            aliases += [
                "fish oil",
                "n 3 fatty acids",
                "eicosapentaenoic acid",
                "docosahexaenoic acid",
            ]
        for alias in aliases:
            terms = _concept_terms(alias)
            if terms:
                intervention_options.append(
                    _concept_query(terms, intervention_fields, require_all=True)
                )
        outcome_options = [_concept_query(outcome, outcome_fields, require_all=False)]
        outcome_aliases = list(pico.get("outcomeAliases", []))[:6]
        population_terms = _population_terms(pico)
        if "pain" in outcome and "knee" in population_terms and "osteoarthrit" in population_terms:
            # WOMAC is a knee/hip OA pain-stiffness-function composite. Retrieve
            # it in this narrow context; retain the actual endpoint for statistics.
            # Instrument validation: https://pubmed.ncbi.nlm.nih.gov/3068365/
            outcome_aliases.append("WOMAC")
        for alias in outcome_aliases:
            terms = _concept_terms(alias)
            if terms:
                outcome_options.append(_concept_query(terms, outcome_fields, require_all=True))
        must = [
            {"bool": {"should": intervention_options, "minimum_should_match": 1}},
            {"bool": {"should": outcome_options, "minimum_should_match": 1}},
        ]
    else:
        # With no reliable PICO pair, retain a transparent original-query fallback.
        must = [
            {
                "multi_match": {
                    "query": idea,
                    "fields": [
                        "title^3",
                        "abstract",
                        "intervention^3",
                        "outcome^2",
                        "population",
                        "search_text",
                    ],
                    "minimum_should_match": "60%",
                }
            }
        ]
    boosts = []
    if pico.get("population"):
        boosts.append(
            {
                "multi_match": {
                    "query": pico["population"],
                    "fields": ["population^2", "title", "abstract"],
                    "boost": 0.3,
                }
            }
        )
    synonyms = [str(term) for term in pico.get("synonyms", []) if term]
    if synonyms:
        boosts.append(
            {
                "multi_match": {
                    "query": " ".join(synonyms),
                    "fields": ["title", "abstract", "intervention", "outcome"],
                    "boost": 0.3,
                }
            }
        )
    candidate = {"bool": {"must": must, "should": boosts, "minimum_should_match": 0}}
    should = [candidate]
    if expanded_ids:
        should.append({"ids": {"values": list(dict.fromkeys(expanded_ids))}})
    required = [{"term": {"record_kind": "study"}}]
    population = population_query(pico)
    if population is not None:
        required.append(population)
    required.extend(filter_clauses(filters))
    return {
        "bool": {
            "filter": required,
            "should": should,
            "minimum_should_match": 1,
        }
    }


def authority_query(query: dict) -> dict:
    """Rank the query's own match set by combined citation authority.

    Each ``rank_feature`` leaf contributes a saturation-scaled function of the
    stored log signals (``authority`` = log citation count, ``pagerank`` =
    within-corpus citation graph), and the ``should`` clauses sum them. The
    main query stays a filter, so authority re-ranks relevant hits only —
    an unrelated famous paper must not enter results on reputation alone.
    Documents with neither signal simply absent themselves from this leg.
    """
    return {
        "bool": {
            "filter": [query],
            "should": [
                {"rank_feature": {"field": "authority"}},
                {"rank_feature": {"field": "pagerank"}},
            ],
            "minimum_should_match": 1,
        }
    }


# How many ranked records one search carries forward. It has to cover the relevance screen
# (pipeline.SCREEN_LIMIT): a shorter list would silently cap what can be screened.
RETRIEVE_LIMIT = 500


def rrf_fuse(*rankings: list[dict], limit: int = RETRIEVE_LIMIT) -> list[dict]:
    scores: dict[str, float] = {}
    docs = {}
    for ranking in rankings:
        for rank, doc in enumerate(ranking, 1):
            key = doc["id"]
            scores[key] = scores.get(key, 0) + 1 / (60 + rank)
            docs[key] = doc
    return [docs[key] for key in sorted(scores, key=lambda key: (-scores[key], key))[:limit]]


# All numerical validation, transformations, reconstruction and failure rules run
# in statistics.assign_bucket at write time. This script only reapplies the query's
# margin to those validated 95% bounds, plus the time-sensitive reporting deadline.
# Pass doc explicitly to functions: Painless functions do not capture script locals.
# An inconclusive row is emitted as 'inconclusive:<reason>' (statistics.INCONCLUSIVE_REASONS)
# so one pass yields both the bucket and why; aggregate() folds the reasons back together.
BUCKET_SCRIPT = """
boolean present(def values, String key) {
    return values.containsKey(key) && values[key].size() > 0;
}
String text(def values, String key) {
    return present(values, key) ? values[key].value.toString() : '';
}
boolean flag(def values, String key) {
    return present(values, key) && values[key].value == true;
}
if (text(doc, 'bucket') == 'failed') { emit('failed'); return; }
if (flag(doc, 'is_review')) { emit('inconclusive:review'); return; }
if (flag(doc, 'reporting_missing') && present(doc, 'reporting_due_date')
    && doc['reporting_due_date'].value.toInstant().toEpochMilli() < params.today) {
    emit('unreported'); return;
}
if (!present(doc, 'analysis_ci_low') || !present(doc, 'analysis_ci_high')) {
    String source = text(doc, 'source');
    boolean comparative = flag(doc, 'has_control') || source == 'ctgov' || source == 'merged'
        || present(doc, 'nct_ids');
    boolean read = text(doc, 'extraction_status') == 'verified';
    String label = comparative && present(doc, 'reported_result') ? text(doc, 'reported_result')
        : (read && !comparative ? '' : text(doc, 'result_label'));
    emit(label == 'positive' ? 'effect' : (label == 'null' ? 'reported_null'
        : (label == 'mixed' ? 'inconclusive:mixed_result' : 'inconclusive:no_result')));
    return;
}
if (params.delta == null) { emit('inconclusive:no_threshold'); return; }
// The stored analysis is on the default scale; a query on another scale falls
// back to the arm-level effect derived for that scale, mirroring _numeric_evidence.
String prefix = 'analysis';
if (text(doc, 'analysis_effect_type') != params.effect) {
    if (params.derived == null || !present(doc, params.derived + '_ci_low')
        || !present(doc, params.derived + '_ci_high')) { emit('inconclusive:other_scale'); return; }
    prefix = params.derived;
}
double low = doc[prefix + '_ci_low'].value;
double high = doc[prefix + '_ci_high'].value;
if (low > -params.delta && high < params.delta) { emit('credible_null'); return; }
if ((low > 0 || high < 0) && present(doc, prefix + '_estimate')
    && Math.abs(doc[prefix + '_estimate'].value) >= params.delta) { emit('effect'); return; }
emit('inconclusive:wide_interval');
"""


def bucket_runtime(sesoi: float, effect_type: str, today: datetime | None = None) -> dict:
    """Runtime field applying BUCKET_SCRIPT with one query's margin and effect scale."""
    today = (today or datetime.now(UTC)).replace(hour=0, minute=0, second=0, microsecond=0)
    scale = _analysis_type(effect_type)
    return {
        "query_bucket": {
            "type": "keyword",
            "script": {
                "source": BUCKET_SCRIPT,
                "params": {
                    "delta": _margin(sesoi, effect_type),
                    "effect": scale,
                    "derived": derived_field(scale, "")[:-1] if scale in DERIVED_SCALES else None,
                    "today": int(today.timestamp() * 1000),
                },
            },
        }
    }


def _prepare_document(incoming: dict, existing: dict | None = None) -> dict:
    """Refresh source data while preserving a cache tied to unchanged source text."""
    old = existing or {}
    doc = {**old, **incoming}
    doc["record_kind"] = old.get("record_kind", incoming.get("record_kind", "study"))
    doc.setdefault("is_review", False)
    # A later registry refresh must not undo a proven paper/trial link.
    if old.get("source") == "merged" and incoming.get("source") == "ctgov":
        for key in _LINK_FIELDS | _PAPER_FIELDS:
            if key in old:
                doc[key] = old[key]
        doc.update(source="merged", has_linked_publication=True)
        for key in ("pmids", "result_pmids", "referenced_works"):
            doc[key] = sorted(set(old.get(key) or []) | set(incoming.get(key) or []))
    # Reingestion of a linked paper must not resurrect a second patient sample.
    if old.get("record_kind") == "linked_publication":
        for key in ("canonical_id", "canonical_ids"):
            if key in old:
                doc[key] = old[key]
    unchanged = old.get("abstract", "") == doc.get("abstract", "")
    cached = bool(old.get("extracted_at"))
    source_has_new_registry_numbers = incoming.get("source") in {"ctgov", "merged"} and bool(
        incoming.get("numeric_source")
    )
    if unchanged and cached and not source_has_new_registry_numbers:
        for key in _CACHE_FIELDS:
            if key in old:
                doc[key] = old[key]
    elif cached and (not unchanged or source_has_new_registry_numbers):
        # Partial updates need explicit nulls, otherwise removed cache keys survive.
        for key in _CACHE_FIELDS:
            doc[key] = incoming.get(key)
        # Registry-linked rows keep searchable literature text and its provenance.
        if old.get("source") == "merged" and incoming.get("source") == "ctgov":
            doc["evidence_span"] = incoming.get("evidence_span") or old.get("evidence_span", "")
    if not unchanged and "embedding" not in incoming:
        doc["embedding"] = None
    elif unchanged and incoming.get("embedding") is None:
        # get/mget may intentionally omit vector source; partial update preserves it.
        doc.pop("embedding", None)
    if not doc.get("publication_date"):
        year = doc.get("year")
        doc["publication_date"] = (
            f"{int(year):04d}-01-01" if year and 1 <= int(year) <= 9999 else None
        )
    doc["search_text"] = _search_normalize(
        " ".join(
            str(doc.get(field) or "")
            for field in (
                "title",
                "abstract",
                "population",
                "intervention",
                "comparator",
                "outcome",
            )
        )
    )
    # Normalization is the single source of truth for query-time numerical rules.
    doc.update(assign_bucket(doc))
    doc.update(stored_analysis(doc))
    completion = _completion_date(doc.get("primary_completion_date"))
    doc["primary_completion_date"] = completion.isoformat() if completion else None
    # Substitute a known past date only to ask the same strict metadata predicate
    # whether results are confirmed absent. Recency is evaluated separately.
    doc["reporting_missing"] = _unreported(
        {**doc, "primary_completion_date": "1900-01-01"}, date(2000, 1, 1)
    )
    doc["reporting_due_date"] = None
    if completion is not None and completion.year < 9999:
        anniversary = completion.replace(
            year=completion.year + 1,
            day=min(completion.day, calendar.monthrange(completion.year + 1, completion.month)[1]),
        )
        doc["reporting_due_date"] = anniversary.isoformat()
    # rank_feature rejects non-positive values, so authority is absent rather
    # than 0. pagerank is written only by the offline citation-graph pass and
    # survives reingest through the {**old, **incoming} merge above.
    cited_by = doc.get("cited_by_count")
    if (
        isinstance(cited_by, (int, float))
        and not isinstance(cited_by, bool)
        and math.isfinite(cited_by)
        and cited_by > 0
    ):
        doc["authority"] = math.log1p(cited_by)
    else:
        doc.pop("authority", None)
    return doc


class ElasticRepository:
    def __init__(self, config: Settings = settings, client: AsyncElasticsearch | None = None):
        self.config = config
        self.index = config.elastic_index
        self.client = client or AsyncElasticsearch(
            config.elastic_url,
            api_key=config.elastic_api_key
            if config.elastic_api_key and not config.elastic_local
            else None,
            request_timeout=30,
            max_retries=2,
            retry_on_timeout=True,
        )
        self._index_ready = False
        self._vector_source_filter = True
        self._index_lock = asyncio.Lock()

    async def close(self) -> None:
        await self.client.close()

    async def ensure_index(self) -> None:
        if self._index_ready:
            return
        async with self._index_lock:
            if self._index_ready:
                return
            mapping = index_mapping(self.config.embedding_dimensions)
            if not await self.client.indices.exists(index=self.index):
                try:
                    await self.client.indices.create(index=self.index, mappings=mapping)
                except BadRequestError as exc:
                    if "resource_already_exists_exception" not in str(exc):
                        raise
                    await self.client.indices.put_mapping(index=self.index, **mapping)
            else:
                # Add newly introduced explicit fields on existing deployments.
                await self.client.indices.put_mapping(index=self.index, **mapping)
            self._index_ready = True

    async def health(self) -> dict:
        await self.client.info()
        exists = await self.client.indices.exists(index=self.index)
        count = (
            await self.client.count(index=self.index, query={"term": {"record_kind": "study"}})
            if exists
            else {"count": 0}
        )
        return {"connected": True, "index": self.index, "studies": count["count"]}

    async def bulk_upsert(self, studies: list[dict]) -> int:
        if not studies:
            return 0
        await self.ensure_index()
        existing = {
            doc["id"]: doc
            for doc in await self.get_many(
                [study["id"] for study in studies], include_vectors=False
            )
        }
        actions = []
        for item in studies:
            doc = _prepare_document(item, existing.get(item["id"]))
            actions.append(
                {
                    "_op_type": "update",
                    "_index": self.index,
                    "_id": doc["id"],
                    "doc": doc,
                    "doc_as_upsert": True,
                }
            )
        count, _ = await async_bulk(self.client, actions, refresh="wait_for")
        return count

    async def get(self, study_id: str) -> dict | None:
        try:
            result = await self.client.get(index=self.index, id=study_id)
            return dict(result["_source"], id=result["_id"])
        except NotFoundError:
            return None

    async def get_many(self, ids: list[str], *, include_vectors: bool = True) -> list[dict]:
        if not ids:
            return []
        identifiers = list(dict.fromkeys(ids))
        if not include_vectors:
            result = await self.client.mget(
                index=self.index, ids=identifiers, source_excludes=["embedding"]
            )
        elif not self._vector_source_filter:
            result = await self.client.mget(index=self.index, ids=identifiers)
        else:
            # 9.2+ defaults to excluding vectors; 9.1 keeps them but does not
            # accept this newer source-filter option inside mget documents.
            try:
                result = await self.client.mget(
                    index=self.index,
                    docs=[
                        {"_id": identifier, "_source": {"exclude_vectors": False}}
                        for identifier in identifiers
                    ],
                )
            except BadRequestError as exc:
                message = str(exc).lower()
                if not any(
                    fragment in message
                    for fragment in ("exclude_vectors", "expected text", "unknown field")
                ):
                    raise
                self._vector_source_filter = False
                result = await self.client.mget(index=self.index, ids=identifiers)
        return [
            dict(item["_source"], id=item["_id"]) for item in result["docs"] if item.get("found")
        ]

    @staticmethod
    def hits(response: Any) -> list[dict]:
        return [dict(h["_source"], id=h["_id"]) for h in response["hits"]["hits"]]

    async def sample_embedded(self, limit: int = 4000, seed: int = 0) -> dict:
        """A deterministic random sample of embedded studies, plus what it omits.

        Random rather than top-cited: a map of where the literature is thin must
        not be drawn from its most popular corner. The counts travel with the
        sample so a caller can say which fraction of the corpus it describes.
        """
        await self.ensure_index()
        query = _EMBEDDED_QUERY
        total = await self.client.count(index=self.index, query=query)
        body = {
            "index": self.index,
            "size": max(0, limit),
            "query": {
                "function_score": {
                    "query": query,
                    "random_score": {"seed": seed, "field": "_seq_no"},
                    # A filter-only query scores every document 0, and the default
                    # multiply boost would turn the random score into 0 as well,
                    # making "the sample" the first `size` documents in index order
                    # — measured against the live index, that returned 4000 registry
                    # records and not one of the 567k papers.
                    "boost_mode": "replace",
                }
            },
            "source_includes": _MAP_FIELDS,
        }
        result = await self._search_with_vectors(body)
        return {"documents": self.hits(result), "corpus": total["count"]}

    async def count_embedded(self) -> int:
        """How many embedded studies exist, so coverage can be stated as a fraction."""
        await self.ensure_index()
        result = await self.client.count(index=self.index, query=_EMBEDDED_QUERY)
        return int(result["count"])

    async def scan_embedded(
        self,
        batch_size: int = 2000,
        slices: int = 8,
        limit: int = 0,
    ) -> AsyncIterator[list[dict]]:
        """Every embedded study, in batches, over one point in time.

        Paginated rather than fetched at once: the embedded corpus is millions of
        384-float vectors, so a single oversized search would neither fit in one
        response nor let a caller show anything before the last document arrived.
        Slices run concurrently because a single `search_after` walk is limited by
        round-trip latency, not by Elasticsearch — measured against the live index,
        one walk moved ~3.2k documents/s and sixteen slices ~11k/s.

        A point in time keeps the walk consistent, so the same document is never
        yielded twice and none is missed while the index is written to.
        """
        await self.ensure_index()
        slices = max(1, slices)
        pit = (await self.client.open_point_in_time(index=self.index, keep_alive=_PIT_KEEP_ALIVE))
        pit_id = pit["id"]
        queue: asyncio.Queue = asyncio.Queue(maxsize=slices * 2)
        budgets = _split_budget(limit, slices)

        async def walk(slice_id: int, budget: int | None) -> None:
            after: list | None = None
            taken = 0
            while budget is None or taken < budget:
                size = batch_size if budget is None else min(batch_size, budget - taken)
                # No index: a point in time already names one, and Elasticsearch
                # rejects a search that specifies both.
                body: dict[str, Any] = {
                    "size": size,
                    "query": _EMBEDDED_QUERY,
                    "pit": {"id": pit_id, "keep_alive": _PIT_KEEP_ALIVE},
                    "sort": [{"_shard_doc": "asc"}],
                    "source_includes": _MAP_FIELDS,
                }
                if slices > 1:
                    body["slice"] = {"id": slice_id, "max": slices}
                if after is not None:
                    body["search_after"] = after
                response = await self._search_with_vectors(body)
                hits = response["hits"]["hits"]
                if not hits:
                    return
                after = hits[-1]["sort"]
                taken += len(hits)
                await queue.put(self.hits(response))

        async def produce() -> None:
            try:
                await asyncio.gather(*(walk(index, budgets[index]) for index in range(slices)))
            finally:
                await queue.put(None)

        producer = asyncio.create_task(produce())
        try:
            while True:
                batch = await queue.get()
                if batch is None:
                    break
                yield batch
            await producer
        finally:
            producer.cancel()
            with contextlib.suppress(asyncio.CancelledError, Exception):
                await producer
            with contextlib.suppress(Exception):  # the walk is over either way
                await self.client.close_point_in_time(id=pit_id)

    async def _search_with_vectors(self, body: dict) -> Any:
        """Search keeping the stored vectors in `_source`, on 9.1 and on 9.2+."""
        if not self._vector_source_filter:
            return await self.client.search(**body)
        try:
            return await self.client.search(**body, source_exclude_vectors=False)
        except (BadRequestError, TypeError):
            self._vector_source_filter = False
            return await self.client.search(**body)

    async def knn_studies(
        self, vector: list[float], limit: int, filters: dict | None = None
    ) -> list[dict]:
        """Nearest studies to an arbitrary vector, each with its own embedding.

        The vectors travel back because the caller reports a cosine per concept,
        which cannot be recomputed without them. Elasticsearch scores a `cosine`
        dense vector as `(1 + cos) / 2`, so the score is converted back to the
        cosine rather than reported as an opaque relevance number.
        """
        await self.ensure_index()
        body = {
            "index": self.index,
            "size": limit,
            "knn": {
                "field": "embedding",
                "query_vector": vector,
                "k": limit,
                "num_candidates": min(10 * limit, 1000),
                "filter": [{"term": {"record_kind": "study"}}, *filter_clauses(filters)],
            },
            "source_includes": [
                "embedding",
                "title",
                "year",
                "url",
                "source",
                "bucket",
                "query_bucket",
                "cited_by_count",
            ],
        }
        result = await self._search_with_vectors(body)
        return [
            dict(hit["_source"], id=hit["_id"], cosine=2 * float(hit["_score"]) - 1)
            for hit in result["hits"]["hits"]
        ]

    async def retrieve(self, query: dict, vector: list[float] | None) -> tuple[list[dict], str]:
        base = {"index": self.index, "size": RETRIEVE_LIMIT, "source_excludes": ["embedding"]}
        # Third ranking signal: citation authority over the same match set.
        authority = authority_query(query)
        if vector is None:
            lexical, ranked = await asyncio.gather(
                self.client.search(**base, query=query),
                self.client.search(**base, query=authority),
            )
            return rrf_fuse(self.hits(lexical), self.hits(ranked)), "bm25"
        knn = {
            "field": "embedding",
            "query_vector": vector,
            "k": RETRIEVE_LIMIT,
            "num_candidates": 2 * RETRIEVE_LIMIT,
            "filter": query,
        }
        try:
            result = await self.client.search(
                **base,
                retriever={
                    "rrf": {
                        "rank_window_size": RETRIEVE_LIMIT,
                        "rank_constant": 60,
                        "retrievers": [
                            {"standard": {"query": query}},
                            {"knn": knn},
                            {"standard": {"query": authority}},
                        ],
                    }
                },
            )
            return self.hits(result), "hybrid_rrf"
        except (BadRequestError, AuthorizationException) as exc:
            # Only a license restriction permits fallback; auth/DSL errors must surface.
            message = str(exc).lower()
            if "license" not in message or not any(
                term in message for term in ("rrf", "rank fusion", "retriever")
            ):
                raise
            lexical, semantic, ranked = await asyncio.gather(
                self.client.search(**base, query=query),
                self.client.search(**base, knn=knn),
                self.client.search(**base, query=authority),
            )
            return (
                rrf_fuse(self.hits(lexical), self.hits(semantic), self.hits(ranked)),
                "hybrid_client_rrf",
            )

    async def screen_population(
        self, ids: list[str], pico: dict, filters: dict | None = None
    ) -> set[str]:
        """Apply the population guard and the corpus filters to expanded rows.

        Review references enter the study list without passing the lexical query,
        so they are screened here against the same conditions the aggregation uses.
        """
        identifiers = list(dict.fromkeys(ids))
        population = population_query(pico)
        restrictions = filter_clauses(filters)
        if not identifiers or (population is None and not restrictions):
            return set(identifiers)
        clauses = [
            {"ids": {"values": identifiers}},
            {"term": {"record_kind": "study"}},
        ]
        if population is not None:
            clauses.append(population)
        clauses.extend(restrictions)
        result = await self.client.search(
            index=self.index,
            size=len(identifiers),
            source=False,
            query={"bool": {"filter": clauses}},
        )
        return {hit["_id"] for hit in result["hits"]["hits"]}

    async def registry_sweep(self, query: dict) -> list[dict]:
        result = await self.client.search(
            index=self.index,
            size=100,
            query={
                "bool": {"must": [query], "filter": [{"terms": {"source": ["ctgov", "merged"]}}]}
            },
            source_excludes=["embedding"],
        )
        return self.hits(result)

    async def cache_extraction(self, study_id: str, extraction: dict) -> None:
        current = await self.get(study_id)
        if current is None:
            raise ValueError("Cannot cache extraction for a missing study")
        updated = _prepare_document({**current, **extraction})
        # Only cache and derived fields change; source vectors remain intact.
        derived = set(assign_bucket(updated)) | {
            "reporting_missing",
            "reporting_due_date",
            "search_text",
        }
        patch = {key: updated.get(key) for key in set(extraction) | derived}
        await self.client.update(index=self.index, id=study_id, doc=patch, refresh="wait_for")

    async def persist_links(self, merged: list[dict]) -> None:
        """Persist each canonical trial, then hide its duplicate publication rows."""
        if not merged:
            return
        paper_ids = list(
            dict.fromkeys(
                paper["id"]
                for trial in merged
                for paper in trial.get("linked_papers", [])
                if paper.get("id")
            )
        )
        papers = {paper["id"]: paper for paper in await self.get_many(paper_ids)}
        canonical_ids = {trial["id"] for trial in merged}
        links: dict[str, set[str]] = {}
        prepared = []
        for trial in merged:
            doc = dict(trial)
            doc["linked_papers"] = list(
                {
                    paper["id"]: paper
                    for paper in trial.get("linked_papers", [])
                    if paper.get("id")
                }.values()
            )
            for paper in doc["linked_papers"]:
                pid = paper["id"]
                if pid in canonical_ids:
                    continue
                links.setdefault(pid, set()).add(doc["id"])
                full_paper = papers.get(pid, {})
                if (
                    full_paper.get("abstract") == doc.get("abstract")
                    and full_paper.get("embedding") is not None
                ):
                    doc["embedding"] = full_paper["embedding"]
            prepared.append(doc)
        await self.bulk_upsert(prepared)
        actions = []
        for paper_id, targets in links.items():
            if paper_id not in papers:
                continue
            existing = papers[paper_id].get("canonical_ids") or []
            if isinstance(existing, str):
                existing = [existing]
            targets.update(existing)
            actions.append(
                {
                    "_op_type": "update",
                    "_index": self.index,
                    "_id": paper_id,
                    "doc": {"record_kind": "linked_publication", "canonical_ids": sorted(targets)},
                }
            )
        if actions:
            await async_bulk(self.client, actions, refresh="wait_for")

    @staticmethod
    def match_set(query: dict) -> dict:
        """The counted population: primary study records matching the query."""
        return {
            "bool": {
                "must": [query],
                "filter": [{"term": {"record_kind": "study"}}],
                "must_not": [{"term": {"is_review": True}}],
            }
        }

    async def count_studies(self, query: dict) -> int:
        """Size of one match set, so a filtered search can say what it removed.

        Uses the same predicate as ``aggregate``; otherwise the difference between
        a filtered and an unfiltered count would not be attributable to the filters.
        """
        await self.ensure_index()
        result = await self.client.count(index=self.index, query=self.match_set(query))
        return result["count"]

    async def aggregate(self, query: dict, sesoi: float, effect_type: str) -> dict:
        today = datetime.now(UTC).replace(hour=0, minute=0, second=0, microsecond=0)
        runtime = bucket_runtime(sesoi, effect_type, today)
        q = self.match_set(query)
        completed = {
            "bool": {
                "filter": [
                    {"term": {"overall_status": "COMPLETED"}},
                    {"terms": {"source": ["ctgov", "merged"]}},
                ]
            }
        }
        result = await self.client.search(
            index=self.index,
            query=q,
            size=0,
            track_total_hits=True,
            runtime_mappings=runtime,
            aggs={
                "buckets": {
                    "terms": {
                        "field": "query_bucket",
                        "size": len(BUCKETS) + len(INCONCLUSIVE_REASONS),
                    }
                },
                "years": {
                    "date_histogram": {
                        "field": "publication_date",
                        "calendar_interval": "year",
                        "min_doc_count": 1,
                    }
                },
                "nulls": {
                    "filter": {"terms": {"query_bucket": ["credible_null", "reported_null"]}},
                    "aggs": {
                        "sample": {
                            "sampler": {"shard_size": 300},
                            "aggs": {
                                "terms": {
                                    "significant_text": {
                                        "field": "abstract",
                                        "size": 12,
                                        "filter_duplicate_text": True,
                                        "background_filter": q,
                                    }
                                }
                            },
                        }
                    },
                },
                "completed": {
                    "filter": completed,
                    "aggs": {
                        "missing": {"filter": {"term": {"reporting_missing": True}}},
                        "overdue": {
                            "filter": {
                                "bool": {
                                    "filter": [
                                        {"term": {"reporting_missing": True}},
                                        {
                                            "range": {
                                                "reporting_due_date": {"lt": today.isoformat()}
                                            }
                                        },
                                    ]
                                }
                            }
                        },
                    },
                },
                "effect_directions": {
                    "filter": {"term": {"query_bucket": "effect"}},
                    "aggs": {
                        "directions": {
                            "terms": {
                                "field": "result_direction",
                                "missing": "unclear",
                                "size": len(EFFECT_DIRECTIONS),
                            }
                        }
                    },
                },
                "spin_candidates": {
                    "filter": {
                        "bool": {
                            "filter": [
                                {"term": {"source": "merged"}},
                                {"term": {"has_results": True}},
                            ]
                        }
                    },
                    "aggs": {"spin": {"filter": {"term": {"possible_abstract_spin": True}}}},
                },
            },
        )
        aggs = result["aggregations"]
        counts = dict.fromkeys(BUCKETS, 0)
        reasons = dict.fromkeys(INCONCLUSIVE_REASONS, 0)
        for row in aggs["buckets"]["buckets"]:
            bucket, _, reason = row["key"].partition(":")
            counts[bucket] += row["doc_count"]
            if reason:
                reasons[reason] += row["doc_count"]
        directions = dict.fromkeys(EFFECT_DIRECTIONS.values(), 0)
        for row in aggs["effect_directions"]["directions"]["buckets"]:
            # An unrecognized direction is not a fourth answer; it is one nobody stated.
            directions[EFFECT_DIRECTIONS.get(row["key"], "unclear")] += row["doc_count"]
        registered = aggs["completed"]["doc_count"]
        unreported = aggs["completed"]["missing"]["doc_count"]
        return {
            "total": result["hits"]["total"]["value"],
            "bucketCounts": counts,
            "inconclusiveReasons": reasons,
            "effectDirections": directions,
            "yearCounts": [
                {"year": int(b["key_as_string"][:4]), "count": b["doc_count"]}
                for b in aggs["years"]["buckets"]
            ],
            "nullTerms": [
                {"term": b["key"], "score": b["score"], "count": b["doc_count"]}
                for b in aggs["nulls"]["sample"]["terms"]["buckets"]
            ],
            "fileDrawer": {
                "completed": registered,
                "unreported": unreported,
                "overdue": aggs["completed"]["overdue"]["doc_count"],
                "share": unreported / registered if registered else None,
            },
            "spin": {
                "eligible": aggs["spin_candidates"]["doc_count"],
                "disagreements": aggs["spin_candidates"]["spin"]["doc_count"],
            },
        }