"""
Tests for CogniQ — run with: pytest tests/ -v
"""

import pytest
import numpy as np
import sys, os
sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))

from cogniq import RAG, SearchResult, AskResult
from cogniq.embedders import TFIDFEmbedder, CustomEmbedder, AutoEmbedder
from cogniq.stores import MemoryStore
from cogniq.chunkers import FixedChunker, SentenceChunker, SmartChunker
from cogniq.graph import KnowledgeGraph
from cogniq.plugins import Plugin, register_plugin, PluginChain, DedupPlugin, MinScorePlugin


# ─────────────────────────────────────────────
# EMBEDDER TESTS
# ─────────────────────────────────────────────

class TestTFIDFEmbedder:
    def test_basic_embed(self):
        e = TFIDFEmbedder(dim=128)
        v = e.embed("hello world")
        assert v.shape == (128,)
        assert abs(np.linalg.norm(v) - 1.0) < 1e-5

    def test_zero_text(self):
        e = TFIDFEmbedder(dim=64)
        v = e.embed("")
        assert v.shape == (64,)
        assert np.linalg.norm(v) == 0.0

    def test_batch_embed(self):
        e = TFIDFEmbedder(dim=64)
        texts = ["hello world", "foo bar", "test document"]
        vecs = e.embed_batch(texts)
        assert vecs.shape == (3, 64)

    def test_similar_texts_closer(self):
        e = TFIDFEmbedder(dim=256)
        docs = ["liquidity risk management", "banking liquidity requirements",
                "fruit salad recipe", "apple mango"]
        e.fit(docs)
        v1 = e.embed("liquidity risk management")
        v2 = e.embed("banking liquidity requirements")
        v3 = e.embed("fruit salad recipe")
        sim12 = float(v1 @ v2)
        sim13 = float(v1 @ v3)
        assert sim12 > sim13, "Related docs should have higher similarity"

    def test_partial_fit(self):
        e = TFIDFEmbedder(dim=64)
        e.partial_fit("new document about finance")
        v = e.embed("finance document")
        assert v.shape == (64,)

    def test_deterministic(self):
        e = TFIDFEmbedder(dim=128)
        t = "the quick brown fox"
        assert np.allclose(e.embed(t), e.embed(t))


class TestCustomEmbedder:
    def test_wrap_function(self):
        def my_fn(text):
            return [0.5] * 32
        e = CustomEmbedder(my_fn, dim=32)
        v = e.embed("hello")
        assert v.shape == (32,)
        assert abs(np.linalg.norm(v) - 1.0) < 1e-5


# ─────────────────────────────────────────────
# STORE TESTS
# ─────────────────────────────────────────────

class TestMemoryStore:
    def test_add_and_search(self):
        store = MemoryStore()
        vecs = np.random.randn(3, 16).astype(np.float32)
        docs = [{"text": f"doc {i}"} for i in range(3)]
        ids = store.add(vecs, docs)
        assert len(ids) == 3
        assert len(store) == 3

        q = vecs[0]  # search with first vector — should be top result
        results = store.search(q, top_k=3)
        assert len(results) == 3
        assert results[0][0]["text"] == "doc 0"
        assert results[0][1] > 0.99  # nearly perfect match

    def test_score_range(self):
        store = MemoryStore()
        vecs = np.random.randn(10, 32).astype(np.float32)
        docs = [{"text": f"d{i}"} for i in range(10)]
        store.add(vecs, docs)
        q = np.random.randn(32).astype(np.float32)
        results = store.search(q, top_k=5)
        for _, score in results:
            assert -1.01 <= score <= 1.01

    def test_capacity_growth(self):
        store = MemoryStore(initial_capacity=4)
        for i in range(20):
            vec = np.random.randn(8).astype(np.float32)
            store.add(vec[None, :], [{"text": f"d{i}"}])
        assert len(store) == 20

    def test_delete(self):
        store = MemoryStore()
        vecs = np.random.randn(5, 8).astype(np.float32)
        docs = [{"text": f"d{i}", "id": f"id{i}"} for i in range(5)]
        ids = store.add(vecs, docs)
        store.delete([ids[0]])
        assert len(store) == 4

    def test_empty_search(self):
        store = MemoryStore()
        q = np.random.randn(16).astype(np.float32)
        results = store.search(q)
        assert results == []


