import os
import boto3
import gradio as gr
import pandas as pd
import torch
import importlib
import shutil
import logging
import fitz  # PyMuPDF for image extraction
import base64
from io import BytesIO
from PIL import Image
from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_aws import ChatBedrock  # Use Bedrock for Claude
from langchain_mistralai.chat_models import ChatMistralAI
from langchain_community.vectorstores import FAISS

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# Environment variables will be loaded from Hugging Face Spaces secrets
MISTRAL_API_KEY = os.environ.get("MISTRAL_API_KEY")
AWS_ACCESS_KEY = os.environ.get("AWS_ACCESS_KEY")
AWS_SECRET_KEY = os.environ.get("AWS_SECRET_KEY")
AWS_REGION = os.environ.get("AWS_REGION", "us-east-1")

# Global variables
use_proprietary = True  # Default to Claude
pdfs_loaded = False
vector_store_loaded = False
chat_history = []
rag_pipeline = None
retriever = None
pdf_image_cache = {}  # Cache for extracted images

# Configure AWS credentials for Bedrock
os.environ["AWS_ACCESS_KEY_ID"] = AWS_ACCESS_KEY
os.environ["AWS_SECRET_ACCESS_KEY"] = AWS_SECRET_KEY
os.environ["AWS_DEFAULT_REGION"] = AWS_REGION

# Function to extract images from PDFs
def extract_images_from_pdf(pdf_path):
    """Extract images from a PDF file and return them as base64 encoded strings."""
    if pdf_path in pdf_image_cache:
        return pdf_image_cache[pdf_path]
    
    logger.info(f"Extracting images from {pdf_path}")
    images = []
    
    try:
        # Open the PDF
        doc = fitz.open(pdf_path)
        
        # For each page
        for page_num, page in enumerate(doc):
            # Get images
            image_list = page.get_images(full=True)
            
            for img_index, img in enumerate(image_list):
                # Get the XREF of the image
                xref = img[0]
                
                # Extract the image bytes
                base_image = doc.extract_image(xref)
                image_bytes = base_image["image"]
                
                # Get the image extension
                image_ext = base_image["ext"]
                
                # Convert to PIL Image
                image = Image.open(BytesIO(image_bytes))
                
                # Convert to base64 for HTML display
                buffered = BytesIO()
                image.save(buffered, format="PNG")
                img_str = base64.b64encode(buffered.getvalue()).decode()
                
                # Store image info
                images.append({
                    "base64": img_str,
                    "page": page_num + 1,
                    "index": img_index
                })
        
        # Cache the results
        pdf_image_cache[pdf_path] = images
        
        return images
    
    except Exception as e:
        logger.error(f"Error extracting images from {pdf_path}: {str(e)}")
        return []

