"""Tests for the per-document toc.json view and the heading_level / tree_depth
fields surfaced in items.jsonl.

The TOC is derived from `DocumentItem.heading_level` (positive only) in
position order. PDF-style corpora (all section_headers at level 1) get a flat
list of siblings; HTML/markdown corpora with real heading hierarchy get a
nested tree. Items with no section_header at all produce `tree: []`.
"""

import asyncio
import json
from pathlib import PurePosixPath

import pytest

from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig
from haiku.rag.sandbox import AnalysisContext, Sandbox
from haiku.rag.store.models.chunk import Chunk
from haiku.rag.store.models.document import Document
from haiku.rag.store.models.document_item import DocumentItem


async def _empty_doc(client, *, uri: str, title: str) -> str:
    """Create a Document row directly via the repository (no chunking, no
    embedder) so these tests can run without a reachable embedding endpoint.
    Returns the document id."""
    doc = await client.document_repository.create(
        Document(content="x", uri=uri, title=title)
    )
    return doc.id


def _para(doc_id: str, pos: int, depth: int = 1) -> DocumentItem:
    return DocumentItem(
        document_id=doc_id,
        position=pos,
        self_ref=f"#/texts/{pos}",
        label="paragraph",
        text=f"para{pos}",
        page_numbers=[1],
        tree_depth=depth,
    )


def _header(
    doc_id: str, pos: int, level: int, text: str, depth: int = 1, page: int = 1
) -> DocumentItem:
    return DocumentItem(
        document_id=doc_id,
        position=pos,
        self_ref=f"#/texts/{pos}",
        label="section_header",
        text=text,
        page_numbers=[page],
        heading_level=level,
        tree_depth=depth,
    )


async def _read_vfs_text(sandbox: Sandbox, path: str) -> str:
    """Read a VFS file the way execute() does: the synchronous reader runs on a
    worker thread and bridges DB access back to the (free) calling loop."""
    vfs = await sandbox._build_vfs()
    sandbox._loop = asyncio.get_running_loop()
    return await asyncio.to_thread(vfs.path_read_text, PurePosixPath(path))


async def _read_toc(sandbox: Sandbox, doc_id: str) -> dict:
    raw = await _read_vfs_text(sandbox, f"/documents/{doc_id}/toc.json")
    return json.loads(raw)


async def _read_items_jsonl(sandbox: Sandbox, doc_id: str) -> list[dict]:
    raw = await _read_vfs_text(sandbox, f"/documents/{doc_id}/items.jsonl")
    return [json.loads(line) for line in raw.strip().splitlines()] if raw else []


def _flatten(tree: list[dict]) -> list[dict]:
    out = []
    for node in tree:
        out.append(node)
        out.extend(_flatten(node["children"]))
    return out


@pytest.mark.asyncio
class TestMetadataJson:
    """metadata.json is served by a reader callback, like the other VFS files."""

    async def test_metadata_reader_returns_document_fields(self, temp_db_path):
        async with HaikuRAG(temp_db_path, create=True) as client:
            doc_id = await _empty_doc(client, uri="test://meta", title="Meta Doc")

        sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
        raw = await _read_vfs_text(sandbox, f"/documents/{doc_id}/metadata.json")
        meta = json.loads(raw)

        assert meta["id"] == doc_id
        assert meta["title"] == "Meta Doc"
        assert meta["uri"] == "test://meta"


