"""The load stage, exercised against the in-memory store with no server.

Two classes of failure are silent in a vector database and are what most of
these tests aim at. Loading the same corpus twice can leave two copies of a
dose, and every query still succeeds — so idempotency is asserted by point
count, not by inspecting the upsert calls. Mixing two corpus generations or two
models into one collection also raises nothing at query time; every search
returns *something*, just from the wrong material. That is what the manifest
gate exists to make loud, and there is a test per way it can be violated.

The last test runs over the real `chunks.jsonl` when it is present, because a
provenance rule that only holds for hand-written records is not evidence.
"""
import json
from pathlib import Path

import pytest

from ingestion.load import (
    ChunkLoader,
    CollectionSpec,
    CorpusManifest,
    CorpusMismatch,
    InMemoryVectorStore,
    PointCountMismatch,
    build_point,
    corpus_sha256,
    count_chunks,
    iter_chunk_records,
    manifest_collection,
    point_id_for,
    read_manifest,
    validate_chunk_record,
)

DIMENSIONS = 4
COLLECTION = "duoc_thu_chunks"
CORPUS_SHA = "a" * 64
OTHER_SHA = "b" * 64
MODEL = "amazon.titan-embed-text-v2:0"

REAL_CHUNKS = (
    Path(__file__).resolve().parents[1] / "data" / "processed" / "chunks.jsonl"
)


def chunk_record(chunk_id="abacavir__lieu_luong__0", **overrides):
    record = {
        "schema_version": 4,
        "chunk_id": chunk_id,
        "drug_id": "abacavir",
        "drug_name": "ABACAVIR",
        "section_key": "lieu_luong_va_cach_dung",
        "section_display_name": "Liều lượng và cách dùng",
        "text": "Người lớn: 300 mg, hai lần mỗi ngày.",
        "source_text": "Người lớn: 300 mg, hai lần mỗi ngày.",
        "heading_physical_page": 100,
        "source_page_range": [100, 102],
        "printed_page_range": [101, 103],
        "atc_codes": ["J05AF06"],
        "part_index": 0,
        "part_count": 1,
        "est_tokens": 14,
        "oversized": False,
        "chunk_kind": "prose",
        "attachments": [],
        "has_quarantined_content": False,
    }
    record.update(overrides)
    return record


def vector(seed=0.1):
    return [seed] * DIMENSIONS


def manifest(**overrides):
    values = {
        "corpus_sha256": CORPUS_SHA,
        "chunk_count": 3,
        "model_id": MODEL,
        "dimensions": DIMENSIONS,
        "input_kind": "document",
        "provider": "titan-v2",
    }
    values.update(overrides)
    return CorpusManifest(**values)


def spec(**overrides):
    values = {"name": COLLECTION, "vector_size": DIMENSIONS}
    values.update(overrides)
    return CollectionSpec(**values)


def loader(store, **overrides):
    return ChunkLoader(
        store,
        overrides.pop("spec", spec()),
        overrides.pop("manifest", manifest()),
        **overrides,
    )


def pairs(count=3):
    return [
        (chunk_record(chunk_id=f"drug__section__{i}"), vector(0.1 * (i + 1)))
        for i in range(count)
    ]


# --- A5: derived ids and idempotency -------------------------------------


def test_point_id_is_derived_from_chunk_id_and_is_stable():
    first = point_id_for("abacavir__lieu_luong__0")
    second = point_id_for("abacavir__lieu_luong__0")

    assert first == second
    assert first != point_id_for("abacavir__lieu_luong__1")


def test_point_id_refuses_an_empty_chunk_id():
    with pytest.raises(ValueError, match="chunk_id is required"):
        point_id_for("   ")


def test_loading_the_same_corpus_twice_leaves_the_point_count_unchanged():
    store = InMemoryVectorStore()
    data = pairs(3)

    first = loader(store).load(data)
    second = loader(store).load(data)

    assert first.collection_created is True
    assert second.collection_created is False
    assert first.collection_count == 3
    assert second.collection_count == 3, "a re-run duplicated points"
    assert second.points_upserted == 3