# Function to load PDFs from local directory
def load_pdfs_from_directory():
    """Load PDFs from multiple possible locations in the Hugging Face Space."""
    logger.info("Loading PDFs from file system...")
    
    # List of directories to check for PDFs
    directories_to_check = [
        "pdf_data",           # Default directory
        ".",                  # Root directory
        "/content",           # Another common location
        "/app",               # HF Spaces app directory
        os.path.expanduser("~")  # Home directory
    ]
    
    pdf_files = []
    pdf_locations = {}
    
    # Search for PDFs in each directory
    for directory in directories_to_check:
        if os.path.exists(directory) and os.path.isdir(directory):
            logger.info(f"Checking directory: {directory}")
            try:
                # Check for PDFs in this directory
                for f in os.listdir(directory):
                    if f.lower().endswith('.pdf'):
                        full_path = os.path.join(directory, f)
                        if os.path.isfile(full_path):
                            pdf_files.append(f)
                            pdf_locations[f] = full_path
                            logger.info(f"Found PDF: {f} at {full_path}")
            except Exception as e:
                logger.warning(f"Error checking directory {directory}: {str(e)}")
    
    if not pdf_files:
        # Try a more aggressive search with glob
        import glob
        logger.info("Performing deep search for PDFs...")
        for directory in directories_to_check:
            if os.path.exists(directory):
                # Recursively search for PDFs
                try:
                    for pdf_path in glob.glob(os.path.join(directory, "**/*.pdf"), recursive=True):
                        if os.path.isfile(pdf_path):
                            f = os.path.basename(pdf_path)
                            pdf_files.append(f)
                            pdf_locations[f] = pdf_path
                            logger.info(f"Deep search found PDF: {f} at {pdf_path}")
                except Exception as e:
                    logger.warning(f"Error in deep search for {directory}: {str(e)}")
    
    # If we found PDFs, ensure they're in the pdf_data directory
    if pdf_files:
        # Create pdf_data directory if it doesn't exist
        os.makedirs("pdf_data", exist_ok=True)
        
        # Copy all found PDFs to pdf_data if they're not already there
        for pdf_file in pdf_files:
            source_path = pdf_locations[pdf_file]
            target_path = os.path.join("pdf_data", pdf_file)
            
            # Skip if already in pdf_data
            if os.path.normpath(source_path) == os.path.normpath(target_path):
                continue
                
            try:
                shutil.copy2(source_path, target_path)
                logger.info(f"Copied PDF to pdf_data: {pdf_file}")
            except Exception as e:
                logger.warning(f"Failed to copy {pdf_file}: {str(e)}")
    
    # Final check - what's actually in pdf_data now?
    if os.path.exists("pdf_data"):
        pdf_data_files = [f for f in os.listdir("pdf_data") if f.lower().endswith('.pdf')]
        if pdf_data_files:
            logger.info(f"PDF data directory now contains {len(pdf_data_files)} PDFs: {pdf_data_files}")
            global pdfs_loaded
            pdfs_loaded = True
            return True, f"Successfully loaded {len(pdf_data_files)} PDFs"
    
    # If we still don't have PDFs, log specific PDFs we're looking for
    expected_pdfs = [
        "ACS580_Catalog_3AUA0000145061_RevP_EN.pdf",
        "ACS580MV_catalog_3BHT490775R0001_RevF_EN.pdf",
        "ACS5000_catalog_3BHT490501R0001_RevN_EN.pdf",
        "ACS6080_catalog_3AUA0000221913_RevC_EN.pdf"
    ]
    
    logger.warning(f"Specifically looking for these PDFs: {expected_pdfs}")
    logger.warning("No PDF files found in any expected directory")
    return False, "No PDF files found. Please ensure PDFs are uploaded to the Hugging Face Space."

# Function to process PDFs and create vector store
def process_pdfs_and_create_vectorstore():
    """Process local PDFs and create a FAISS vector store."""
    logger.info("Starting processing of PDFs and creating vector store...")
    
    # Check if PDFs are loaded
    if not pdfs_loaded:
        success, message = load_pdfs_from_directory()
        if not success:
            return False, message
    
    # Create directories
    os.makedirs("processed_data", exist_ok=True)
    
    # Get all PDF files in the pdf_data directory
    pdf_files = [f for f in os.listdir("pdf_data") if f.endswith('.pdf')]
    
    if not pdf_files:
        logger.warning("No PDF files found. Please upload PDFs to the pdf_data directory.")
        return False, "No PDF files found. Please upload PDFs to the pdf_data directory."
    
    # Initialize text splitter with improved parameters for technical content
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200,
        separators=["\n\n", "\n", ". ", " ", ""]
    )
    
    # Load and process each PDF
    all_chunks = []
    
    for i, pdf_file in enumerate(pdf_files):
        pdf_path = os.path.join("pdf_data", pdf_file)
        logger.info(f"Processing PDF {i+1}/{len(pdf_files)}: {pdf_file}")
        
        try:
            loader = PyPDFLoader(pdf_path)
            documents = loader.load()
            
            # Enhance metadata
            for doc in documents:
                doc.metadata["source"] = pdf_file
                doc.metadata["page"] = doc.metadata.get("page", 0) + 1  # Make page numbers 1-indexed
                doc.metadata["total_pages"] = len(documents)
                doc.metadata["title"] = pdf_file.replace(".pdf", "").replace("_", " ").title()
                doc.metadata["pdf_path"] = pdf_path
            
            # Split into chunks
            chunks = text_splitter.split_documents(documents)
            all_chunks.extend(chunks)
            
            # Extract images
            extract_images_from_pdf(pdf_path)
            
        except Exception as e:
            logger.error(f"Error processing {pdf_file}: {str(e)}")
    
    if not all_chunks:
        logger.warning("No content was extracted from the PDFs")
        return False, "No content was extracted from the PDFs"
    
    logger.info(f"Extracted {len(all_chunks)} text chunks from {len(pdf_files)} PDFs")
    logger.info("Generating embeddings for semantic search...")
    
    # Use a Sentence Transformer model for embeddings
    embeddings = HuggingFaceEmbeddings(
        model_name="sentence-transformers/all-MiniLM-L6-v2",
        model_kwargs={'device': 'cuda' if torch.cuda.is_available() else 'cpu'}
    )
    
    logger.info("Building vector database for semantic search...")
    
    # Create FAISS vector store
    vectorstore = FAISS.from_documents(all_chunks, embeddings)
    
    # Save the vector store
    vectorstore.save_local("processed_data/faiss_index")
    
    logger.info("Vector database created and saved successfully")
    
    global vector_store_loaded
    vector_store_loaded = True
    
    return True, vectorstore

