import gradio as gr
import os
import time
import tempfile
import shutil
from pathlib import Path
from collections import defaultdict
import json

# LlamaIndex imports
from llama_index.core import (
    VectorStoreIndex, 
    SimpleDirectoryReader, 
    Settings,
    StorageContext,
    load_index_from_storage
)
from llama_index.core.node_parser import SentenceSplitter
from llama_index.llms.openrouter import OpenRouter
from llama_index.embeddings.huggingface import HuggingFaceEmbedding

# Global variables
current_index = None
current_query_engine = None
query_stats = defaultdict(list)
temp_doc_dir = None

# Available models
AVAILABLE_MODELS = {
    "GPT-4o": "openai/gpt-4o",
    "GPT-4o Mini": "openai/gpt-4o-mini",
    "Claude 3.5 Sonnet": "anthropic/claude-3.5-sonnet",
    "Claude 3 Haiku": "anthropic/claude-3-haiku",
    "Llama 3.1 70B": "meta-llama/llama-3.1-70b-instruct",
    "Llama 3.1 8B": "meta-llama/llama-3.1-8b-instruct",
    "Mistral Large": "mistralai/mistral-large",
    "Gemini Pro": "google/gemini-pro"
}

def initialize_embeddings():
    """Initialize the embedding model"""
    try:
        Settings.embed_model = HuggingFaceEmbedding(
            model_name="BAAI/bge-small-en-v1.5"
        )
        return "✓ Embedding model initialized successfully"
    except Exception as e:
        return f"❌ Error initializing embeddings: {str(e)}"

def setup_llm(api_key, model_name, temperature=0.1, max_tokens=512):
    """Setup the language model"""
    try:
        if not api_key:
            return None, "❌ API key is required"
        
        model_id = AVAILABLE_MODELS.get(model_name, "openai/gpt-4o")
        
        llm = OpenRouter(
            api_key=api_key,
            max_tokens=max_tokens,
            context_window=4096,
            model=model_id,
            temperature=temperature
        )
        
        # Test the connection
        test_response = llm.complete("Hello")
        return llm, f"✓ {model_name} configured successfully"
    except Exception as e:
        return None, f"❌ Error setting up LLM: {str(e)}"

def process_uploaded_files(files, progress=gr.Progress()):
    """Process uploaded files and create document index"""
    global current_index, temp_doc_dir
    
    if not files:
        return "❌ No files uploaded", "", ""
    
    try:
        # Create temporary directory for documents
        if temp_doc_dir and os.path.exists(temp_doc_dir):
            shutil.rmtree(temp_doc_dir)
        
        temp_doc_dir = tempfile.mkdtemp()
        
        # Copy uploaded files to temp directory
        file_info = []
        progress(0, desc="Copying files...")
        
        for i, file in enumerate(files):
            if file is not None:
                file_path = Path(file.name)
                dest_path = os.path.join(temp_doc_dir, file_path.name)
                shutil.copy2(file.name, dest_path)
                file_info.append(f"• {file_path.name} ({file_path.suffix})")
                progress((i + 1) / len(files) * 0.3, desc=f"Copying {file_path.name}...")
        
        progress(0.3, desc="Loading documents...")
        
        # Load documents
        documents = SimpleDirectoryReader(
            input_dir=temp_doc_dir,
            exclude_hidden=True,
            recursive=True
        ).load_data()
        
        if not documents:
            return "❌ No readable documents found", "", ""
        
        progress(0.5, desc="Creating text chunks...")
        
        # Configure text splitting
        text_splitter = SentenceSplitter(
            chunk_size=512,
            chunk_overlap=50
        )
        
        progress(0.7, desc="Building vector index...")
        
        # Create vector index
        current_index = VectorStoreIndex.from_documents(
            documents,
            transformations=[text_splitter],
            show_progress=False
        )
        
        progress(1.0, desc="Index created successfully!")
        
        # Calculate statistics
        total_chars = sum(len(doc.text) for doc in documents)
        
        status = f"✓ Successfully processed {len(documents)} documents"
        file_list = "\n".join(file_info)
        stats = f"📊 Total content: ~{total_chars:,} characters\n📁 Files processed: {len(files)}"
        
        return status, file_list, stats
        
    except Exception as e:
        return f"❌ Error processing files: {str(e)}", "", ""

