import os
import json
import streamlit as st
from openai import OpenAI
from PyPDF2 import PdfReader
from pinecone import Pinecone
from dotenv import load_dotenv
import io

load_dotenv()

# Set up OpenAI client
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

# Set up Pinecone
pc = Pinecone(api_key=os.getenv("PINECONE_API_KEY"))
contract_rag_index = pc.Index(os.getenv("PINECONE_INDEX_NAME"), url=os.getenv("PINECONE_INDEX_URL"))
match_index = pc.Index("match")

# Load matches.json
with open('matches.json', 'r') as f:
    matches = json.load(f)

def get_embedding(text):
    response = client.embeddings.create(input=text, model="text-embedding-3-large")
    return response.data[0].embedding

def get_relevant_context(query, top_k=10):
    query_embedding = get_embedding(query)
    search_results = contract_rag_index.query(vector=query_embedding, top_k=top_k, include_metadata=True)
    sorted_results = sorted(search_results['matches'], key=lambda x: x['score'], reverse=True)
    context = "\n".join([result['metadata']['text'] for result in sorted_results])
    return context, sorted_results

def match_letter(letter_content):
    query_embedding = get_embedding(letter_content)
    search_results = match_index.query(vector=query_embedding, top_k=1, include_metadata=True)
    
    if search_results['matches']:
        best_match = search_results['matches'][0]
        return best_match['metadata']['file_name']
    else:
        return None

def validate_letter(letter_content):
    context, results = get_relevant_context(letter_content)
    matched_file = match_letter(letter_content)
    
    validation_prompt = f"""
    Analyze the following letter and compare it with the provided context from the contract database:
    
    Letter:
    {letter_content}
    
    Context from contracts:
    {context}
    
    Tasks:
    1. Identify the specific claims and statements made in the letter.
    2. For each claim or statement, find the corresponding clause in the contract database context.
    3. Determine if each claim or statement is valid according to the actual contract terms.
    4. If there are any discrepancies or invalid claims, explain why they are invalid and reference the relevant contract terms.
    5. If all claims are valid, confirm their validity and provide supporting evidence from the contracts.
    
    Provide your analysis in the following format:
    Validity: [Valid / Invalid]
    Claims and Analysis:
    1. [Claim 1]: [Valid/Invalid]
       - Relevant Contract Clause: [Quote the relevant clause]
       - Analysis: [Explain why the claim is valid or invalid based on the actual contract terms]
    2. [Claim 2]: [Valid/Invalid]
       - Relevant Contract Clause: [Quote the relevant clause]
       - Analysis: [Explain why the claim is valid or invalid based on the actual contract terms]
    ... (continue for all claims)
    
    Overall Analysis: [Provide a summary of the overall validity of the letter based on the contract terms]
    Discrepancies: [List any discrepancies found between the letter and the actual contract terms]
    Recommendation: [Your recommendation based on the analysis]
    """
    
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": "You are a legal assistant specializing in contract validation. Your task is to strictly compare the claims in the letter with the actual contract terms provided in the context. Do not assume any information not explicitly stated in the contract clauses."},
            {"role": "user", "content": validation_prompt}
        ]
    )
    
    return response.choices[0].message.content, results, matched_file

def generate_output_letter(output_template, validation_result):
    prompt = f"""
    Based on the following validation result, modify the output letter template to reflect the analysis:

    Validation Result:
    {validation_result}

    Output Letter Template:
    {output_template}

    Instructions:
    1. Use the output letter template as a base.
    2. Modify the content to reflect the validity of the claims and any discrepancies found.
    3. Include references to specific contract clauses where relevant.
    4. Maintain a professional and formal tone throughout the letter.
    5. Ensure the letter addresses all the points raised in the original letter, as analyzed in the validation result.

    Please provide the complete modified output letter.
    """

    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": "You are a legal assistant tasked with drafting formal response letters based on contract validations."},
            {"role": "user", "content": prompt}
        ]
    )

    return response.choices[0].message.content

