import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
from threading import Thread
import pdfplumber
from docx import Document
import pandas as pd
import io
from typing import List, Tuple, Optional
import json

# Memory Agent Class
class MemoryAgent:
    def __init__(self):
        self.memory_store = []
        self.documents = {}
    
    def add(self, text: str, source: str = "conversation"):
        """Add text to memory with source tracking"""
        self.memory_store.append({
            "content": text,
            "source": source
        })
    
    def add_document(self, filename: str, content: str):
        """Store document content"""
        self.documents[filename] = content
        # Store a summary in memory, but keep full content in documents dict
        summary = content[:1000] + "..." if len(content) > 1000 else content
        self.add(f"Document '{filename}' uploaded (Length: {len(content)} chars). Preview:\n{summary}", source=f"document:{filename}")
    
    def get_context(self, query: Optional[str] = None, max_items: int = 10) -> str:
        """Retrieve relevant context from memory"""
        if not self.memory_store:
            return ""
        
        # Return recent memory items
        recent_items = self.memory_store[-max_items:]
        context_parts = []
        
        for item in recent_items:
            if item["source"].startswith("document:"):
                context_parts.append(f"[From {item['source'].split(':')[1]}]\n{item['content']}")
            else:
                context_parts.append(item["content"])
        
        return "\n\n".join(context_parts)
    
    def clear(self):
        """Clear all memory"""
        self.memory_store = []
        self.documents = {}

# Document Processing Functions
def extract_text_from_pdf(file_path: str) -> str:
    """Extract text from PDF using pdfplumber"""
    text_parts = []
    try:
        with pdfplumber.open(file_path) as pdf:
            for page in pdf.pages:
                text = page.extract_text()
                if text:
                    text_parts.append(text)
        return "\n\n".join(text_parts)
    except Exception as e:
        return f"Error extracting PDF: {str(e)}"

def extract_text_from_docx(file_path: str) -> str:
    """Extract text from DOCX"""
    try:
        doc = Document(file_path)
        return "\n\n".join([para.text for para in doc.paragraphs if para.text.strip()])
    except Exception as e:
        return f"Error extracting DOCX: {str(e)}"

def extract_text_from_txt(file_path: str) -> str:
    """Extract text from TXT file"""
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            return f.read()
    except Exception as e:
        return f"Error extracting TXT: {str(e)}"

def extract_text_from_csv(file_path: str) -> str:
    """Extract text from CSV"""
    try:
        df = pd.read_csv(file_path)
        return df.to_string()
    except Exception as e:
        return f"Error extracting CSV: {str(e)}"

def process_uploaded_file(file_path: str) -> Tuple[str, str]:
    """Process uploaded file and extract text"""
    filename = file_path.split('/')[-1]
    extension = filename.split('.')[-1].lower()
    
    if extension == 'pdf':
        content = extract_text_from_pdf(file_path)
    elif extension == 'docx':
        content = extract_text_from_docx(file_path)
    elif extension == 'txt':
        content = extract_text_from_txt(file_path)
    elif extension == 'csv':
        content = extract_text_from_csv(file_path)
    else:
        content = f"Unsupported file type: {extension}"
    
    return filename, content

# Global model cache
model_cache = {}
tokenizer_cache = {}

def load_model(model_name: str):
    """Load model and tokenizer with caching"""
    if model_name not in model_cache:
        try:
            tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
            # Set padding token if not exists
            if tokenizer.pad_token is None:
                tokenizer.pad_token = tokenizer.eos_token
            
            model = AutoModelForCausalLM.from_pretrained(
                model_name,
                device_map="auto",
                torch_dtype=torch.bfloat16,
                trust_remote_code=True,
                low_cpu_mem_usage=True
            )
            model_cache[model_name] = model
            tokenizer_cache[model_name] = tokenizer
        except Exception as e:
            raise Exception(f"Error loading model: {str(e)}")
    
    return model_cache[model_name], tokenizer_cache[model_name]

def format_chat_history(history: List[List[str]]) -> str:
    """Format chat history for context"""
    formatted = []
    for user_msg, assistant_msg in history:
        if user_msg:
            formatted.append(f"User: {user_msg}")
        if assistant_msg:
            formatted.append(f"Assistant: {assistant_msg}")
    return "\n".join(formatted)

