import gradio as gr
import time
import sys
from llama_index.llms.google_genai import GoogleGenAI
from llama_index.core import Settings
from config import *
from document_processor import *
from llama_index.core.chat_engine import CondensePlusContextChatEngine

query_engine = None
chunks_df = None
chat_engine = None
chat_history = []

def answer_question(question, history):
    global query_engine, chat_engine
    
    if query_engine is None:
        return "<div style='background-color: #e53e3e; color: white; padding: 20px; border-radius: 10px;'>❌ System not initialized or document database is empty</div>", "", history
    
    try:
        start_time = time.time()
        
        # Initialize chat engine if not exists
        if chat_engine is None:
            chat_engine = CondensePlusContextChatEngine.from_defaults(
                retriever=query_engine.retriever,
                response_synthesizer=query_engine.response_synthesizer
            )
        
        response = chat_engine.chat(question)
        retrieved_nodes = query_engine.retriever.retrieve(question)
        
        end_time = time.time()
        processing_time = end_time - start_time
        
        sources_html = generate_sources_html(retrieved_nodes)
        
        answer_with_time = f"""<div style='background-color: #2d3748; color: white; padding: 20px; border-radius: 10px; margin-bottom: 10px;'>
        <h3 style='color: #63b3ed; margin-top: 0;'>📋 Answer:</h3>
        <div style='line-height: 1.6; font-size: 16px;'>{response.response}</div>
        <div style='margin-top: 15px; padding-top: 10px; border-top: 1px solid #4a5568; font-size: 14px; color: #a0aec0;'>
        ⏱️ Processing time: {processing_time:.2f} sec
        </div>
        </div>"""
        
        # Update chat history (keep last 6 messages - 3 exchanges)
        new_history = history + [[question, response.response]]
        if len(new_history) > 3:
            new_history = new_history[-3:]
        
        return answer_with_time, sources_html, new_history
        
    except Exception as e:
        error_msg = f"<div style='background-color: #e53e3e; color: white; padding: 20px; border-radius: 10px;'>❌ Error processing question: {str(e)}</div>"
        return error_msg, "", history
    


def generate_sources_html(nodes):
    html = "<div style='background-color: #2d3748; color: white; padding: 20px; border-radius: 10px; max-height: 400px; overflow-y: auto;'>"
    html += "<h3 style='color: #63b3ed; margin-top: 0;'>📚 Sources:</h3>"
    
    unique_docs = {}
    for node in nodes:
        metadata = node.metadata if hasattr(node, 'metadata') else {}
        doc_name = metadata.get('document_name', 'unknown')
        doc_link = metadata.get('document_link', '')
        doc_key = f"{doc_name}||{doc_link}"
        if doc_key not in unique_docs:
            unique_docs[doc_key] = []
        unique_docs[doc_key].append(node)
    
    for doc_key, doc_nodes in unique_docs.items():
        doc_name, doc_link = doc_key.split('||', 1)
        html += f"<div style='margin-bottom: 15px; padding: 15px; border: 1px solid #4a5568; border-radius: 8px; background-color: #1a202c;'>"
        if doc_link:
            html += f"<h4 style='margin: 0 0 10px 0; color: #63b3ed;'>📄 <a href='{doc_link}' target='_blank' style='color: #63b3ed;'>{doc_name}</a></h4>"
        else:
            html += f"<h4 style='margin: 0 0 10px 0; color: #63b3ed;'>📄 {doc_name}</h4>"
        html += f"<p style='margin: 0; color: #a0aec0; font-size: 14px;'>Found relevant fragments: {len(doc_nodes)}</p>"
        html += "</div>"
    
    html += "</div>"
    return html

def get_documents_display():
    documents = get_existing_documents()
    
    if not documents:
        return "<div style='padding: 20px; text-align: center; color: #666;'>No documents in system yet</div>"
    
    html = f"<div style='background-color: #f8f9fa; padding: 20px; border-radius: 10px;'>"
    html += f"<h3 style='color: #2d3748; margin-top: 0;'>📚 {len(documents)} documents in the system:</h3>"
    html += "<div style='max-height: 400px; overflow-y: auto;'>"
    
    for i, doc_name in enumerate(documents, 1):
        html += f"<div style='padding: 8px; margin: 5px 0; background-color: white; border-radius: 5px; border-left: 4px solid #63b3ed;'>"
        html += f"{i}. {doc_name}"
        html += "</div>"
    
    html += "</div></div>"
    return html

def upload_and_process_file(files, doc_names, doc_links):
    global query_engine, chunks_df, chat_engine
    
    if not files:
        return "No files selected", get_documents_display()
    
    if len(files) != len(doc_names) or len(files) != len(doc_links):
        return "Error: Number of files must match number of document names and links", get_documents_display()
    
    existing_docs = get_existing_documents()
    results = []
    
    for i, file in enumerate(files):
        doc_name = doc_names[i].strip() if i < len(doc_names) else ""
        doc_link = doc_links[i].strip() if i < len(doc_links) else ""
        
        if not doc_name:
            doc_name = file.name.split('/')[-1].replace('.txt', '').replace('.pdf', '')
        
        # Check if document already exists
        if doc_name in existing_docs:
            results.append(f"⚠️ {doc_name}: Document already exists in the system")
            continue
        
        log_message(f"🔄 Starting processing of file {i+1}/{len(files)}: {file.name}")
        
        file_info, error = process_uploaded_file(file.name, file.name.split('/')[-1], doc_name, doc_link)
        
        if error:
            results.append(f"❌ {file.name.split('/')[-1]}: {error}")
            continue
        
        query_engine, chunks_df, error = add_to_vector_index(file_info['chunks'], file_info, chunks_df)
        
        if error:
            results.append(f"❌ {file_info['file_name']}: Error adding to database - {error}")
        else:
            results.append(f"✅ {file_info['document']}: Successfully processed and added to database")
            log_message(f"✅ Completed processing: {file_info['document']}")
            # Reset chat engine to include new documents
            chat_engine = None
    
    return "\n".join(results), get_documents_display()