def create_query_engine(api_key, model_name, temperature, max_tokens, similarity_k):
    """Create query engine with current settings"""
    global current_query_engine
    
    if not current_index:
        return None, "❌ No document index available. Please upload documents first."
    
    llm, llm_status = setup_llm(api_key, model_name, temperature, max_tokens)
    if not llm:
        return None, llm_status
    
    try:
        current_query_engine = current_index.as_query_engine(
            llm=llm,
            similarity_top_k=similarity_k,
            response_mode="tree_summarize",
            verbose=False
        )
        return current_query_engine, f"✓ Query engine ready with {model_name}"
    except Exception as e:
        return None, f"❌ Error creating query engine: {str(e)}"

def query_documents(question, api_key, model_name, temperature, max_tokens, similarity_k, show_sources):
    """Query the document index"""
    global query_stats
    
    if not question.strip():
        return "Please enter a question.", "", ""
    
    if not current_index:
        return "❌ No documents loaded. Please upload documents first.", "", ""
    
    # Create/update query engine
    query_engine, status = create_query_engine(api_key, model_name, temperature, max_tokens, similarity_k)
    if not query_engine:
        return status, "", ""
    
    try:
        start_time = time.time()
        
        # Query the documents
        response = query_engine.query(question)
        query_time = time.time() - start_time
        
        # Track statistics
        query_stats['response_times'].append(query_time)
        query_stats['questions'].append(question)
        
        # Format response
        answer = str(response)
        
        # Format sources if requested
        sources_text = ""
        if show_sources and hasattr(response, 'source_nodes'):
            sources_list = []
            for i, node in enumerate(response.source_nodes, 1):
                file_name = node.metadata.get('file_name', 'Unknown')
                score = getattr(node, 'score', 0)
                content_preview = node.text[:150] + "..." if len(node.text) > 150 else node.text
                sources_list.append(f"**Source {i}:** {file_name} (relevance: {score:.3f})\n{content_preview}")
            sources_text = "\n\n".join(sources_list)
        
        # Performance info
        perf_info = f"⏱️ Response time: {query_time:.2f}s | Model: {model_name}"
        
        return answer, sources_text, perf_info
        
    except Exception as e:
        return f"❌ Error during query: {str(e)}", "", ""

def get_performance_stats():
    """Get performance statistics"""
    if not query_stats['response_times']:
        return "No queries performed yet."
    
    times = query_stats['response_times']
    stats = f"""📊 **Performance Statistics** (based on {len(times)} queries)
    
• Average response time: {sum(times)/len(times):.2f}s
• Fastest response: {min(times):.2f}s  
• Slowest response: {max(times):.2f}s
• Total queries: {len(times)}
"""
    return stats

def clear_all_data():
    """Clear all data and reset the application"""
    global current_index, current_query_engine, query_stats, temp_doc_dir
    
    current_index = None
    current_query_engine = None
    query_stats = defaultdict(list)
    
    if temp_doc_dir and os.path.exists(temp_doc_dir):
        shutil.rmtree(temp_doc_dir)
        temp_doc_dir = None
    
    return "✓ All data cleared", "", "", "", ""

# Initialize embeddings on startup
embedding_status = initialize_embeddings()