# ─────────────────────────────────────────────
# CHUNKER TESTS
# ─────────────────────────────────────────────

class TestChunkers:
    LONG_TEXT = " ".join(["word"] * 600)

    def test_fixed_chunker(self):
        c = FixedChunker(size=100, overlap=20)
        chunks = c.chunk(self.LONG_TEXT)
        assert len(chunks) > 1
        for ch in chunks:
            assert len(ch["text"].split()) <= 100

    def test_sentence_chunker(self):
        text = "First sentence. Second sentence! Third sentence? Fourth one. Fifth one."
        c = SentenceChunker(target_sentences=2)
        chunks = c.chunk(text)
        assert len(chunks) >= 1

    def test_smart_chunker_short(self):
        text = "Short text."
        c = SmartChunker()
        chunks = c.chunk(text, source="test")
        assert len(chunks) >= 1
        assert chunks[0]["source"] == "test"

    def test_smart_chunker_headings(self):
        text = "# Section 1\nContent here.\n\n## Section 2\nMore content here."
        c = SmartChunker(min_words=1, max_words=50)
        chunks = c.chunk(text)
        assert len(chunks) >= 1

    def test_chunk_metadata(self):
        c = FixedChunker(size=50)
        chunks = c.chunk("hello world " * 100, source="my_doc")
        for ch in chunks:
            assert "text" in ch
            assert "chunk_idx" in ch
            assert ch["source"] == "my_doc"


# ─────────────────────────────────────────────
# KNOWLEDGE GRAPH TESTS
# ─────────────────────────────────────────────

class TestKnowledgeGraph:
    def test_ingest_and_entities(self):
        g = KnowledgeGraph()
        g.ingest("Banks must maintain LCR. Failure to maintain LCR causes penalty.")
        stats = g.stats()
        assert stats["entities"] > 0

    def test_causal_extraction(self):
        g = KnowledgeGraph()
        g.ingest("Risk causes liquidity stress. Liquidity stress leads to penalty.")
        # Should have extracted causal relations
        stats = g.stats()
        assert stats["relations"] >= 0  # may or may not match patterns

    def test_score_relevance(self):
        g = KnowledgeGraph()
        g.ingest("LCR is the Liquidity Coverage Ratio used by RBI for banks.")
        s1 = g.score_relevance("LCR requirements for banks", "LCR is used by banks")
        s2 = g.score_relevance("LCR requirements for banks", "fruit salad recipe")
        assert s1 >= s2

    def test_add_custom_relation(self):
        g = KnowledgeGraph()
        g.add_relation("stress", "default", "triggers")
        paths = g.find_paths("stress", "default")
        assert len(paths) >= 0  # path exists


# ─────────────────────────────────────────────
# PLUGIN TESTS
# ─────────────────────────────────────────────

class TestPlugins:
    def test_dedup_plugin(self):
        p = DedupPlugin(threshold=0.5)
        results = [
            ({"text": "hello world foo"}, 0.9),
            ({"text": "hello world bar"}, 0.85),  # near-dup
            ({"text": "completely different text xyz"}, 0.7),
        ]
        filtered = p.post_search(results, "query")
        # Should keep 1 of the two similar ones + the different one
        assert len(filtered) <= 3

    def test_min_score_plugin(self):
        p = MinScorePlugin(threshold=0.5)
        results = [
            ({"text": "a"}, 0.8),
            ({"text": "b"}, 0.3),  # below threshold
            ({"text": "c"}, 0.6),
        ]
        filtered = p.post_search(results, "query")
        assert len(filtered) == 2
        assert all(s >= 0.5 for _, s in filtered)

    def test_plugin_chain_from_names(self):
        chain = PluginChain.build(["dedup", "min_score"])
        assert len(chain.plugins) == 2

    def test_custom_plugin(self):
        @register_plugin("test_plugin_xyz")
        class MyPlugin(Plugin):
            def post_search(self, results, query, **kw):
                return results[:1]  # only top 1

        chain = PluginChain.build(["test_plugin_xyz"])
        results = [({"text": f"r{i}"}, 0.9 - i * 0.1) for i in range(5)]
        out = chain.post_search(results, "q")
        assert len(out) == 1