@pytest.mark.asyncio
class TestTocShape:
    """toc.json builds a section tree from heading_level + position."""

    async def test_multilevel_tree(self, temp_db_path):
        async with HaikuRAG(temp_db_path, create=True) as client:
            doc_id = await _empty_doc(client, uri="test://multi", title="TOC Test Doc")
            items = [
                _header(doc_id, 0, 1, "Intro"),
                _para(doc_id, 1),
                _header(doc_id, 2, 2, "Background"),
                _para(doc_id, 3),
                _header(doc_id, 4, 3, "Prior Work"),
                _para(doc_id, 5),
                _header(doc_id, 6, 2, "Approach"),
                _para(doc_id, 7),
                _header(doc_id, 8, 1, "Methods"),
                _para(doc_id, 9),
            ]
            await client.document_item_repository.create_items(doc_id, items)

        sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
        toc = await _read_toc(sandbox, doc_id)

        assert toc["doc_id"] == doc_id
        assert toc["title"] == "TOC Test Doc"
        tree = toc["tree"]
        # Two roots: Intro (children: Background>{Prior Work}, Approach) and Methods
        assert [n["title"] for n in tree] == ["Intro", "Methods"]
        intro = tree[0]
        assert intro["level"] == 1
        assert intro["item_range"] == [0, 8]  # ends at "Methods" position
        assert [c["title"] for c in intro["children"]] == ["Background", "Approach"]

        background = intro["children"][0]
        assert background["level"] == 2
        # Background covers positions 2..5; "Approach" begins at 6 (same-level sibling)
        assert background["item_range"] == [2, 6]
        assert [c["title"] for c in background["children"]] == ["Prior Work"]

        prior = background["children"][0]
        assert prior["level"] == 3
        # Prior Work has no descendants and the next same-or-shallower header is
        # "Approach" at level 2, position 6.
        assert prior["item_range"] == [4, 6]
        assert prior["children"] == []

        approach = intro["children"][1]
        assert approach["item_range"] == [6, 8]

        methods = tree[1]
        assert methods["item_range"] == [8, 10]  # to end of items
        assert methods["children"] == []

    async def test_flat_pdf_style(self, temp_db_path):
        """All section_headers at level 1 (PDF reality) -> flat sibling list."""
        async with HaikuRAG(temp_db_path, create=True) as client:
            doc_id = await _empty_doc(client, uri="test://pdf-style", title="Flat PDF")
            items = [
                _header(doc_id, 0, 1, "Chapter 1"),
                _para(doc_id, 1),
                _para(doc_id, 2),
                _header(doc_id, 3, 1, "Chapter 2"),
                _para(doc_id, 4),
                _header(doc_id, 5, 1, "Chapter 3"),
                _para(doc_id, 6),
            ]
            await client.document_item_repository.create_items(doc_id, items)

        sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
        toc = await _read_toc(sandbox, doc_id)

        tree = toc["tree"]
        assert [n["title"] for n in tree] == ["Chapter 1", "Chapter 2", "Chapter 3"]
        assert all(n["level"] == 1 and n["children"] == [] for n in tree)
        assert tree[0]["item_range"] == [0, 3]
        assert tree[1]["item_range"] == [3, 5]
        assert tree[2]["item_range"] == [5, 7]

    async def test_no_headers(self, temp_db_path):
        async with HaikuRAG(temp_db_path, create=True) as client:
            doc_id = await _empty_doc(
                client, uri="test://no-headers", title="No Headers"
            )
            await client.document_item_repository.create_items(
                doc_id, [_para(doc_id, i) for i in range(5)]
            )

        sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
        toc = await _read_toc(sandbox, doc_id)
        assert toc["tree"] == []

    async def test_skip_header_with_zero_level(self, temp_db_path):
        """A section_header with ``heading_level == 0`` is skipped."""
        async with HaikuRAG(temp_db_path, create=True) as client:
            doc_id = await _empty_doc(client, uri="test://zero", title="Zero Level")
            items = [
                _header(doc_id, 0, 1, "Real H1"),
                _para(doc_id, 1),
                _header(doc_id, 2, 0, "Pre-migration ghost"),
                _para(doc_id, 3),
            ]
            await client.document_item_repository.create_items(doc_id, items)

        sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
        toc = await _read_toc(sandbox, doc_id)
        titles = [n["title"] for n in _flatten(toc["tree"])]
        assert titles == ["Real H1"]
        assert toc["tree"][0]["item_range"] == [0, 4]

    async def test_node_shape_has_chunk_ids_not_position(self, temp_db_path):
        async with HaikuRAG(temp_db_path, create=True) as client:
            doc_id = await _empty_doc(client, uri="test://shape", title="Shape")
            await client.document_item_repository.create_items(
                doc_id, [_header(doc_id, 0, 1, "Only"), _para(doc_id, 1)]
            )

        sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
        toc = await _read_toc(sandbox, doc_id)
        node = toc["tree"][0]
        expected = {
            "self_ref",
            "level",
            "title",
            "page_numbers",
            "item_range",
            "chunk_ids",
            "children",
        }
        assert expected <= set(node)
        assert "position" not in node
        assert node["chunk_ids"] == []