# Function to load existing vector store
def load_vectorstore():
    """Load an existing FAISS vector store or create if not exists."""
    logger.info("Attempting to load existing vector store...")
    
    if not os.path.exists("processed_data/faiss_index"):
        logger.info("No existing vector database found. Creating new one...")
        return process_pdfs_and_create_vectorstore()
    
    try:
        # Initialize embeddings
        embeddings = HuggingFaceEmbeddings(
            model_name="sentence-transformers/all-MiniLM-L6-v2",
            model_kwargs={'device': 'cuda' if torch.cuda.is_available() else 'cpu'}
        )
        
        # Load the vector store
        vectorstore = FAISS.load_local("processed_data/faiss_index", embeddings)
        
        global vector_store_loaded
        vector_store_loaded = True
        
        logger.info("Vector database loaded successfully")
        return True, vectorstore
    
    except Exception as e:
        logger.error(f"Error loading vector database: {str(e)}")
        logger.info("Attempting to create new vector store...")
        return process_pdfs_and_create_vectorstore()

# Function to initialize the RAG pipeline
def initialize_rag_pipeline(vectorstore):
    """Initialize the RAG pipeline with either AWS Bedrock Claude or Mistral AI."""
    logger.info(f"Initializing RAG pipeline with {'AWS Bedrock Claude' if use_proprietary else 'Mistral AI'}")
    
    retriever = vectorstore.as_retriever(
        search_type="mmr",  # Use Maximum Marginal Relevance for diverse results
        search_kwargs={"k": 5, "fetch_k": 10}
    )
    
    if use_proprietary:
        # Initialize Claude from AWS Bedrock
        llm = ChatBedrock(
            model_id="anthropic.claude-3-sonnet-20240229-v1:0",
            model_kwargs={
                "temperature": 0.3,
                "max_tokens": 1024
            },
            region_name=AWS_REGION
        )
    else:
        # Initialize Mistral AI model
        llm = ChatMistralAI(
            model="mistral-large-latest",
            temperature=0.3,
            mistral_api_key=MISTRAL_API_KEY
        )
    
    # Create a template for the RAG prompt
    template = """
    You are Ginnie, an expert AI assistant specializing in ABB industrial products and solutions.
    
    <context>
    {context}
    </context>
    
    Human: {question}
    
    Assistant: 
    """
    
    # Create the prompt
    prompt = PromptTemplate.from_template(template)
    
    # Create the chain
    rag_chain = (
        {"context": retriever, "question": RunnablePassthrough()}
        | prompt
        | llm
        | StrOutputParser()
    )
    
    return rag_chain, retriever

# Function to get source documents from retriever
def get_source_documents(retriever, query):
    """Get source documents for a query."""
    docs = retriever.get_relevant_documents(query)
    sources = []
    
    for i, doc in enumerate(docs):
        source_info = {
            "title": doc.metadata.get("title", "Unknown"),
            "source": doc.metadata.get("source", "Unknown"),
            "page": doc.metadata.get("page", "Unknown"),
            "pdf_path": doc.metadata.get("pdf_path", ""),
            "excerpt": doc.page_content[:200] + "..." if len(doc.page_content) > 200 else doc.page_content
        }
        sources.append(source_info)
    
    return sources

# Function to format source citations and include relevant images
def format_sources_with_images(sources, include_images=True):
    """Format sources for display with optional images."""
    if not sources:
        return ""
    
    source_text = "\n\n**Sources:**\n"
    
    # Create a set to track unique sources
    unique_sources = set()
    images_html = ""
    
    for source in sources:
        source_key = f"{source['source']}_{source['page']}"
        if source_key not in unique_sources:
            unique_sources.add(source_key)
            source_text += f"- **{source['title']}** (Page {source['page']})\n"
            
            # Add images if requested and available
            if include_images and source.get("pdf_path") and os.path.exists(source["pdf_path"]):
                # Find images for this page
                page_images = [img for img in extract_images_from_pdf(source["pdf_path"]) 
                             if img["page"] == source["page"]]
                
                # Add up to 2 images per page to avoid clutter
                for i, img in enumerate(page_images[:2]):
                    images_html += f'<div class="source-image"><img src="data:image/png;base64,{img["base64"]}" alt="Image from {source["title"]} page {source["page"]}" /><p>Source: {source["title"]} (Page {source["page"]})</p></div>'
    
    # Add images section if any images were found
    if images_html:
        source_text += "\n\n**Relevant Visuals:**\n"
        source_text += f"<div class='image-container'>{images_html}</div>"
    
    return source_text

