import re
from typing import List

class MockDocument:
    def __init__(self, content, source, folder_name=None):
        self.page_content = content
        self.metadata = {"source": source}
        if folder_name:
            self.metadata["folder_name"] = folder_name

def test_filtering(active_instructions, docs):
    # Regex logic from rag_service.py
    allowed_folders = []
    patterns_only = [
        r"(?:from|in|use|of)\s+([a-z0-9_-]+)\s+folder",
        r"([a-z0-9_-]+)\s+folder\s+only",
        r"([a-z0-9_-]+)\s+file\s+only",
        r"only\s+([a-z0-9_-]+)"
    ]
    for pat in patterns_only:
        matches = re.findall(pat, active_instructions.lower())
        if matches:
            allowed_folders.extend([m.strip() for m in matches if m.strip() not in ["the", "this", "my"]])
    
    ignored_folders = []
    patterns_ignore = [
        r"ignore\s+([a-z0-9_-]+)\s+folder",
        r"ignore\s+([a-z0-9_-]+)\s+file",
        r"don't\s+use\s+([a-z0-9_-]+)",
        r"do\s+not\s+use\s+([a-z0-9_-]+)"
    ]
    for pat in patterns_ignore:
        matches = re.findall(pat, active_instructions.lower())
        if matches:
            ignored_folders.extend([m.strip() for m in matches if m.strip() not in ["the", "this", "my"]])

    print(f"Prompt: {active_instructions}")
    print(f"Allowed Folders: {allowed_folders}")
    print(f"Ignored Folders: {ignored_folders}")

    filtered_docs = []
    for d in docs:
        source_lower = d.metadata.get("source", "").lower()
        doc_folder = d.metadata.get("folder_name", "").lower()
        
        is_allowed = True
        
        if allowed_folders:
            is_allowed = False
            for folder in allowed_folders:
                if folder == doc_folder or folder in source_lower:
                    is_allowed = True
                    break
        
        if is_allowed and ignored_folders:
            for folder in ignored_folders:
                if folder == doc_folder or folder in source_lower:
                    is_allowed = False
                    break
        
        if is_allowed:
            filtered_docs.append(d)
        else:
            print(f">> Filtered out: {source_lower} (Folder: {doc_folder})")

    return filtered_docs

# Test Data
docs = [
    MockDocument("Data from NGROK", "ngrok_info.pdf", "NGROK"),
    MockDocument("Legal Data", "legal_doc.pdf", "Legal"),
    MockDocument("Some other data", "random.pdf", "Main")
]

print("--- TEST 1: Only NGROK ---")
instructions1 = "Hey give answer from NGROK folder only."
results1 = test_filtering(instructions1, docs)
assert len(results1) == 1
assert results1[0].metadata["folder_name"] == "NGROK"
print("SUCCESS: Only NGROK kept.")

print("\n--- TEST 2: Ignore Legal ---")
instructions2 = "Please ignore Legal folder and answer normally."
results2 = test_filtering(instructions2, docs)
assert len(results2) == 2
assert all(d.metadata["folder_name"] != "Legal" for d in results2)
print("SUCCESS: Legal ignored.")

print("\n--- TEST 3: Multiple folders ---")
instructions3 = "Use NGROK folder only and also check items in Main folder only."
results3 = test_filtering(instructions3, docs)
# Logic check: if two "only" folders are specified, it should allow both
assert len(results3) == 2
print("SUCCESS: Multiple folders handled.")

print("\n--- TEST 4: Case insensitive and underscore ---")
instructions4 = "from ngrok folder only"
results4 = test_filtering(instructions4, docs)
assert len(results4) == 1
print("SUCCESS: Case insensitive works.")