@pytest.mark.asyncio
class TestTocChunkIdsAggregation:
    """toc.json nodes carry the union of chunk_ids covered by their item_range."""

    async def test_chunk_ids_union_over_item_range(self, temp_db_path, monkeypatch):
        async with HaikuRAG(temp_db_path, create=True) as client:
            doc_id = await _empty_doc(client, uri="test://chunks", title="Chunks")
            await client.document_item_repository.create_items(
                doc_id,
                [
                    _header(doc_id, 0, 1, "Intro"),
                    _para(doc_id, 1),
                    _para(doc_id, 2),
                    _header(doc_id, 3, 2, "Background"),
                    _para(doc_id, 4),
                    _header(doc_id, 5, 1, "Methods"),
                    _para(doc_id, 6),
                ],
            )

        from haiku.rag.store.repositories.chunk import ChunkRepository

        # self_ref → list[chunk_id]. Intro covers #/texts/0..2, Background
        # covers #/texts/3..4, Methods covers #/texts/5..6. Item at #/texts/2
        # belongs to two chunks (cA + cB) — verifying dedup-preserving-order.
        fake_index = {
            "#/texts/1": ["cA"],
            "#/texts/2": ["cA", "cB"],
            "#/texts/3": ["cB"],
            "#/texts/4": ["cC"],
            "#/texts/5": ["cD"],
            "#/texts/6": ["cD"],
        }

        async def fake_grouped(self, document_ids):
            return {doc_id: fake_index}

        monkeypatch.setattr(
            ChunkRepository,
            "get_chunk_ids_by_self_ref_grouped",
            fake_grouped,
        )

        sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
        toc = await _read_toc(sandbox, doc_id)

        intro = toc["tree"][0]
        assert intro["title"] == "Intro"
        # Intro spans positions 0..2 (header) plus its child Background's
        # range — item_range is [0, 5].
        assert intro["item_range"] == [0, 5]
        assert intro["chunk_ids"] == ["cA", "cB", "cC"]

        background = intro["children"][0]
        assert background["title"] == "Background"
        assert background["item_range"] == [3, 5]
        assert background["chunk_ids"] == ["cB", "cC"]

        methods = toc["tree"][1]
        assert methods["title"] == "Methods"
        assert methods["item_range"] == [5, 7]
        assert methods["chunk_ids"] == ["cD"]


@pytest.mark.asyncio
class TestTocCaching:
    """items + toc reads for a doc share one items fetch and one chunk-index fetch."""

    async def test_items_and_chunk_index_fetched_once_per_doc(
        self, temp_db_path, monkeypatch
    ):
        async with HaikuRAG(temp_db_path, create=True) as client:
            doc_id = await _empty_doc(client, uri="test://cache", title="Cache")
            await client.document_item_repository.create_items(
                doc_id,
                [_header(doc_id, 0, 1, "Only"), _para(doc_id, 1), _para(doc_id, 2)],
            )

        sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())

        from haiku.rag.store.repositories.chunk import ChunkRepository
        from haiku.rag.store.repositories.document_item import DocumentItemRepository

        items_calls = {"n": 0}
        chunk_calls = {"n": 0}
        original_items = DocumentItemRepository.get_all_items
        original_chunks = ChunkRepository.get_chunk_ids_by_self_ref_grouped

        async def counting_items(self, document_id):
            items_calls["n"] += 1
            return await original_items(self, document_id)

        async def counting_chunks(self, document_ids):
            chunk_calls["n"] += 1
            return await original_chunks(self, document_ids)

        monkeypatch.setattr(DocumentItemRepository, "get_all_items", counting_items)
        monkeypatch.setattr(
            ChunkRepository, "get_chunk_ids_by_self_ref_grouped", counting_chunks
        )

        # Items + toc share `_doc_items` and `_doc_chunk_index`. Four reads →
        # one items fetch + one chunk-index fetch.
        _ = await _read_toc(sandbox, doc_id)
        _ = await _read_items_jsonl(sandbox, doc_id)
        _ = await _read_toc(sandbox, doc_id)
        _ = await _read_items_jsonl(sandbox, doc_id)

        assert items_calls["n"] == 1
        assert chunk_calls["n"] == 1


@pytest.mark.asyncio
class TestItemsJsonlSurfacesNewFields:
    """items.jsonl row shape: heading_level is always present (0 on non-headers);
    chunk_ids surfaces each item's containing chunks; position and tree_depth
    are not exposed."""

    async def test_jsonl_row_shape(self, temp_db_path):
        async with HaikuRAG(temp_db_path, create=True) as client:
            doc_id = await _empty_doc(client, uri="test://jsonl-fields", title="Fields")
            items = [
                _header(doc_id, 0, 1, "H1", depth=2),
                _para(doc_id, 1, depth=3),
                _header(doc_id, 2, 2, "H2", depth=4),
            ]
            await client.document_item_repository.create_items(doc_id, items)

        sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
        rows = await _read_items_jsonl(sandbox, doc_id)

        assert len(rows) == 3
        assert rows[0]["heading_level"] == 1
        assert rows[1]["heading_level"] == 0
        assert rows[2]["heading_level"] == 2
        for r in rows:
            expected = {
                "self_ref",
                "label",
                "text",
                "page_numbers",
                "heading_level",
                "chunk_ids",
            }
            assert expected <= set(r)
            assert "position" not in r
            assert "tree_depth" not in r


