"""
Retrieval tests using bundled apple tree corpus.
"""

import tempfile
from pathlib import Path

import pytest

from justembed.chunker import chunk_text
from justembed.embedder import E5Embedder
from justembed import db


# Apple tree corpus (from spec 09-examples.md)
APPLE_TREE_TEXT = """
Our apple tree stands at the edge of the garden, leaning slightly toward the morning sun. Its branches spread wide, forming a soft canopy that changes character with every season. In spring, the tree is covered in pale blossoms that sway gently in the breeze.

By early summer, the blossoms give way to small green fruits. These apples grow slowly, absorbing sunlight and moisture from the soil. Their color deepens as the weeks pass, shifting from light green to a warm yellowish tone.

When autumn arrives, the apples reach their full size and turn a bright red. The fruit hangs heavily from the branches, and the tree becomes a favorite spot for birds looking for a sweet snack. The red apples stand out vividly against the fading leaves.

The tree itself is old but resilient. Its bark is rough, and its trunk twists in a way that tells a long story of storms and seasons. Despite its age, it continues to produce fruit every year, surprising us with its consistency.

Harvesting the apples is a small celebration in our home. We gather the red fruit in baskets, wash them carefully, and enjoy them fresh or baked into pies. The apple tree has become a symbol of patience, growth, and the quiet rewards of nature.
""".strip()


EXAMPLE_QUERIES = [
    "Which fruit on the tree turns red in autumn",
    "Tell me about the color changes of the apples",
    "What happens to the blossoms in summer",
    "Describe the tree's appearance through the seasons",
    "Where do birds gather for sweet snacks",
]


def test_chunker_produces_chunks():
    """Chunker produces non-empty chunks from apple tree text."""
    chunks = chunk_text(APPLE_TREE_TEXT, file="apple_tree.txt")
    assert len(chunks) >= 1
    assert len(chunks) <= 500
    for c in chunks:
        assert "file" in c and c["file"] == "apple_tree.txt"
        assert "chunk_index" in c
        assert "text" in c and len(c["text"]) > 0


def test_embedder_returns_384_dim():
    """E5 embedder returns 384-dim vectors."""
    emb = E5Embedder()
    vecs = emb.embed(["hello world"])
    assert len(vecs) == 1
    assert len(vecs[0]) == 384


def test_embed_query_returns_384_dim():
    """E5 embed_query returns 384-dim vector."""
    emb = E5Embedder()
    vec = emb.embed_query("red fruit")
    assert len(vec) == 384


def test_full_retrieval_pipeline():
    """Full pipeline: chunk -> embed -> store -> query -> retrieve."""
    from justembed.config import ensure_workspace_structure

    with tempfile.TemporaryDirectory() as tmpdir:
        workspace = Path(tmpdir)
        ensure_workspace_structure(workspace)

        # Create KB
        db.create_kb(workspace, "test_kb")

        # Chunk
        chunks = chunk_text(APPLE_TREE_TEXT, file="apple_tree.txt")
        assert len(chunks) >= 1

        # Embed
        emb = E5Embedder()
        embeddings = emb.embed([c["text"] for c in chunks])

        # Store
        conn = db.get_connection(workspace, "test_kb")
        db.insert_chunks(conn, "apple_tree.txt", chunks, embeddings)
        conn.close()

        # Query: "Which fruit on the tree turns red in autumn"
        q_emb = emb.embed_query("Which fruit on the tree turns red in autumn")
        conn = db.get_connection(workspace, "test_kb")
        results = db.query_chunks(conn, q_emb, top_k=3, kb_name="test_kb")
        conn.close()

        assert len(results) >= 1
        # Top result should mention red/apple/autumn
        top_text = results[0]["text"].lower()
        assert any(w in top_text for w in ["red", "apple", "autumn", "fruit"])


def test_example_queries_retrieve_relevant():
    """Example queries retrieve relevant chunks (top-3 contains expected content)."""
    with tempfile.TemporaryDirectory() as tmpdir:
        workspace = Path(tmpdir)
        (workspace / "models" / "core").mkdir(parents=True, exist_ok=True)
        (workspace / "models" / "custom").mkdir(parents=True, exist_ok=True)
        (workspace / "kb").mkdir(parents=True, exist_ok=True)
        (workspace / "log").mkdir(parents=True, exist_ok=True)

        db.create_kb(workspace, "apple")
        chunks = chunk_text(APPLE_TREE_TEXT, file="apple_tree.txt")
        emb = E5Embedder()
        embeddings = emb.embed([c["text"] for c in chunks])
        conn = db.get_connection(workspace, "apple")
        db.insert_chunks(conn, "apple_tree.txt", chunks, embeddings)
        conn.close()

        # Query each example
        conn = db.get_connection(workspace, "apple")
        for query in EXAMPLE_QUERIES[:2]:  # First 2 to keep test fast
            q_emb = emb.embed_query(query)
            results = db.query_chunks(conn, q_emb, top_k=3, kb_name="apple")
            assert len(results) >= 1
            combined = " ".join(r["text"].lower() for r in results)
            # At least one result should have relevant content
            assert len(combined) > 20
        conn.close()