def test_a_reloaded_chunk_overwrites_its_own_point_rather_than_adding_one():
    store = InMemoryVectorStore()
    record = chunk_record()
    loader(store).load([(record, vector(0.1))])

    edited = chunk_record(text="Người lớn: 600 mg, một lần mỗi ngày.")
    loader(store).load([(edited, vector(0.9))])

    assert store.count(COLLECTION) == 1
    stored = store.retrieve(COLLECTION, point_id_for(record["chunk_id"]))
    assert stored.payload["text"] == "Người lớn: 600 mg, một lần mỗi ngày."
    assert stored.vector == vector(0.9)


def test_records_are_upserted_in_batches_of_the_configured_size():
    store = InMemoryVectorStore()

    report = loader(store, batch_size=2).load(pairs(5))

    assert report.batches == 3
    assert report.points_upserted == 5
    assert report.collection_count == 5


def test_batch_size_must_be_positive():
    with pytest.raises(ValueError, match="batch_size must be positive"):
        ChunkLoader(InMemoryVectorStore(), spec(), manifest(), batch_size=0)


# --- A4: collection shape and payload ------------------------------------


def test_collection_is_created_with_the_declared_size_and_payload_indexes():
    store = InMemoryVectorStore()

    loader(store).load(pairs(1))

    assert store.spec(COLLECTION).vector_size == DIMENSIONS
    assert store.spec(COLLECTION).distance == "Cosine"
    indexed = dict(store.indexed_fields(COLLECTION))
    assert indexed["drug_id"] == "keyword"
    assert indexed["section_key"] == "keyword"
    assert indexed["atc_codes"] == "keyword"
    assert indexed["chunk_kind"] == "keyword"
    assert indexed["has_quarantined_content"] == "bool"


def test_payload_carries_every_provenance_field_of_the_chunk_record():
    store = InMemoryVectorStore()
    record = chunk_record()

    loader(store).load([(record, vector())])

    payload = store.retrieve(COLLECTION, point_id_for(record["chunk_id"])).payload
    assert payload == record


def test_a_field_added_by_a_future_chunker_flows_through_untouched():
    store = InMemoryVectorStore()
    record = chunk_record(
        population_tags=["Người lớn", "Suy thận"], printed_page_range=[101, 103]
    )

    loader(store).load([(record, vector())])

    payload = store.retrieve(COLLECTION, point_id_for(record["chunk_id"])).payload
    assert payload["population_tags"] == ["Người lớn", "Suy thận"]
    assert payload["printed_page_range"] == [101, 103]


@pytest.mark.parametrize(
    "field",
    [
        "chunk_id",
        "drug_id",
        "section_key",
        "source_page_range",
        "printed_page_range",
        "text",
    ],
)
def test_a_chunk_missing_a_required_provenance_field_is_refused(field):
    with pytest.raises(ValueError, match="missing required provenance fields"):
        validate_chunk_record(chunk_record(**{field: None}))


# --- failing closed on incomplete provenance ------------------------------
#
# Every case below passed an earlier version of this validator. The cost of
# that is not an exception at load time — it is paying for an embedding run and
# then discovering every answer abstains because the chunks cannot be cited.


@pytest.mark.parametrize("field", ["source_page_range", "printed_page_range"])
def test_an_empty_page_range_is_missing_not_present(field):
    """`[] in (None, "")` is False, which is exactly how this slipped through."""
    with pytest.raises(ValueError, match="missing required provenance fields"):
        validate_chunk_record(chunk_record(**{field: []}))


def test_an_unknown_old_or_future_schema_is_refused_fail_closed():
    for version in (3, 5):
        with pytest.raises(ValueError, match="supports exactly v4"):
            validate_chunk_record(chunk_record(schema_version=version))


def test_a_chunk_with_no_schema_version_at_all_is_refused():
    record = chunk_record()
    del record["schema_version"]
    with pytest.raises(ValueError, match="declares schema_version None"):
        validate_chunk_record(record)