def generate_response(
    message: str,
    history: List[List[str]],
    model_name: str,
    temperature: float,
    max_tokens: int,
    system_prompt: str,
    memory_agent: MemoryAgent
):
    """Generate streaming response from SaulLM model"""
    try:
        # Load model
        model, tokenizer = load_model(model_name)
        
        # Build context from memory
        memory_context = memory_agent.get_context(max_items=15)
        
        # Build full prompt with enhanced structure
        context_parts = [system_prompt]
        
        # Add uploaded documents context if available
        if memory_agent.documents:
            context_parts.append(f"\n### UPLOADED DOCUMENTS:")
            context_parts.append("You have access to the following documents that the user uploaded:")
            for doc_name, doc_content in memory_agent.documents.items():
                # Include more document content for better context
                content_preview = doc_content[:5000] if len(doc_content) > 5000 else doc_content
                context_parts.append(f"\n**Document: {doc_name}**\n{content_preview}")
                if len(doc_content) > 5000:
                    context_parts.append(f"[Document continues... Total length: {len(doc_content)} characters]")
            context_parts.append("\nWhen the user asks about documents, refer to the specific content above and provide detailed, accurate answers based on what you can see in the uploaded documents.")
        
        # Add general memory/context
        if memory_context and not memory_agent.documents:
            context_parts.append(f"\n### Conversation Context:\n{memory_context}")
        
        # Add chat history
        if history:
            recent_history = history[-4:]  # Last 4 exchanges to leave room for documents
            context_parts.append(f"\n### Recent Conversation:\n{format_chat_history(recent_history)}")
        
        context_parts.append(f"\n### Current User Question:\n{message}\n\n### Your Response (provide a detailed, professional answer):")
        
        full_prompt = "\n".join(context_parts)
        
        # Tokenize
        inputs = tokenizer(full_prompt, return_tensors="pt", truncation=True, max_length=4096)
        inputs = {k: v.to(model.device) for k, v in inputs.items()}
        
        # Setup streamer
        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
        
        # Generation parameters
        generation_kwargs = dict(
            inputs,
            streamer=streamer,
            max_new_tokens=max_tokens,
            temperature=temperature,
            do_sample=temperature > 0,
            top_p=0.95,
            repetition_penalty=1.1
        )
        
        # Start generation in thread
        thread = Thread(target=model.generate, kwargs=generation_kwargs)
        thread.start()
        
        # Stream response
        partial_response = ""
        for new_text in streamer:
            partial_response += new_text
            # Clean up special tokens and formatting
            cleaned_response = partial_response.replace("[/INST]", "").replace("[INST]", "").replace("</s>", "").replace("<s>", "").strip()
            yield cleaned_response
        
        # Add to memory (use cleaned version)
        memory_agent.add(f"User: {message}\nAssistant: {cleaned_response}")
        
    except Exception as e:
        yield f"Error generating response: {str(e)}"

def handle_file_upload(files, memory_agent: MemoryAgent):
    """Handle multiple file uploads"""
    if not files:
        return "No files uploaded", []
    
    uploaded_docs = []
    for file in files:
        filename, content = process_uploaded_file(file.name)
        memory_agent.add_document(filename, content)
        uploaded_docs.append(filename)
    
    doc_list = "\n".join([f"✓ {doc}" for doc in uploaded_docs])
    return f"Uploaded {len(uploaded_docs)} file(s):\n{doc_list}", uploaded_docs

# Default system prompt
DEFAULT_SYSTEM_PROMPT = """You are SaulLM, a helpful and knowledgeable AI assistant.

Answer questions directly and professionally across all domains including law, finance, securitization, technology, science, and general knowledge.

When documents are uploaded:
- Read the document content carefully
- Answer questions based on the specific information in the documents
- Reference relevant parts of the document when answering
- Be precise and accurate with document details

Response style:
- Provide direct answers without unnecessary preambles
- Don't say "as an AI" or similar phrases
- Be clear, concise, and professional
- Give complete, helpful responses
- Use your expertise to provide accurate information"""