# Create Gradio interface
with gr.Blocks(title="Wikipedia Stub Generator", theme=gr.themes.Soft()) as app:
    gr.Markdown("""
    # 📚 Women in Religion Stub Generator
    
    Upload your source documents and generate a stub article from a fine-tuned large language model. 
    Built with LlamaIndex and powered by multiple LLM providers through OpenRouter.
    """)
    
    with gr.Tab("📤 Upload Documents"):
        gr.Markdown("### Upload Your Source Documents")
        gr.Markdown("Supported formats: PDF, TXT, DOCX, MD, and more")
        
        file_upload = gr.File(
            label="Choose files",
            file_count="multiple",
            file_types=[".pdf", ".txt", ".docx", ".md", ".csv", ".json"]
        )
        
        upload_btn = gr.Button("Process Documents", variant="primary")
        
        with gr.Row():
            with gr.Column():
                upload_status = gr.Textbox(label="Status", interactive=False)
                file_list = gr.Textbox(label="Uploaded Files", lines=5, interactive=False)
            with gr.Column():
                doc_stats = gr.Textbox(label="Document Statistics", lines=5, interactive=False)
    
    with gr.Tab("💬 Generate Stubs"):
        gr.Markdown("### Generate Stub Article")
        
        with gr.Row():
            with gr.Column(scale=2):
                question_input = gr.Textbox(
                    label="Your Question",
                    placeholder="What would you like to know about the documents?",
                    lines=2
                )
                
                with gr.Row():
                    query_btn = gr.Button("Ask Question", variant="primary")
                    clear_btn = gr.Button("Clear All Data", variant="stop")
                
                answer_output = gr.Textbox(
                    label="Answer",
                    lines=10,
                    interactive=False
                )
                
                performance_info = gr.Textbox(
                    label="Performance Info",
                    interactive=False
                )
            
            with gr.Column(scale=1):
                gr.Markdown("### Settings")
                
                api_key_input = gr.Textbox(
                    label="OpenRouter API Key",
                    placeholder="Enter your API key or leave empty to use HF secret",
                    type="password"
                )
                
                model_dropdown = gr.Dropdown(
                    label="Model",
                    choices=list(AVAILABLE_MODELS.keys()),
                    value="GPT-4o Mini",
                    interactive=True
                )
                
                temperature_slider = gr.Slider(
                    label="Temperature",
                    minimum=0.0,
                    maximum=2.0,
                    value=0.1,
                    step=0.1
                )
                
                max_tokens_slider = gr.Slider(
                    label="Max Tokens",
                    minimum=100,
                    maximum=2000,
                    value=512,
                    step=50
                )
                
                similarity_k_slider = gr.Slider(
                    label="Sources to Retrieve",
                    minimum=1,
                    maximum=10,
                    value=5,
                    step=1
                )
                
                show_sources_checkbox = gr.Checkbox(
                    label="Show Sources",
                    value=True
                )
        
        with gr.Accordion("📖 Sources Used", open=False):
            sources_output = gr.Textbox(
                label="Source Documents",
                lines=8,
                interactive=False
            )
    
    with gr.Tab("📊 Performance"):
        gr.Markdown("### Performance Statistics")
        
        stats_btn = gr.Button("Refresh Stats")
        performance_stats = gr.Textbox(
            label="Statistics",
            lines=10,
            interactive=False
        )
    
    with gr.Tab("ℹ️ Help"):
        gr.Markdown("""
        ### How to Use This Application
        
        1. **Upload Documents**: Go to the "Upload Documents" tab and select your files
        2. **Process**: Click "Process Documents" to create the searchable index
        3. **Ask Questions**: Use the "Ask Questions" tab to query your documents
        4. **Adjust Settings**: Modify model parameters for different response styles
        
        ### Best Practices for Questions
        
        - 🎯 **Be specific**: "What does Smith say about feminist theology?" vs "Tell me about feminism"
        - 📚 **Ask about concepts**: "What is religious authority?" rather than just names
        - 🔍 **Use comparative questions**: "How do different scholars approach this topic?"
        - 📊 **Request analysis**: "What are the main arguments presented?"
        - 🏛️ **Ask about methodology**: "What research methods are discussed?"
        
        ### API Key Setup
        
        You can provide your OpenRouter API key in two ways:
        1. Enter it directly in the "API Key" field
        2. Set it as a Hugging Face Space secret named `OPENROUTER_API_KEY`
        
        ### Model Information
        
        Different models have different strengths:
        - **GPT-4o**: Best overall performance, most accurate
        - **Claude 3.5 Sonnet**: Excellent reasoning and analysis
        - **Llama models**: Open source, good performance
        - **Mistral**: Strong multilingual capabilities
        """)
    
    # Event handlers
    upload_btn.click(
        fn=process_uploaded_files,
        inputs=[file_upload],
        outputs=[upload_status, file_list, doc_stats],
        show_progress=True
    )
    
    query_btn.click(
        fn=query_documents,
        inputs=[
            question_input, api_key_input, model_dropdown, 
            temperature_slider, max_tokens_slider, similarity_k_slider, 
            show_sources_checkbox
        ],
        outputs=[answer_output, sources_output, performance_info]
    )
    
    clear_btn.click(
        fn=clear_all_data,
        inputs=[],
        outputs=[upload_status, file_list, doc_stats, answer_output, sources_output]
    )
    
    stats_btn.click(
        fn=get_performance_stats,
        inputs=[],
        outputs=[performance_stats]
    )
    
    # Auto-refresh API key from environment if not provided
    def get_api_key():
        return os.getenv("OPENROUTER_API_KEY", "")
    
    app.load(
        fn=get_api_key,
        inputs=[],
        outputs=[api_key_input]
    )

# Launch the app
if __name__ == "__main__":
    app.launch()