@pytest.mark.parametrize(
    "value", [[101], [101, 102, 103], "101-103", 101, {"start": 101}]
)
def test_a_page_range_that_is_not_a_pair_is_refused(value):
    with pytest.raises(ValueError, match=r"expected a \[start, end\] pair"):
        validate_chunk_record(chunk_record(printed_page_range=value))


def test_a_page_range_running_backwards_is_refused():
    with pytest.raises(ValueError, match="running backwards"):
        validate_chunk_record(chunk_record(printed_page_range=[103, 101]))


def test_a_non_integer_page_is_refused():
    with pytest.raises(ValueError, match="non-integer page"):
        validate_chunk_record(chunk_record(printed_page_range=[101.5, 103]))


def test_boolean_pages_are_not_accepted_as_python_integers():
    with pytest.raises(ValueError, match="non-integer page"):
        validate_chunk_record(chunk_record(printed_page_range=[False, True]))


@pytest.mark.parametrize(
    ("field", "value"),
    [("source_page_range", [-1, 0]), ("printed_page_range", [0, 1])],
)
def test_page_ranges_reject_impossible_lower_bounds(field, value):
    with pytest.raises(ValueError, match="pages must start"):
        validate_chunk_record(chunk_record(**{field: value}))


def test_page_zero_and_false_are_values_not_absences():
    """Physical pages are 0-indexed; a falsiness test would reject real records."""
    validate_chunk_record(
        chunk_record(
            heading_physical_page=0,
            source_page_range=[0, 0],
            printed_page_range=[1, 1],
            has_quarantined_content=False,
            oversized=False,
            part_index=0,
        )
    )


def test_a_wrong_sized_vector_is_refused_before_anything_is_upserted():
    store = InMemoryVectorStore()

    with pytest.raises(ValueError, match="expects 4"):
        loader(store).load([(chunk_record(), [0.1, 0.2])])

    assert store.count(COLLECTION) == 0


def test_manifest_dimensions_must_agree_with_the_collection_spec():
    with pytest.raises(ValueError, match="manifest declares 8 dimensions"):
        ChunkLoader(InMemoryVectorStore(), spec(), manifest(dimensions=8))


def test_collection_spec_refuses_a_nonpositive_vector_size():
    with pytest.raises(ValueError, match="vector_size must be positive"):
        CollectionSpec(name=COLLECTION, vector_size=0)


# --- A6: the corpus binding gate -----------------------------------------


def test_the_manifest_lives_beside_the_data_so_the_point_count_stays_exact():
    store = InMemoryVectorStore()

    loader(store).load(pairs(3))

    assert store.count(COLLECTION) == 3, "the manifest must not inflate the count"
    assert store.count(manifest_collection(COLLECTION)) == 1
    stored = read_manifest(store, COLLECTION)
    assert stored.corpus_sha256 == CORPUS_SHA
    assert stored.model_id == MODEL
    assert stored.provider == "titan-v2"


def test_a_second_corpus_generation_is_refused_and_nothing_is_written():
    store = InMemoryVectorStore()
    loader(store).load(pairs(3))

    with pytest.raises(CorpusMismatch, match="does not match"):
        loader(store, manifest=manifest(corpus_sha256=OTHER_SHA)).load(pairs(2))

    assert store.count(COLLECTION) == 3
    assert read_manifest(store, COLLECTION).corpus_sha256 == CORPUS_SHA


def test_a_second_model_is_refused_even_when_the_corpus_matches():
    store = InMemoryVectorStore()
    loader(store).load(pairs(1))

    with pytest.raises(CorpusMismatch, match="cohere.embed-v4:0"):
        loader(store, manifest=manifest(model_id="cohere.embed-v4:0")).load(pairs(1))


def test_a_query_subspace_vector_is_refused_for_a_document_collection():
    store = InMemoryVectorStore()
    loader(store).load(pairs(1))

    with pytest.raises(CorpusMismatch, match="input kind query"):
        loader(store, manifest=manifest(input_kind="query")).load(pairs(1))