# Build Gradio Interface
with gr.Blocks(title="SaulLM Chat", theme=gr.themes.Soft()) as demo:
    # State variables
    memory_state = gr.State(lambda: MemoryAgent())
    uploaded_docs_state = gr.State([])
    
    gr.Markdown("# SaulLM Chat")

    
    with gr.Row():
        # Left Sidebar
        with gr.Column(scale=1, min_width=300):
            gr.Markdown("### Configuration")
            
            model_dropdown = gr.Dropdown(
                choices=[
                    "Equall/Saul-7B-Instruct-v1",
                    "Equall/SaulLM-54B-Instruct",
                    "Equall/SaulLM-141B-Instruct"
                ],
                value="Equall/Saul-7B-Instruct-v1",
                label="Model Selection",
                interactive=True
            )
            
            load_model_btn = gr.Button("Load Model", variant="primary")
            model_load_status = gr.Textbox(
                label="Model Status",
                value="No model loaded yet",
                interactive=False,
                lines=2
            )
            
            temperature_slider = gr.Slider(
                minimum=0.0,
                maximum=1.0,
                value=0.7,
                step=0.1,
                label="Temperature",
                info="Higher = more creative"
            )
            
            max_tokens_slider = gr.Slider(
                minimum=128,
                maximum=2048,
                value=1024,
                step=128,
                label="Max Tokens",
                info="Maximum response length"
            )
            
            system_prompt_box = gr.Textbox(
                value=DEFAULT_SYSTEM_PROMPT,
                label="System Prompt",
                lines=10,
                max_lines=15,
                info="Define the assistant's behavior and expertise"
            )
            
            update_prompt_btn = gr.Button("Update System Prompt", variant="secondary")
            prompt_update_status = gr.Textbox(
                label="Prompt Status",
                value="Using default system prompt",
                interactive=False,
                lines=1
            )
            
            clear_btn = gr.Button("🗑️ Clear Chat & Memory", variant="stop")
        
        # Main Chat Area
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(
                label="Chat",
                height=600,
                bubble_full_width=False,
                show_copy_button=True
            )
            
            with gr.Row():
                msg_input = gr.Textbox(
                    label="Message",
                    placeholder="Ask me...",
                    lines=2,
                    scale=9
                )
                send_btn = gr.Button("Send", variant="primary", scale=1)
            
            with gr.Row():
                regenerate_btn = gr.Button("Regenerate Response")
                clear_input_btn = gr.Button("Clear Input")
            
            gr.Markdown("### Document Upload")
            file_upload = gr.File(
                label="Upload Documents (PDF, DOCX, TXT, CSV)",
                file_count="multiple",
                file_types=[".pdf", ".docx", ".txt", ".csv"]
            )
            
            with gr.Row():
                upload_status = gr.Textbox(
                    label="Upload Status",
                    interactive=False,
                    lines=2,
                    scale=1
                )
                uploaded_docs_display = gr.Textbox(
                    label="Uploaded Documents",
                    interactive=False,
                    lines=2,
                    scale=1
                )
    
    # Event Handlers
    def user_submit(message, history):
        """Handle user message submission"""
        return "", history + [[message, None]]
    
    def bot_response(history, model_name, temperature, max_tokens, system_prompt, memory_agent):
        """Generate bot response"""
        if not history or history[-1][1] is not None:
            return history
        
        user_message = history[-1][0]
        history[-1][1] = ""
        
        for partial_response in generate_response(
            user_message,
            history[:-1],
            model_name,
            temperature,
            max_tokens,
            system_prompt,
            memory_agent
        ):
            history[-1][1] = partial_response
            yield history
    
    def regenerate_last(history, model_name, temperature, max_tokens, system_prompt, memory_agent):
        """Regenerate the last response"""
        if not history:
            return history
        
        # Remove last assistant response
        if history[-1][1]:
            history[-1][1] = None
        
        # Generate new response
        for updated_history in bot_response(history, model_name, temperature, max_tokens, system_prompt, memory_agent):
            yield updated_history
    
    def clear_chat(memory_agent):
        """Clear chat and memory"""
        memory_agent.clear()
        return [], "", "Chat and memory cleared"
    
    def update_uploaded_docs_display(docs):
        """Update the uploaded documents display"""
        if not docs:
            return ""
        return "\n".join([f"✓ {doc}" for doc in docs])
    
    def load_model_handler(model_name):
        """Handle model loading button click"""
        try:
            load_model(model_name)
            return f"✓ Model '{model_name}' loaded successfully!"
        except Exception as e:
            return f" Error loading model: {str(e)}"
    
    def update_system_prompt_handler(prompt):
        """Handle system prompt update"""
        if prompt.strip():
            return "✓ System prompt updated successfully!"
        else:
            return "System prompt cannot be empty"
    
    # Wire up events
    load_model_btn.click(
        load_model_handler,
        [model_dropdown],
        model_load_status
    )
    
    update_prompt_btn.click(
        update_system_prompt_handler,
        [system_prompt_box],
        prompt_update_status
    )
    
    msg_input.submit(
        user_submit,
        [msg_input, chatbot],
        [msg_input, chatbot],
        queue=False
    ).then(
        bot_response,
        [chatbot, model_dropdown, temperature_slider, max_tokens_slider, system_prompt_box, memory_state],
        chatbot
    )
    
    send_btn.click(
        user_submit,
        [msg_input, chatbot],
        [msg_input, chatbot],
        queue=False
    ).then(
        bot_response,
        [chatbot, model_dropdown, temperature_slider, max_tokens_slider, system_prompt_box, memory_state],
        chatbot
    )
    
    regenerate_btn.click(
        regenerate_last,
        [chatbot, model_dropdown, temperature_slider, max_tokens_slider, system_prompt_box, memory_state],
        chatbot
    )
    
    clear_input_btn.click(
        lambda: "",
        None,
        msg_input
    )
    
    file_upload.change(
        handle_file_upload,
        [file_upload, memory_state],
        [upload_status, uploaded_docs_state]
    ).then(
        update_uploaded_docs_display,
        [uploaded_docs_state],
        uploaded_docs_display
    )
    
    clear_btn.click(
        clear_chat,
        [memory_state],
        [chatbot, uploaded_docs_display, upload_status]
    )

# Launch
if __name__ == "__main__":
    demo.queue(max_size=20)
    demo.launch()