# System setup function
def setup_system():
    """Perform complete system setup with improved error handling."""
    global rag_pipeline, retriever
    
    logger.info("Starting system setup...")
    
    # Step 1: Load PDFs if needed
    if not pdfs_loaded:
        success, message = load_pdfs_from_directory()
        if not success:
            logger.warning(f"PDF loading failed: {message}")
            # List files in current directory for debugging
            try:
                logger.info(f"Files in current directory: {os.listdir('.')}")
                if os.path.exists("pdf_data"):
                    logger.info(f"Files in pdf_data directory: {os.listdir('pdf_data')}")
            except Exception as e:
                logger.error(f"Error listing directories: {str(e)}")
    
    # Step 2: Load or create vector store
    success, result = load_vectorstore()
    
    if success and isinstance(result, FAISS):
        # Step 3: Initialize RAG pipeline
        rag_pipeline, retriever = initialize_rag_pipeline(result)
        logger.info("RAG pipeline initialized successfully")
        return True
    else:
        logger.error("Failed to set up the system")
        # Print some system information for debugging
        logger.info(f"Current working directory: {os.getcwd()}")
        logger.info(f"Environment variables: PDF_PATH={os.environ.get('PDF_PATH')}")
        return False

# Message processing function
def process_message(message, chatbot_history):
    """Process user message and generate response with optional images."""
    global chat_history, rag_pipeline, retriever
    
    if not message:
        return chatbot_history
    
    # Add user message to history
    chatbot_history.append((message, ""))
    
    # Check if system is ready
    if not vector_store_loaded or rag_pipeline is None or retriever is None:
        # Try to setup the system
        if setup_system():
            response = "I've just finished setting up the ABB product information system. I can now answer your question."
        else:
            response = "I'm having trouble setting up the system. Please check the logs for more information."
            chatbot_history[-1] = (message, response)
            return chatbot_history
    
    try:
        # Get sources
        sources = get_source_documents(retriever, message)
        
        # Check if the query is about images
        image_request = any(term in message.lower() for term in ["image", "picture", "photo", "visual", "diagram", "figure", "show me"])
        
        # Generate response
        response = rag_pipeline.invoke(message)
        
        # Format response with sources and images if requested
        formatted_response = response + format_sources_with_images(sources, include_images=image_request)
        
        # Update chatbot history
        chatbot_history[-1] = (message, formatted_response)
        
    except Exception as e:
        # Handle errors
        error_message = f"I encountered an error: {str(e)}. Please try again."
        chatbot_history[-1] = (message, error_message)
    
    return chatbot_history

# Function to switch between models
def switch_model(choice):
    """Switch between proprietary and open source models."""
    global use_proprietary, rag_pipeline, retriever
    
    use_proprietary = choice == "Proprietary (Claude AI via AWS Bedrock)"
    logger.info(f"Model switched to {choice}")
    
    # Reinitialize the pipeline if vector store is loaded
    if vector_store_loaded:
        success, vectorstore = load_vectorstore()
        if success:
            rag_pipeline, retriever = initialize_rag_pipeline(vectorstore)
    
    return f"Model switched to {choice}"

# Function to reset chat
def reset_chat(chatbot_history):
    """Reset the chat history."""
    return []

# Function to setup and update status
def setup_and_update():
    success = setup_system()
    if success:
        return "✅ System is ready! You can now ask questions about ABB products."
    else:
        return "⚠️ System setup encountered issues. Some features may be limited."