def test_a_dimension_change_is_refused():
    store = InMemoryVectorStore()
    loader(store).load(pairs(1))

    eight = CollectionSpec(name=COLLECTION, vector_size=8)
    with pytest.raises(CorpusMismatch, match="8 dimensions"):
        ChunkLoader(store, eight, manifest(dimensions=8)).load(
            [(chunk_record(), [0.1] * 8)]
        )


def test_an_existing_collection_with_no_manifest_is_refused():
    store = InMemoryVectorStore()
    store.create_collection(spec())
    store.upsert(COLLECTION, [build_point(chunk_record(), vector())])

    with pytest.raises(CorpusMismatch, match="has no manifest"):
        loader(store).load(pairs(1))


def test_the_same_corpus_and_model_is_allowed_through():
    store = InMemoryVectorStore()
    loader(store).load(pairs(3))

    report = loader(store).load(pairs(3))

    assert report.collection_count == 3
    assert report.count_matches is True


# --- the v1 point-count gate ---------------------------------------------


def test_assert_point_count_passes_when_every_chunk_has_exactly_one_point():
    store = InMemoryVectorStore()
    active = loader(store)
    active.load(pairs(3))

    assert active.assert_point_count(3) == 3


def test_assert_point_count_raises_when_the_collection_is_short():
    store = InMemoryVectorStore()
    active = loader(store)
    active.load(pairs(3))

    with pytest.raises(PointCountMismatch, match="holds 3 points but the corpus has 4"):
        active.assert_point_count(4)


# --- mode A: exhaustive filter retrieval ---------------------------------


def section_pairs(drug_id, section_key, parts):
    return [
        (
            chunk_record(
                chunk_id=f"{drug_id}__{section_key}__{i}",
                drug_id=drug_id,
                section_key=section_key,
                part_index=i,
                part_count=parts,
            ),
            vector(0.1),
        )
        for i in range(parts)
    ]


def test_a_filter_returns_every_part_of_a_section_not_a_top_k():
    """The rule mode A exists for: two of five contraindications is worse than none."""
    store = InMemoryVectorStore()
    loader(store).load(
        section_pairs("metformin", "chong_chi_dinh", 5)
        + section_pairs("metformin", "lieu_luong_va_cach_dung", 3)
        + section_pairs("pantoprazol", "chong_chi_dinh", 2)
    )

    found = store.find_by_payload(
        COLLECTION, {"drug_id": "metformin", "section_key": "chong_chi_dinh"}
    )

    assert len(found) == 5
    assert sorted(p.payload["part_index"] for p in found) == [0, 1, 2, 3, 4]
    assert {p.payload["drug_id"] for p in found} == {"metformin"}


def test_a_filter_never_leaks_a_neighbouring_drugs_section():
    store = InMemoryVectorStore()
    loader(store).load(
        section_pairs("pantoprazol", "chong_chi_dinh", 2)
        + section_pairs("omeprazol", "chong_chi_dinh", 2)
    )

    found = store.find_by_payload(
        COLLECTION, {"drug_id": "pantoprazol", "section_key": "chong_chi_dinh"}
    )

    assert len(found) == 2
    assert {p.payload["drug_id"] for p in found} == {"pantoprazol"}


def test_a_list_valued_field_matches_on_any_element():
    store = InMemoryVectorStore()
    record = chunk_record(atc_codes=["A10BA02", "A10BD20"])
    loader(store).load([(record, vector())])

    assert len(store.find_by_payload(COLLECTION, {"atc_codes": "A10BD20"})) == 1
    assert len(store.find_by_payload(COLLECTION, {"atc_codes": "J05AF06"})) == 0


def test_a_filter_matching_nothing_returns_empty_rather_than_raising():
    store = InMemoryVectorStore()
    loader(store).load(pairs(2))

    assert store.find_by_payload(COLLECTION, {"drug_id": "khong_ton_tai"}) == []


def test_a_filter_with_no_condition_is_refused():
    store = InMemoryVectorStore()
    loader(store).load(pairs(1))

    with pytest.raises(ValueError, match="at least one condition"):
        store.find_by_payload(COLLECTION, {})


