from llama_index.core import VectorStoreIndex
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.retrievers.bm25 import BM25Retriever
from llama_index.core.retrievers import QueryFusionRetriever
from llama_index.core.response_synthesizers import get_response_synthesizer
from my_logging import log_message

def create_vector_index(documents):
    """Create vector index from documents"""
    log_message(f"Building vector index from {len(documents)} documents...")
    index = VectorStoreIndex.from_documents(documents)
    log_message("✓ Index created")
    return index

def keyword_filter_nodes(query, nodes, min_keyword_matches=1):
    """Return nodes that contain at least one keyword from the query."""
    keywords = [w.lower() for w in query.split() if len(w) > 2]
    filtered = []
    for node in nodes:
        text = node.text.lower()
        if any(k in text for k in keywords):
            filtered.append(node)
    return filtered

import re

def extract_doc_id_from_query(query):
    """Extract document IDs from query text"""
    # Match patterns like: ГОСТ 59023.2, НП-104, ГОСТ Р 50.04.07-2022
    patterns = [
        r'(?:ГОСТ\s*Р?\s*)[\d\.]+(?:-\d{4})?',  # ГОСТ patterns
        r'НП-\d+(?:-\d+)?',                       # НП patterns
        r'МУ[_\s][\d\.]+',                        # МУ patterns
    ]
    
    found_ids = []
    for pattern in patterns:
        matches = re.findall(pattern, query, re.IGNORECASE)
        found_ids.extend(matches)
    
    # Normalize spacing
    normalized = [re.sub(r'\s+', ' ', id.strip()) for id in found_ids]
    return normalized


def filter_nodes_by_doc_id(nodes, doc_ids, threshold=0.8):
    """Keep nodes that match any of the document IDs"""
    if not doc_ids:
        return nodes
    
    from difflib import SequenceMatcher
    
    filtered = []
    for node in nodes:
        node_doc_id = node.metadata.get('document_id', '').upper()
        
        for query_doc_id in doc_ids:
            query_doc_id = query_doc_id.upper()
            
            # Exact substring match
            if query_doc_id in node_doc_id or node_doc_id in query_doc_id:
                filtered.append(node)
                break
            
            # Fuzzy match for close variants
            similarity = SequenceMatcher(None, query_doc_id, node_doc_id).ratio()
            if similarity >= threshold:
                filtered.append(node)
                break
    
    return filtered

def russian_tokenizer(text):
    """Better tokenizer for Russian document IDs and technical terms"""
    import re
    
    # Keep document ID patterns intact
    text = re.sub(r'(ГОСТ\s*Р?\s*[\d\.]+(?:-\d{4})?)', r' \1 ', text)
    text = re.sub(r'(НП-\d+(?:-\d+)?)', r' \1 ', text)
    text = re.sub(r'(МУ[_\s][\d\.]+)', r' \1 ', text)
    
    # Split on whitespace and punctuation, but keep numbers with decimals
    tokens = re.findall(r'\d+\.\d+|\w+', text.lower())
    
    return tokens


def create_query_engine(vector_index):
    """Create hybrid retrieval engine with document ID filtering"""
    log_message("Creating query engine...")

    vector_retriever = VectorIndexRetriever(
        index=vector_index,
        similarity_top_k=50
    )
    bm25_retriever = BM25Retriever.from_defaults(
        docstore=vector_index.docstore,
        similarity_top_k=50,
        tokenizer=russian_tokenizer  # Add custom tokenizer

    )
    hybrid_retriever = QueryFusionRetriever(
        [vector_retriever, bm25_retriever],
        similarity_top_k=60,
        num_queries=1
    )

    class DeduplicatedQueryEngine(RetrieverQueryEngine):
        def retrieve(self, query):
            nodes = hybrid_retriever.retrieve(query)
            log_message(f"Hybrid retrieval returned: {len(nodes)} nodes")
            
            # Extract document IDs from query
            doc_ids = extract_doc_id_from_query(query)
            if doc_ids:
                log_message(f"Detected document IDs in query: {doc_ids}")
                
                # Filter by document ID
                doc_filtered = filter_nodes_by_doc_id(nodes, doc_ids, threshold=0.7)
                log_message(f"After doc ID filter: {len(doc_filtered)} nodes")
                
                # If we found matching documents, use only those
                if doc_filtered:
                    nodes = doc_filtered
                else:
                    log_message("WARNING: No nodes matched document IDs, using all results")
            
            # Deduplication
            seen_hashes = set()
            unique_nodes = []
            doc_type_counts = {'text': 0, 'table': 0, 'image': 0}
            
            for node in nodes:
                text_hash = hash(node.text[:500])
                
                if text_hash not in seen_hashes:
                    seen_hashes.add(text_hash)
                    unique_nodes.append(node)
                    
                    node_type = node.metadata.get('type', 'text')
                    doc_type_counts[node_type] = doc_type_counts.get(node_type, 0) + 1
            
            log_message(f"After dedup: {len(unique_nodes)} unique nodes")
            log_message(f"Types: text={doc_type_counts.get('text', 0)}, "
                       f"table={doc_type_counts.get('table', 0)}, "
                       f"image={doc_type_counts.get('image', 0)}")
            
            # Log which documents we're returning
            returned_docs = set(n.metadata.get('document_id', 'unknown') for n in unique_nodes[:50])
            log_message(f"Returning nodes from: {sorted(returned_docs)}")
            
            return unique_nodes[:50]

    response_synthesizer = get_response_synthesizer()
    
    query_engine = DeduplicatedQueryEngine(
        retriever=hybrid_retriever,
        response_synthesizer=response_synthesizer
    )

    log_message("✓ Query engine created with doc ID filtering")
    return query_engine