@pytest.mark.asyncio
class TestVfsReadPaths:
    """The synchronous VFS readers bridge back to the event loop; drive them
    through a worker thread the way execute() does."""

    async def test_content_txt_is_read_lazily_per_document(self, temp_db_path):
        async with HaikuRAG(temp_db_path, create=True) as client:
            doc = await client.document_repository.create(
                Document(content="the stored body", uri="test://body", title="Body")
            )
            doc_id = doc.id

        sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())

        # Build the VFS first, then change the stored content. A lazy
        # CallbackFile reads through at access time and sees the new body;
        # an eager MemoryFile mount would have captured the old one.
        vfs = await sandbox._build_vfs()
        sandbox._loop = asyncio.get_running_loop()

        # Rewrite via the repository rather than client.update_document: the
        # latter re-chunks and re-embeds, which this file deliberately avoids
        # so these tests need no embedding endpoint.
        async with HaikuRAG(temp_db_path, create=False) as client:
            doc.content = "the rewritten body"
            await client.document_repository.update(doc)

        content = await asyncio.to_thread(
            vfs.path_read_text, PurePosixPath(f"/documents/{doc_id}/content.txt")
        )

        assert content == "the rewritten body"

    async def test_document_files_are_read_only(self, temp_db_path):
        async with HaikuRAG(temp_db_path, create=True) as client:
            doc = await client.document_repository.create(
                Document(content="x", uri="test://ro", title="RO")
            )
            doc_id = doc.id

        sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
        vfs = await sandbox._build_vfs()
        sandbox._loop = asyncio.get_running_loop()

        with pytest.raises(PermissionError, match="read-only"):
            await asyncio.to_thread(
                vfs.path_write_text,
                PurePosixPath(f"/documents/{doc_id}/content.txt"),
                "nope",
            )

    async def test_chunks_jsonl_lists_chunks_in_order_with_their_metadata(
        self, temp_db_path
    ):
        """One row per chunk, in chunk order, carrying the stored metadata as
        is; the second read of a document is served from the sandbox's cache."""
        config = AppConfig()
        dim = config.embeddings.model.vector_dim
        async with HaikuRAG(temp_db_path, create=True) as client:
            doc_id = await _empty_doc(client, uri="test://paras", title="Paras")
            for order, para_no in enumerate(["13", "14"]):
                await client.chunk_repository.create(
                    Chunk(
                        document_id=doc_id,
                        content=f"Paragraph {para_no}.",
                        embedding=[0.1] * dim,
                        order=order,
                        metadata={"para_no": para_no, "doc_item_refs": []},
                    )
                )

        sandbox = Sandbox(temp_db_path, config, AnalysisContext())
        first = await _read_vfs_text(sandbox, f"/documents/{doc_id}/chunks.jsonl")
        rows = [json.loads(line) for line in first.split("\n")]

        assert [row["metadata"]["para_no"] for row in rows] == ["13", "14"]
        assert all(set(row) == {"chunk_id", "metadata"} for row in rows)
        assert rows[0]["metadata"] == {"para_no": "13", "doc_item_refs": []}
        assert sandbox._chunks_jsonl_cache[doc_id] == first
        assert (
            await _read_vfs_text(sandbox, f"/documents/{doc_id}/chunks.jsonl") == first
        )

    async def test_item_range_is_a_line_slice_into_items_jsonl(self, temp_db_path):
        """`item_range` indexes lines of items.jsonl, as documented, not item
        positions: a gap in positions must not pull the next heading into a
        section."""
        async with HaikuRAG(temp_db_path, create=True) as client:
            doc_id = await _empty_doc(client, uri="test://slice", title="Slice")
            items = [
                _header(doc_id, 0, 1, "Intro"),
                _para(doc_id, 1),
                _header(doc_id, 3, 1, "Methods"),
                _para(doc_id, 4),
            ]
            await client.document_item_repository.create_items(doc_id, items)

        sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
        toc = await _read_toc(sandbox, doc_id)
        raw = await _read_vfs_text(sandbox, f"/documents/{doc_id}/items.jsonl")
        lines = raw.split("\n")

        intro, methods = toc["tree"]
        assert intro["item_range"] == [0, 2]
        assert methods["item_range"] == [2, 4]
        start, end = intro["item_range"]
        assert [json.loads(line)["self_ref"] for line in lines[start:end]] == [
            "#/texts/0",
            "#/texts/1",
        ]

    async def test_toc_skips_gaps_in_item_positions(self, temp_db_path):
        """Positions need not be contiguous — a heading's span may cover
        positions that carry no item."""
        async with HaikuRAG(temp_db_path, create=True) as client:
            doc_id = await _empty_doc(client, uri="test://gaps", title="Gaps")
            # Positions 1 and 2 are absent between the header and the paragraph.
            items = [
                _header(doc_id, 0, 1, "Intro"),
                _para(doc_id, 3),
            ]
            await client.document_item_repository.create_items(doc_id, items)

        sandbox = Sandbox(temp_db_path, AppConfig(), AnalysisContext())
        toc = await _read_toc(sandbox, doc_id)

        assert [n["title"] for n in toc["tree"]] == ["Intro"]