def test_parts_reassemble_in_order_into_the_whole_section():
    store = InMemoryVectorStore()
    bodies = ["Phần một.", "Phần hai.", "Phần ba."]
    records = [
        (
            chunk_record(
                chunk_id=f"metformin__chong_chi_dinh__{i}",
                drug_id="metformin",
                section_key="chong_chi_dinh",
                text=body,
                part_index=i,
                part_count=len(bodies),
            ),
            vector(),
        )
        for i, body in enumerate(bodies)
    ]
    loader(store).load(records)

    found = store.find_by_payload(
        COLLECTION, {"drug_id": "metformin", "section_key": "chong_chi_dinh"}
    )
    ordered = sorted(found, key=lambda p: p.payload["part_index"])

    assert [p.payload["text"] for p in ordered] == bodies
    assert {p.payload["part_count"] for p in found} == {3}


# --- corpus digest and reading -------------------------------------------


def test_corpus_sha256_changes_when_a_single_byte_changes(tmp_path):
    path = tmp_path / "chunks.jsonl"
    path.write_text(json.dumps(chunk_record()) + "\n", encoding="utf-8")
    before = corpus_sha256(path)

    path.write_text(
        json.dumps(chunk_record(text="Người lớn: 301 mg.")) + "\n", encoding="utf-8"
    )

    assert corpus_sha256(path) != before
    assert len(before) == 64


def test_the_same_data_hashes_the_same_under_crlf_and_lf(tmp_path):
    """A gate that cries wolf gets switched off.

    A raw-byte digest made a Windows CRLF checkout and a Linux LF checkout of
    identical data disagree, so A6 would refuse a CI load against the very
    corpus it was built from.
    """
    body = json.dumps(chunk_record()) + "\n" + json.dumps(chunk_record("b")) + "\n"
    lf = tmp_path / "lf.jsonl"
    crlf = tmp_path / "crlf.jsonl"
    lf.write_bytes(body.encode("utf-8"))
    crlf.write_bytes(body.replace("\n", "\r\n").encode("utf-8"))

    assert corpus_sha256(lf) == corpus_sha256(crlf)
    assert crlf.stat().st_size > lf.stat().st_size, "the files really do differ"


def test_blank_lines_are_skipped_and_bad_json_names_its_line(tmp_path):
    path = tmp_path / "chunks.jsonl"
    path.write_text(
        json.dumps(chunk_record()) + "\n\n" + json.dumps(chunk_record("b")) + "\n",
        encoding="utf-8",
    )

    assert count_chunks(path) == 2

    broken = tmp_path / "broken.jsonl"
    broken.write_text(json.dumps(chunk_record()) + "\n{oops\n", encoding="utf-8")
    with pytest.raises(ValueError, match="line 2 is not valid JSON"):
        list(iter_chunk_records(broken))


# --- the real artifact ----------------------------------------------------


@pytest.mark.skipif(
    not REAL_CHUNKS.exists(), reason="chunks.jsonl has not been generated"
)
def test_every_real_chunk_satisfies_the_loader_provenance_contract():
    """Whole-artifact scope: all records in `data/processed/chunks.jsonl`."""
    seen_ids = set()
    seen_points = set()
    total = 0
    for record in iter_chunk_records(REAL_CHUNKS):
        validate_chunk_record(record)
        point = point_id_for(record["chunk_id"])
        assert point not in seen_points, f"point id collision on {record['chunk_id']}"
        seen_points.add(point)
        seen_ids.add(record["chunk_id"])
        total += 1

    assert total == len(seen_ids), "duplicate chunk_id in the artifact"
    assert total == len(seen_points)
    # A floor, not the exact count: the corpus is regenerated as `segment/`
    # changes, but a truncated or half-written artifact must not pass as
    # whole-artifact evidence. Measured 15,066 records on 2026-08-04.
    assert total > 10_000, f"chunks.jsonl looks truncated: only {total} records"