from llama_index.core import VectorStoreIndex
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.postprocessor import SimilarityPostprocessor
from utils.llm_client import llm_client

def compare_proposal_to_specs(specs_index: VectorStoreIndex, proposals_index: VectorStoreIndex, st):
    print(f"- Starting comparison between specs and proposals")

    discrepancies = []

    # Create retrievers with similarity threshold
    specs_retriever = VectorIndexRetriever(
        index=specs_index,
        similarity_top_k=3,
    )

    # Create query engines with post-processing
    specs_query_engine = RetrieverQueryEngine(
        retriever=specs_retriever,
        node_postprocessors=[SimilarityPostprocessor(similarity_cutoff=0.2)]
    )

    print(f"There is {len(proposals_index.docstore.docs)} proposal")
    text_data = f"How to read ?\n You have 3 parts per proposal:\n 1. System prompt\n2. User message\n3. Discrepancy result\n---\n"
    
    # Iterate over the proposal nodes
    index = 0
    for proposal_node_id, proposal_node in proposals_index.docstore.docs.items():
        proposal_text = proposal_node.get_content()
        index += 1
        print(f"\nProcessing proposal {proposal_node_id} (len: {len(proposal_text)}): {proposal_text}\n")
        print("-------------")
        # Retrieve relevant specs sections
        query_result = specs_query_engine.query(proposal_text)

        print(f"\nSpec query result {proposal_node_id}: {query_result}\n")
        
        if not query_result.source_nodes:
            print("No source_nodes")
            # text_data += "No source_nodes"

            continue  # Skip if no relevant specs found

        spec_text = "\n".join([node.get_content() for node in query_result.source_nodes])

        # Use the language model to compare the proposal section to the specs
        system_prompt = "You are an expert in construction contract compliance."
        user_messages = (
            f"Compare the following proposal section with the related specifications and identify any discrepancies.\n\n"
            f"**Proposal Section:**\n{proposal_text}\n\n"
            f"**Related Specifications:**\n{spec_text}\n\n"
            "List any discrepancies you find."
        )
        text_data += f"### Proposal {index}\n"
        text_data += "1. System Prompt\n"
        text_data += f"```\n{system_prompt}\n```\n"
        text_data += f"2. User message\n"
        text_data += f"```\n{user_messages}\n```\n"
        
        
        discrepancies_result = llm_client(system_prompt, user_messages)
        discrepancies.append({
            'specification': spec_text,
            'proposal': proposal_text,
            'discrepancy': discrepancies_result,
            'spec_node_ids': [node.node_id for node in query_result.source_nodes],
            'proposal_node_id': proposal_node_id
        })
        
        st.session_state.state.discrepancies = discrepancies
        text_data += f"\n---\n```\n{discrepancies_result}\n```\n"

        
    return text_data