#!/usr/bin/env python3
"""
Tests for MDRetriever module
"""

import unittest
import tempfile
import os
import sys
from pathlib import Path

# Add src to path
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))

from redflagbpm.mcp.md_retriever import MDRetriever


class TestMDRetriever(unittest.TestCase):
    """Test cases for MDRetriever class"""
    
    def setUp(self):
        """Set up test fixtures"""
        # Create a temporary database for each test
        self.temp_db = tempfile.NamedTemporaryFile(delete=False, suffix='.db')
        self.temp_db.close()
        self.db_path = self.temp_db.name
        self.retriever = MDRetriever(db_path=self.db_path)
    
    def tearDown(self):
        """Clean up after tests"""
        if os.path.exists(self.db_path):
            os.unlink(self.db_path)
    
    def test_init_database(self):
        """Test database initialization"""
        # Database should be created
        self.assertTrue(os.path.exists(self.db_path))
        
        # Should be able to query the table
        import sqlite3
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='md_docs'")
        result = cursor.fetchone()
        conn.close()
        self.assertIsNotNone(result)
    
    def test_process_simple_document(self):
        """Test processing a simple markdown document"""
        # Create a temporary markdown file
        with tempfile.NamedTemporaryFile(mode='w', suffix='.md', delete=False, encoding='utf-8') as f:
            f.write("""# First Section [section1]
> Keywords for section 1

Content of section 1.

# Second Section [section2]
> Keywords for section 2

Content of section 2.
""")
            temp_path = f.name
        
        try:
            index = "TEST/guide.md"
            self.retriever.process_document(index, temp_path)
            
            # Check that sections were stored
            doc1 = self.retriever.get_document_by_id(f"{index}#section1")
            self.assertIsNotNone(doc1)
            self.assertEqual(doc1['title'], "First Section")
            self.assertEqual(doc1['keywords'], "Keywords for section 1")
            self.assertIn("Content of section 1", doc1['content'])
            
            doc2 = self.retriever.get_document_by_id(f"{index}#section2")
            self.assertIsNotNone(doc2)
            self.assertEqual(doc2['title'], "Second Section")
        finally:
            os.unlink(temp_path)
    
    def test_process_document_without_keywords(self):
        """Test processing document sections without keywords"""
        with tempfile.NamedTemporaryFile(mode='w', suffix='.md', delete=False, encoding='utf-8') as f:
            f.write("""# Section Without Keywords [section-no-keywords]

Content without keywords.
Some more content here.
""")
            temp_path = f.name
        
        try:
            index = "TEST/guide.md"
            self.retriever.process_document(index, temp_path)
            
            doc = self.retriever.get_document_by_id(f"{index}#section-no-keywords")
            self.assertIsNotNone(doc)
            # Title should be used as keywords when keywords are not provided
            self.assertEqual(doc['keywords'], "Section Without Keywords")
        finally:
            os.unlink(temp_path)
    
    def test_process_document_with_code_blocks(self):
        """Test processing document with code blocks"""
        with tempfile.NamedTemporaryFile(mode='w', suffix='.md', delete=False, encoding='utf-8') as f:
            f.write("""# Code Section [code-section]
> Section with code

Here is some code:

```python
def hello():
    print("Hello")
```

More content after code block.
""")
            temp_path = f.name
        
        try:
            index = "TEST/guide.md"
            self.retriever.process_document(index, temp_path)
            
            doc = self.retriever.get_document_by_id(f"{index}#code-section")
            self.assertIsNotNone(doc)
            # Code block should be included in content
            self.assertIn("```python", doc['content'])
            self.assertIn("def hello():", doc['content'])
            self.assertIn("More content", doc['content'])
        finally:
            os.unlink(temp_path)
    
    def test_get_index_info(self):
        """Test getting index information"""
        with tempfile.NamedTemporaryFile(mode='w', suffix='.md', delete=False, encoding='utf-8') as f:
            f.write("""# Section 1 [section1]
> Keywords 1

Content 1.

# Section 2 [section2]
> Keywords 2

Content 2.
""")
            temp_path = f.name
        
        try:
            agent_id = "TEST/AGENT"
            md_resources = ["guide.md"]
            # Index must match the format: agent_id/md_resource
            index = f"{agent_id}/{md_resources[0]}"
            
            self.retriever.process_document(index, temp_path)
            
            index_info = self.retriever.get_index_info(agent_id, md_resources)
            
            # Should contain section information
            self.assertIn("section1", index_info)
            self.assertIn("section2", index_info)
            self.assertIn("Section 1", index_info)
            self.assertIn("Section 2", index_info)
        finally:
            os.unlink(temp_path)
    
    def test_drop_index(self):
        """Test dropping an index"""
        with tempfile.NamedTemporaryFile(mode='w', suffix='.md', delete=False, encoding='utf-8') as f:
            f.write("""# Section [section1]

Content.
""")
            temp_path = f.name
        
        try:
            index = "TEST/guide.md"
            self.retriever.process_document(index, temp_path)
            
            # Verify document exists
            doc = self.retriever.get_document_by_id(f"{index}#section1")
            self.assertIsNotNone(doc)
            
            # Drop index
            deleted_count = self.retriever.drop_index(index)
            self.assertEqual(deleted_count, 1)
            
            # Verify document is gone
            doc = self.retriever.get_document_by_id(f"{index}#section1")
            self.assertIsNone(doc)
        finally:
            os.unlink(temp_path)
    
    def test_get_document_by_id_not_found(self):
        """Test getting a document that doesn't exist"""
        doc = self.retriever.get_document_by_id("NONEXISTENT#section")
        self.assertIsNone(doc)


if __name__ == '__main__':
    unittest.main()