# ─────────────────────────────────────────────
# END-TO-END RAG TESTS
# ─────────────────────────────────────────────

class TestRAGEndToEnd:
    DOCS = [
        "The Liquidity Coverage Ratio (LCR) requires banks to hold enough HQLA to cover 30-day outflows.",
        "Banks failing to maintain minimum LCR face penalties and supervisory action from RBI.",
        "HQLA includes government securities, central bank reserves, and high-quality corporate bonds.",
        "The Net Stable Funding Ratio (NSFR) ensures long-term funding stability over one year.",
        "Python is a high-level programming language used in data science and AI development.",
    ]

    def setup_method(self):
        self.rag = RAG(embedder="tfidf", graph=True)
        for i, doc in enumerate(self.DOCS):
            self.rag.add(doc, source=f"doc_{i}")

    def test_len(self):
        assert len(self.rag) >= len(self.DOCS)  # may be more (chunked)

    def test_search_returns_results(self):
        results = self.rag.search("LCR requirements")
        assert len(results) > 0
        assert all(isinstance(r, SearchResult) for r in results)

    def test_search_relevance(self):
        results = self.rag.search("LCR penalty")
        # Should return LCR-related docs, not Python doc
        top_text = results[0].text.lower()
        assert any(kw in top_text for kw in ["lcr", "penalty", "bank", "liquidity"])

    def test_ask_no_llm(self):
        result = self.rag.ask("What is LCR?")
        assert isinstance(result, AskResult)
        assert len(result.answer) > 0
        assert isinstance(result.latency_ms, float)
        assert len(result.sources) > 0

    def test_ask_with_custom_llm(self):
        def mock_llm(question, context, **kw):
            return f"MOCK: {question}"

        result = self.rag.ask("What is LCR?", llm=mock_llm)
        assert result.answer.startswith("MOCK:")

    def test_operator_add(self):
        r = RAG(embedder="tfidf")
        r += "new document about banking"
        assert len(r) > 0

    def test_operator_call(self):
        result = self.rag("What happens if LCR is not maintained?")
        assert isinstance(result, AskResult)

    def test_operator_getitem(self):
        results = self.rag["LCR"]
        assert isinstance(results, list)

    def test_save_load(self, tmp_path):
        path = tmp_path / "rag.pkl"
        self.rag.save(path)
        loaded = RAG.load(path)
        assert len(loaded) == len(self.rag)
        results = loaded.search("LCR")
        assert len(results) > 0

    def test_info(self):
        info = self.rag.info()
        assert "documents" in info
        assert info["documents"] > 0

    def test_repr(self):
        r = repr(self.rag)
        assert "RAG" in r

    def test_plugins_in_rag(self):
        r = RAG(embedder="tfidf", plugins=["dedup"])
        for doc in self.DOCS:
            r += doc
        results = r.search("LCR")
        assert len(results) >= 1

    def test_global_api(self):
        import cogniq
        cogniq.reset()
        cogniq.add("Test document for global API.", source="test")
        results = cogniq.search("document")
        assert len(results) > 0
        cogniq.reset()


# ─────────────────────────────────────────────
# PERFORMANCE BENCHMARK
# ─────────────────────────────────────────────

class TestPerformance:
    def test_bulk_add_speed(self):
        import time
        r = RAG(embedder="tfidf", graph=False)
        docs = [f"Document {i} about banking and finance and liquidity risk management." for i in range(100)]
        t0 = time.perf_counter()
        for doc in docs:
            r += doc
        elapsed = time.perf_counter() - t0
        assert elapsed < 30.0, f"100 docs took {elapsed:.1f}s — too slow"
        print(f"\n✓ 100 docs added in {elapsed:.2f}s ({100/elapsed:.0f} docs/sec)")

    def test_search_speed(self):
        import time
        r = RAG(embedder="tfidf", graph=False)
        for i in range(200):
            r += f"Document {i} with financial keywords about risk and liquidity."
        t0 = time.perf_counter()
        for _ in range(50):
            r.search("liquidity risk", top_k=5)
        elapsed = time.perf_counter() - t0
        qps = 50 / elapsed
        print(f"\n✓ 50 queries over 200 docs: {elapsed*20:.2f}ms/query ({qps:.0f} QPS)")
        assert elapsed < 10.0, f"50 queries took {elapsed:.2f}s"


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