def create_interface():
    with gr.Blocks(title="AIEXP - AI Expert for Regulatory Documentation", theme=gr.themes.Soft()) as demo:
        
        gr.Markdown("""
        # AIEXP - Artificial Intelligence Expert
        
        ## Tool for working with regulatory documentation
        """)
        
        with gr.Tab("🔍 Document Search"):
            gr.Markdown("### Ask a question about the uploaded documentation")
            
            with gr.Row():
                with gr.Column(scale=3):
                    chatbot = gr.Chatbot(
                        label="Chat History",
                        height=400,
                        show_label=True
                    )
                    
                    question_input = gr.Textbox(
                        label="Your question to the knowledge base",
                        placeholder="Enter your question about the documents...",
                        lines=3
                    )
                    ask_btn = gr.Button("🔍 Find Answer", variant="primary", size="lg")
                    clear_btn = gr.Button("🗑️ Clear History", variant="secondary")
                    
                    gr.Examples(
                        examples=[
                            "Какой стандарт устанавливает порядок признания протоколов испытаний продукции в области использования атомной энергии?",
                            "Кто несет ответственность за организацию и проведение признания протоколов испытаний продукции?",
                            "В каких случаях могут быть признаны протоколы испытаний, проведенные лабораториями, не включенными в перечисления?",
                        ],
                        inputs=question_input
                    )
            
                with gr.Column(scale=1):
                    answer_output = gr.HTML(
                        label="",
                        value="<div style='background-color: #2d3748; color: white; padding: 20px; border-radius: 10px; text-align: center;'>The answer to your question will appear here...</div>",
                    )
                    
                    sources_output = gr.HTML(
                        label="",
                        value="<div style='background-color: #2d3748; color: white; padding: 20px; border-radius: 10px; text-align: center;'>Sources will appear here...</div>",
                    )
            
            ask_btn.click(
                fn=answer_question,
                inputs=[question_input, chatbot],
                outputs=[answer_output, sources_output, chatbot]
            ).then(
                lambda: "", inputs=None, outputs=question_input
            )
            
            question_input.submit(
                fn=answer_question,
                inputs=[question_input, chatbot],
                outputs=[answer_output, sources_output, chatbot]
            ).then(
                lambda: "", inputs=None, outputs=question_input
            )
            
            clear_btn.click(
                lambda: [], inputs=None, outputs=chatbot
            )
        
        with gr.Tab("📚 Document Management"):
            gr.Markdown("### Document database and adding new files")
            
            with gr.Row():
                with gr.Column(scale=2):
                    documents_display = gr.HTML(
                        label="Document list",
                        value=get_documents_display()
                    )
                    
                    refresh_btn = gr.Button("🔄 Refresh List", variant="secondary")
                
                with gr.Column(scale=1, elem_id="upload-column"):
                    gr.Markdown("#### Upload new documents", elem_classes=["upload-header"])
                    gr.Markdown("Supported formats: PDF, TXT", elem_classes=["upload-info"])
                    
                    file_upload = gr.File(
                        file_count="multiple",
                        file_types=[".pdf", ".txt"],
                        label="Select files to upload",
                        elem_classes=["upload-file"]
                    )
                    
                    doc_names_input = gr.Textbox(
                        label="Document names (one per line)",
                        placeholder="Enter document names, one per line...",
                        lines=5,
                        elem_classes=["upload-input"]
                    )
                    
                    doc_links_input = gr.Textbox(
                        label="Document links (one per line)",
                        placeholder="Enter document links, one per line...",
                        lines=5,
                        elem_classes=["upload-input"]
                    )
                    
                    upload_btn = gr.Button("📤 Upload and Process", variant="primary", elem_classes=["upload-btn"])
                    
                    upload_status = gr.Textbox(
                        label="Upload status",
                        lines=8,
                        max_lines=10,
                        interactive=False,
                        elem_classes=["upload-status"]
                    )
            
            def process_names_and_links(names_text, links_text):
                names = [name.strip() for name in names_text.split('\n') if name.strip()]
                links = [link.strip() for link in links_text.split('\n') if link.strip()]
                return names, links
            
            upload_btn.click(
                fn=lambda files, names, links: upload_and_process_file(
                    files, 
                    *process_names_and_links(names, links)
                ),
                inputs=[file_upload, doc_names_input, doc_links_input],
                outputs=[upload_status, documents_display]
            )
            
            refresh_btn.click(
                fn=lambda: get_documents_display(),
                outputs=[documents_display]
            )
    
    # Add CSS to fix white background in upload tab
    demo.css = """
    #upload-column {
        background-color: #f8f9fa !important;
        padding: 20px !important;
        border-radius: 10px !important;
        border: 1px solid #e9ecef !important;
    }
    .upload-header h4 {
        color: #2d3748 !important;
        margin-bottom: 10px !important;
    }
    .upload-info {
        color: #666 !important;
        margin-bottom: 15px !important;
    }
    .upload-file, .upload-input, .upload-status {
        background-color: white !important;
        border: 1px solid #ced4da !important;
        border-radius: 5px !important;
    }
    .upload-btn {
        margin-top: 10px !important;
    }
    """
    
    return demo