# -*- coding: utf-8 -*-
"""
Updated Pytest for Refactored BibleRAG (June 2026)
@author: David DiPaola (assisted by Grok)

!pytest test_BibleRAG_refactored.py -v -s (inside of spyder console, -v gives result - pass / fail, -s displays printed items)
"""
import os
import warnings
import time
import pytest
import numpy as np
from datasets import Dataset

# Suppress warnings aggressively
os.environ["PYTEST_ADDOPTS"] = "-p no:warnings --disable-warnings"
os.environ["PYTHONWARNINGS"] = "ignore"

import pytest
from Bible_Rag_Refactor import BibleRAG, NotesManager  # Adjust filename if needed

from langchain_community.llms import LlamaCpp
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_core.callbacks.base import BaseCallbackHandler


def pytest_configure(config):
    config.option.warnfilters = ["ignore::DeprecationWarning", "ignore::UserWarning"]
    config.pluginmanager.set_blocked("warnings")


# ==================== SPY CALLBACK ====================
class VerboseCallback(BaseCallbackHandler):
    def on_llm_start(self, serialized, prompts, **kwargs):
        print("\n" + "="*80)
        print(">>> RAW PROMPT TO LLM:")
        print(prompts[0])
        print("="*80)

    def on_llm_end(self, response, **kwargs):
        print("\n" + "="*80)
        print(">>> RAW LLM RESPONSE:")
        print(response.generations[0][0].text)
        print("="*80)


# ==================== FIXTURES ====================
@pytest.fixture(scope="module")
def rag_instance():
    """Main BibleRAG instance (loads once per test session)."""
    rag = BibleRAG()  # Uses defaults from your refactored class
    return rag


@pytest.fixture(scope="module")
def notes_manager():
    """NotesManager fixture."""
    nm = NotesManager(db_path="E:/Python/bible/sovereign_notes.db")
    # Optional: Add a test note
    nm.add_note(
        content="I realized today that peace isn't the absence of conflict, but the presence of God in it.",
        tags="peace, trials",
        verses="John 14:27, Mark 4:38"
    )
    return nm


@pytest.fixture(scope="module")
def engines(rag_instance):
    llm = LlamaCpp(
        model_path=rag_instance.llm_model_path,
        n_ctx=4096,
        n_batch=8,
        n_gpu_layers=-1,
        temperature=0.1,
        repeat_penalty=1.1,
        verbose=False,
        f16_kv=False,
        model_kwargs={
            "offload_kqv": False,
            "flash_attn": False,
        }
    )

    emb = HuggingFaceEmbeddings(
        model_name=rag_instance.embed_model_name,
        model_kwargs={"device": "cpu"}
    )

    return {
        "llm": LangchainLLMWrapper(llm),
        "emb": LangchainEmbeddingsWrapper(emb)
    }

@pytest.fixture
def sample_query():
    return "What does the Bible and my personal notes say about finding peace in trials?"


# ==================== UNIT TESTS ====================
def test_bible_data_loaded(rag_instance):
    assert len(rag_instance.bible_data) > 30000
    assert rag_instance.index.ntotal == len(rag_instance.bible_data)
    sample = rag_instance.bible_data[0]
    assert "metadata" in sample
    assert sample["metadata"]["book"] == "Genesis"
    print("✓ Bible data loaded correctly")


def test_notes_indexing(rag_instance, notes_manager):
    rag_instance.index_notes(notes_manager)
    assert rag_instance.notes_index is not None
    assert len(rag_instance.notes_content) > 0
    print("✓ Notes indexing successful")


def test_retrieve_bible(rag_instance, sample_query):
    results = rag_instance.retrieve_bible(sample_query, k=5, rerank=True)
    assert len(results) == 5
    assert "content" in results[0] or "metadata" in results[0]
    assert "bi_score" in results[0]
    print("✓ Bible retrieval works")


def test_retrieve_notes(rag_instance, notes_manager, sample_query):
    rag_instance.index_notes(notes_manager)
    results = rag_instance.retrieve_notes(sample_query, k_notes=2)
    assert isinstance(results, list)
    if results:
        assert "notes_content" in results[0]
    print("✓ Notes retrieval works")


def test_context_expansion(rag_instance, sample_query):
    docs = rag_instance.retrieve_bible(sample_query, k=3)
    context = rag_instance._expand_and_merge_context(docs, window=3)
    assert isinstance(context, str)
    assert "** Main verse:" in context
    assert len(context) > 100
    print("✓ Context expansion works")


def test_prompt_template(rag_instance):
    template = rag_instance._create_prompt_template()
    assert template is not None
    print(f"✓ Prompt template loaded for model_type: {rag_instance.model_type}")


def test_format_notes(rag_instance, notes_manager):
    rag_instance.index_notes(notes_manager)
    notes = rag_instance.retrieve_notes("peace", k_notes=1)
    formatted = rag_instance.format_notes_for_llm(notes)
    assert isinstance(formatted, str)
    print("✓ Notes formatting for LLM works")


# ==================== PERFORMANCE ====================
def test_retrieval_latency(rag_instance, sample_query):
    start = time.time()
    _ = rag_instance.retrieve_bible(sample_query, k=5, rerank=True)
    latency = time.time() - start
    print(f"Retrieval latency: {latency:.2f}s")
    assert latency < 5.0  # Adjust threshold to your hardware


def test_end_to_end_query(rag_instance, notes_manager, sample_query):
    rag_instance.index_notes(notes_manager)
    rag_instance.initialize_llm()

    start = time.time()
    response = rag_instance.query(sample_query, b_k=5, n_k=2)
    latency = time.time() - start

    assert isinstance(response, str)
    assert len(response) > 100
    print(f"End-to-end query latency: {latency:.2f}s")
    print("Sample response snippet:", response[:200])

if __name__ == "__main__":
    pytest.main([__file__, "-v", "-s", "--tb=no"])