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

import re

import re
from difflib import SequenceMatcher


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


def normalize_doc_id(doc_id: str) -> str:
    """Normalize document ID - KEEP dots for numeric parts"""
    doc_id = doc_id.upper().strip()
    doc_id = re.sub(r'\s+', '', doc_id)  # Remove spaces only
    doc_id = doc_id.replace("ГОСТР", "ГОСТ")
    doc_id = doc_id.replace("GOSTR", "ГОСТ")
    return doc_id

def base_number(doc_id: str) -> str:
    """Extract full numeric pattern including all parts (e.g., '59023.6' from 'ГОСТ 59023.6')"""
    # Match: 59023.6 or 59023.4 or 50.05.01 etc.
    m = re.search(r'(\d+(?:\.\d+)*)', doc_id)
    return m.group(1) if m else ""

def filter_nodes_by_doc_id(nodes, doc_ids, threshold=0.85):
    """Filter nodes by document ID with strict numeric matching"""
    if not doc_ids:
        return nodes

    filtered = []
    doc_ids_norm = [normalize_doc_id(d) for d in doc_ids]
    doc_ids_base = [base_number(d) for d in doc_ids_norm]

    for node in nodes:
        node_doc_id = normalize_doc_id(node.metadata.get('document_id', ''))
        node_base = base_number(node_doc_id)

        for q_doc, q_base in zip(doc_ids_norm, doc_ids_base):
            # STRICT: base number must match exactly
            if q_base and node_base and q_base == node_base:
                filtered.append(node)
                break
            
            # STRICT: full normalized ID must match exactly or have very high similarity
            elif SequenceMatcher(None, node_doc_id, q_doc).ratio() >= threshold:
                filtered.append(node)
                break

    return filtered if filtered else nodes


def extract_doc_id_from_query(query):
    """Extract document IDs from query text with better pattern matching"""
    patterns = [
        r'ГОСТ\s*Р?\s*\d+(?:\.\d+)*(?:-\d{4})?',  # ГОСТ 59023.4, ГОСТ Р 50.05.01-2018
        r'НП-\d+(?:-\d+)?',                         # НП-104-18
        r'МУ[_\s]\d+(?:\.\d+)+(?:\.\d+)*(?:-\d{4})?',  # МУ 1.2.3.07.0057-2018
    ]
    
    found_ids = []
    for pattern in patterns:
        matches = re.findall(pattern, query, re.IGNORECASE)
        found_ids.extend(matches)
    
    # Normalize spacing and preserve dots
    normalized = [re.sub(r'\s+', ' ', id.strip().upper()) for id in found_ids]
    return normalized
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=100
    )
    bm25_retriever = BM25Retriever.from_defaults(
        docstore=vector_index.docstore,
        similarity_top_k=100,
        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}")
                before = len(nodes)
                nodes = filter_nodes_by_doc_id(nodes, doc_ids)
                after = len(nodes)
                log_message(f"Filtered by doc ID: {after}/{before} nodes kept (fallback safe)")

            
            # 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