# Add CSS for image display
custom_css = """
.image-container {
    display: flex;
    flex-wrap: wrap;
    gap: 10px;
    margin-top: 15px;
}

.source-image {
    max-width: 300px;
    margin-bottom: 10px;
}

.source-image img {
    width: 100%;
    border: 1px solid #ddd;
    border-radius: 4px;
    padding: 5px;
}

.source-image p {
    font-size: 0.8rem;
    color: #666;
    margin-top: 5px;
}

.app-header {
    display: flex;
    align-items: center;
    margin-bottom: 20px;
    background-color: #f8f9fa;
    padding: 10px;
    border-radius: 10px;
}

.app-title {
    margin: 0;
    color: #d00d2d;
    font-size: 2.5rem;
}

.app-subtitle {
    margin: 0;
    color: #666;
}

.content-card, .status-card {
    background: white;
    border-radius: 10px;
    padding: 15px;
    box-shadow: 0 2px 10px rgba(0,0,0,0.1);
    margin-bottom: 15px;
}

.primary-button {
    background-color: #d00d2d !important;
    color: white !important;
}

.secondary-button {
    background-color: #f0f0f0 !important;
    color: #333 !important;
}

.input-area {
    margin-top: 10px;
}
"""

# Main Gradio application
def create_gradio_app():
    # Create the Gradio interface
    with gr.Blocks(css=custom_css) as app:
        # Setup status variable
        setup_status = gr.State("System is setting up. Please wait...")
        status_display = gr.Markdown("System is setting up. Please wait...")
        
        with gr.Column(scale=1):
            # Modern header
            with gr.Row(elem_classes="app-header"):
                with gr.Column(scale=1):
                    gr.Image(value="img/ABB-Logo.png", 
                             width=120,
                             height=120, 
                             interactive=False, 
                             label="ABB Logo")
                with gr.Column(scale=3):
                    gr.HTML('<h1 class="app-title">Ginnie</h1>')
                    gr.HTML('<p class="app-subtitle">Your AI assistant for ABB product information</p>')
            
            # Chat interface
            with gr.Row():
                with gr.Column(scale=3):
                    # Chat interface with custom styling
                    gr.HTML('<div class="content-card">')
                    chatbot = gr.Chatbot(
                        value=[],
                        elem_id="chatbot",
                        height=500,
                        show_copy_button=True,
                        avatar_images=["https://ui-avatars.com/api/?name=You&background=0D8ABC&color=fff", 
                                      "https://ui-avatars.com/api/?name=Ginnie&background=d00d2d&color=fff"],
                        render_markdown=True
                    )
                    
                    # Message input with better styling
                    with gr.Row(elem_classes="input-area"):
                        msg = gr.Textbox(
                            placeholder="Ask about ABB products...",
                            label="",
                            lines=2,
                            max_lines=5,
                            show_label=False
                        )
                        
                        send_btn = gr.Button("Send", elem_classes="primary-button")
                    
                    with gr.Row():
                        clear_btn = gr.Button("Clear Chat", elem_classes="secondary-button")
                    gr.HTML('</div>')
                
                with gr.Column(scale=1):
                    # Quick tips card
                    gr.HTML('<div class="status-card">')
                    gr.HTML('''
                    <h3>Quick Tips</h3>
                    <ul>
                        <li>Ask about specific ABB products</li>
                        <li>Inquire about technical specifications</li>
                        <li>Ask about installation and maintenance</li>
                        <li>Get help with troubleshooting</li>
                        <li>Ask to see images of specific products</li>
                    </ul>
                    ''')
                    gr.HTML('</div>')
                    
                    # System status
                    gr.HTML('<div class="status-card">')
                    status_display = gr.Markdown("System is setting up...")
                    gr.HTML('</div>')
                    
                    # Hidden model selection for admins (not primary focus)
                    with gr.Accordion("Admin Settings", open=False):
                        model_radio = gr.Radio(
                            ["Proprietary (Claude AI via AWS Bedrock)", "Open Source (Mistral AI)"],
                            label="Select AI Model",
                            value="Proprietary (Claude AI via AWS Bedrock)"
                        )
                        model_switch_btn = gr.Button("Switch Model", elem_classes="secondary-button")
            
            # Set up event handlers
            send_btn.click(
                process_message,
                [msg, chatbot],
                [chatbot],
                api_name="send_message"
            )
            
            msg.submit(
                process_message,
                [msg, chatbot],
                [chatbot],
                api_name="send_message_enter"
            )
            
            clear_btn.click(
                reset_chat,
                [chatbot],
                [chatbot],
                api_name="clear_chat"
            )
            
            model_switch_btn.click(
                switch_model,
                [model_radio],
                [status_display],
                api_name="switch_model"
            )

        # Add the system setup to run when the app loads
        app.load(setup_and_update, None, status_display)
    
    return app

# Main execution function
def main():
    # Create and launch the Gradio app
    app = create_gradio_app()
    
    # Launch the application
    app.queue()
    app.launch()

# Launch the application - make sure you're using the correct function name
if __name__ == "__main__":
    main()