def clean_output(generated_letter):
    prompt = f"""
    Please clean and format the following generated letter:

    {generated_letter}

    Instructions:
    1. Ensure proper formatting and layout.
    2. Correct any grammatical or spelling errors.
    3. Improve clarity and conciseness where possible.
    4. Maintain a professional and formal tone.
    5. Ensure all references to contract clauses and legal terms are accurate.
    6. Format the letter with proper paragraphs, headings, and spacing.

    Please provide the cleaned and formatted letter.
    """

    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": "You are a professional editor specializing in legal and business correspondence."},
            {"role": "user", "content": prompt}
        ]
    )

    return response.choices[0].message.content

# Streamlit UI
st.set_page_config(layout="wide", page_title="Contract Validation System")

st.title("Contract Validation System")

# Letter validation
st.header("Validate Letter Against Contracts")
uploaded_letter = st.file_uploader("Upload a letter to validate (PDF)", type="pdf")

if uploaded_letter and st.button("Validate Letter"):
    with st.spinner("Validating letter..."):
        pdf_reader = PdfReader(io.BytesIO(uploaded_letter.read()))
        letter_content = ""
        for page in pdf_reader.pages:
            letter_content += page.extract_text()
        
        validation_result, sources, matched_file = validate_letter(letter_content)
        
        # Create two columns
        col1, col2 = st.columns(2)
        
        with col1:
            st.subheader("Validation Result")
            st.write(validation_result)
            
            if matched_file:
                st.subheader("Matched Input File")
                st.write(matched_file)
                
                # Find the corresponding output file
                output_file = matches.get(matched_file)
                if output_file:
                    st.subheader("Corresponding Output File")
                    st.write(output_file)
                else:
                    st.error("No corresponding output file found in matches.json")
            else:
                st.warning("No matching input file found in the database")
            
            with st.expander("View Sources (Contract Clauses)"):
                for i, source in enumerate(sources, 1):
                    st.markdown(f"**Source {i} - {source['metadata'].get('doc_name', 'N/A')} - Score: {source['score']:.2f}**")
                    st.markdown(f"Chunk Index: {source['metadata'].get('chunk_index', 'N/A')}")
                    st.text(source['metadata']['text'])
                    st.markdown("---")
        
        with col2:
            if matched_file and output_file:
                output_path = os.path.join('docs', 'out', output_file)
                if os.path.exists(output_path):
                    with open(output_path, 'rb') as f:
                        output_pdf = PdfReader(f)
                        output_template = ""
                        for page in output_pdf.pages:
                            output_template += page.extract_text()
                    
                    st.subheader("Generated Output Letter")
                    generated_letter = generate_output_letter(output_template, validation_result)
                    cleaned_letter = clean_output(generated_letter)
                    st.text_area("", cleaned_letter, height=1000)
                else:
                    st.error(f"Output file not found: {output_path}")

# Chat functionality
st.header("Chat with AI")

# Initialize chat history
if "messages" not in st.session_state:
    st.session_state.messages = []

# Display chat messages from history on app rerun
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# React to user input
if prompt := st.chat_input("What is your question?"):
    # Display user message in chat message container
    st.chat_message("user").markdown(prompt)
    # Add user message to chat history
    st.session_state.messages.append({"role": "user", "content": prompt})

    # Get context from the contract-rag database
    context, _ = get_relevant_context(prompt)

    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[
            {"role": "system", "content": "You are a helpful assistant with knowledge of the contract database."},
            {"role": "user", "content": f"Context from the contract database:\n{context}\n\nUser question: {prompt}"}
        ]
    )
    
    # Display assistant response in chat message container
    with st.chat_message("assistant"):
        st.markdown(response.choices[0].message.content)
    # Add assistant response to chat history
    st.session_state.messages.append({"role": "assistant", "content": response.choices[0].message.content})