import streamlit as st
import numpy as np
import faiss
import os
import PyPDF2
import docx
import pandas as pd
import requests
import json
from datetime import datetime
import time
import uuid
import sys
import tempfile

# Disable sentence-transformers for now (causing issues on HF)
# We'll use a simpler approach for embeddings if needed

# Initialize session state
def init_session_state():
    """Initialize all session state variables"""
    if 'conversations' not in st.session_state:
        st.session_state.conversations = {}
    if 'current_conv_id' not in st.session_state:
        st.session_state.current_conv_id = None
    if 'documents' not in st.session_state:
        st.session_state.documents = []
    if 'is_ready' not in st.session_state:
        st.session_state.is_ready = False
    if 'embedder' not in st.session_state:
        st.session_state.embedder = None  # Disable for now

# ✅ Load API key
def get_api_key():
    # Try secrets in Streamlit Cloud
    if hasattr(st, 'secrets'):
        try:
            api_key = st.secrets.get("MISTRAL_API_KEY", "").strip()
            if api_key:
                return api_key
        except:
            pass
    
    # Try environment variable
    api_key = os.getenv("MISTRAL_API_KEY", "").strip()
    return api_key

MISTRAL_API_KEY = get_api_key()
MISTRAL_API_URL = "https://api.mistral.ai/v1/chat/completions"

# --------------------------- Document Processing (Simplified) ---------------------------
def auto_process_document(file):
    """Automatically process uploaded document - simplified version"""
    try:
        if file is None:
            return False
        
        # Save uploaded file temporarily
        with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(file.name)[1]) as tmp_file:
            tmp_file.write(file.getbuffer())
            tmp_path = tmp_file.name
        
        ext = os.path.splitext(tmp_path)[1].lower()
        
        chunks = []
        if ext == '.txt':
            with open(tmp_path, 'r', encoding='utf-8') as f:
                content = f.read()
            chunks = split_text_into_chunks(content)
        elif ext == '.csv':
            df = pd.read_csv(tmp_path)
            chunks = dataframe_to_chunks(df)
        elif ext == '.pdf':
            chunks = pdf_to_chunks(tmp_path)
        elif ext in ['.docx', '.doc']:
            chunks = docx_to_chunks(tmp_path)
        
        # Clean up temp file
        os.unlink(tmp_path)
        
        if chunks:
            st.session_state.documents = chunks
            st.session_state.is_ready = True
            
            # Store document info in current conversation
            if st.session_state.current_conv_id:
                conv = st.session_state.conversations[st.session_state.current_conv_id]
                conv["has_document"] = True
                conv["document_name"] = file.name
            
            return True
        
        return False
        
    except Exception as e:
        st.error(f"Document processing error: {e}")
        return False

def split_text_into_chunks(text, chunk_size=400):
    sentences = text.split('.')
    chunks = []
    current_chunk = ""
    for sentence in sentences:
        sentence = sentence.strip()
        if not sentence:
            continue
        if len(current_chunk) + len(sentence) < chunk_size:
            current_chunk += sentence + '. '
        else:
            chunks.append(current_chunk.strip())
            current_chunk = sentence + '. '
    if current_chunk:
        chunks.append(current_chunk.strip())
    return chunks if chunks else [text[:500]]

def dataframe_to_chunks(df):
    chunks = []
    chunks.append(f"Columns: {' | '.join(df.columns)}")
    for idx, row in df.iterrows():
        data = " | ".join([f"{col}: {val}" for col, val in zip(df.columns, row) if pd.notna(val)])
        chunks.append(data[:500])
    return chunks

def pdf_to_chunks(path):
    try:
        with open(path, 'rb') as f:
            reader = PyPDF2.PdfReader(f)
            text = ""
            for i, page in enumerate(reader.pages):
                t = page.extract_text()
                if t:
                    text += f"--- Page {i+1} ---\n{t}\n"
            return split_text_into_chunks(text)
    except Exception as e:
        return [f"PDF content: Error reading PDF"]

def docx_to_chunks(path):
    try:
        doc = docx.Document(path)
        text = "\n".join([p.text for p in doc.paragraphs if p.text.strip()])
        return split_text_into_chunks(text)
    except Exception as e:
        return [f"DOCX content: Error reading document"]

# --------------------------- Conversation Management ---------------------------
def generate_conversation_title(first_message):
    """Generate conversation title from first message"""
    words = first_message.split()[:5]
    title = " ".join(words)
    if len(title) > 30:
        title = title[:27] + "..."
    return title if title.strip() else "New Chat"

def create_new_conversation():
    """Create a new conversation"""
    conv_id = str(uuid.uuid4())
    now = datetime.now().strftime("%H:%M")
    
    st.session_state.conversations[conv_id] = {
        "id": conv_id,
        "title": "New Chat",
        "created_at": now,
        "messages": [],
        "has_document": False,
        "document_name": None
    }
    
    st.session_state.current_conv_id = conv_id
    return conv_id

def get_current_conversation():
    if st.session_state.current_conv_id:
        return st.session_state.conversations[st.session_state.current_conv_id]
    return None

def add_message_to_conversation(role, content):
    """Add message to current conversation"""
    if not st.session_state.current_conv_id:
        create_new_conversation()
    
    conv = st.session_state.conversations[st.session_state.current_conv_id]
    
    # Generate title from first user message
    if role == "user" and len(conv["messages"]) == 0:
        conv["title"] = generate_conversation_title(content)
    
    conv["messages"].append({
        "role": role,
        "content": content,
        "timestamp": datetime.now().strftime("%H:%M")
    })

def get_conversation_context(turns=5):
    """Get recent conversation context"""
    conv = get_current_conversation()
    if not conv or not conv["messages"]:
        return ""
    
    messages = conv["messages"]
    recent = messages[-(turns*2):] if len(messages) > turns*2 else messages
    
    context = ""
    for msg in recent:
        if msg["role"] == "user":
            context += f"User: {msg['content']}\n"
        elif msg["role"] == "assistant":
            context += f"Assistant: {msg['content']}\n"
    
    return context.strip()

def get_document_context(query):
    """Simple document context - return first few chunks if query matches"""
    if not st.session_state.is_ready or not st.session_state.documents:
        return ""
    
    # Simple keyword matching (you can enhance this later)
    query_words = set(query.lower().split())
    relevant_chunks = []
    
    for chunk in st.session_state.documents[:3]:  # Return first 3 chunks
        chunk_lower = chunk.lower()
        if any(word in chunk_lower for word in query_words if len(word) > 3):
            relevant_chunks.append(chunk)
    
    if relevant_chunks:
        return "\n".join([f"• {chunk}" for chunk in relevant_chunks[:2]])
    
    return ""

# --------------------------- Mistral Integration ---------------------------
def call_mistral_llm(prompt, context=""):
    """Call Mistral API"""
    if not MISTRAL_API_KEY:
        return "⚠️ Please add your MISTRAL_API_KEY in the Space settings (Secrets section)."
    
    try:
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {MISTRAL_API_KEY}"
        }

        messages = [
            {"role": "system", "content": "You are a helpful AI assistant. Provide clear, concise, and accurate responses."}
        ]
        
        if context:
            messages.append({"role": "system", "content": f"Context:\n{context}"})
        
        messages.append({"role": "user", "content": prompt})

        payload = {
            "model": "mistral-large-latest",
            "messages": messages,
            "temperature": 0.7,
            "max_tokens": 1000
        }

        response = requests.post(MISTRAL_API_URL, headers=headers, json=payload, timeout=30)
        
        if response.status_code == 200:
            data = response.json()
            return data["choices"][0]["message"]["content"].strip()
        else:
            return f"⚠️ API Error: {response.status_code}"
            
    except Exception as e:
        return f"⚠️ Error: {str(e)}"

def generate_response(query):
    """Generate response with context"""
    # Get conversation context
    conversation_context = get_conversation_context(4)
    
    # Get document context if available
    doc_context = get_document_context(query) if st.session_state.is_ready else ""
    
    # Combine contexts
    full_context = ""
    if conversation_context:
        full_context += f"Previous conversation:\n{conversation_context}\n\n"
    if doc_context:
        full_context += f"Relevant document content:\n{doc_context}"
    
    # Generate response
    response = call_mistral_llm(query, full_context)
    
    # Save messages
    add_message_to_conversation("user", query)
    add_message_to_conversation("assistant", response)
    
    return response

# --------------------------- Main App ---------------------------
def main():
    # Initialize session state
    init_session_state()
    
    # Page config
    st.set_page_config(
        page_title="AI Assistant",
        page_icon="🤖",
        layout="wide",
        initial_sidebar_state="expanded"
    )
    
    # Custom CSS for ChatGPT-like design
    st.markdown("""
        <style>
        /* Main container */
        .main {
            padding: 0rem 1rem;
        }
        
        /* Sidebar */
        [data-testid="stSidebar"] {
            background-color: #171717;
        }
        
        /* Chat containers */
        [data-testid="stChatMessage"] {
            padding: 1rem;
            border-radius: 0.75rem;
            margin-bottom: 1rem;
        }
        
        /* User message */
        [data-testid="stChatMessage"][data-message-author="user"] {
            background-color: #2d2d2d;
        }
        
        /* Assistant message */
        [data-testid="stChatMessage"][data-message-author="assistant"] {
            background-color: #1a1a1a;
        }
        
        /* Input area */
        .stTextInput>div>div>input {
            background-color: #2d2d2d;
            color: white;
            border: 1px solid #4a4a4a;
        }
        
        /* Button styling */
        .stButton>button {
            background-color: #10a37f;
            color: white;
            border: none;
            padding: 0.5rem 1rem;
            border-radius: 0.5rem;
            font-weight: 500;
        }
        
        .stButton>button:hover {
            background-color: #0d8c6f;
        }
        
        /* Hide Streamlit branding */
        #MainMenu {visibility: hidden;}
        footer {visibility: hidden;}
        .stDeployButton {display:none;}
        
        /* Chat container */
        .stChatInputContainer {
            position: fixed;
            bottom: 20px;
            width: 50%;
            left: 25%;
        }
        </style>
    """, unsafe_allow_html=True)
    
    # Check API key
    if not MISTRAL_API_KEY:
        st.warning("""
        ⚠️ **API Key Required**
        
        To use the AI features, please add your MISTRAL_API_KEY:
        
        1. Go to your Space **Settings**
        2. Find **"Repository secrets"** section
        3. Add key: `MISTRAL_API_KEY`
        4. Add value: your_mistral_api_key_here
        
        Without the API key, the app will still work but AI responses won't be generated.
        """)
    
    # Sidebar - Chat History
    with st.sidebar:
        st.title("💬 Conversations")
        
        # New chat button
        if st.button("➕ New Chat", use_container_width=True, type="primary"):
            create_new_conversation()
            st.rerun()
        
        st.divider()
        
        # Chat history list
        st.subheader("History")
        
        # Display all conversations
        conv_items = list(st.session_state.conversations.items())
        
        if not conv_items:
            st.caption("No conversations yet")
        else:
            for conv_id, conv in conv_items[::-1]:  # Show newest first
                col1, col2 = st.columns([4, 1])
                
                with col1:
                    # Display conversation title with icon
                    icon = "📄" if conv.get("has_document", False) else "💬"
                    title_display = f"{icon} {conv['title']}"
                    
                    is_active = conv_id == st.session_state.current_conv_id
                    button_type = "primary" if is_active else "secondary"
                    
                    if st.button(
                        title_display,
                        key=f"conv_{conv_id}",
                        use_container_width=True,
                        type=button_type,
                        help=f"Created at {conv['created_at']}"
                    ):
                        st.session_state.current_conv_id = conv_id
                        st.rerun()
                
                with col2:
                    # Delete button
                    if st.button("🗑️", key=f"delete_{conv_id}", type="secondary"):
                        if conv_id == st.session_state.current_conv_id:
                            st.session_state.current_conv_id = None
                        del st.session_state.conversations[conv_id]
                        st.rerun()
        
        # Document upload in sidebar
        st.divider()
        st.subheader("📁 Documents")
        
        uploaded_file = st.file_uploader(
            "Upload document (TXT, PDF, DOCX, CSV)",
            type=['txt', 'pdf', 'docx', 'csv'],
            label_visibility="collapsed",
            key="file_uploader"
        )
        
        if uploaded_file is not None:
            with st.spinner("Processing document..."):
                if auto_process_document(uploaded_file):
                    st.success(f"✅ {uploaded_file.name} loaded")
                    if st.session_state.current_conv_id:
                        conv = st.session_state.conversations[st.session_state.current_conv_id]
                        conv["has_document"] = True
                        conv["document_name"] = uploaded_file.name
                else:
                    st.error("Failed to process document")
        
        st.divider()
        st.caption("AI Assistant v1.0")
    
    # Main chat area
    col1, col2, col3 = st.columns([1, 2, 1])
    
    with col2:
        # Header
        st.markdown("<h1 style='text-align: center; margin-bottom: 2rem;'>🤖 AI Assistant</h1>", 
                   unsafe_allow_html=True)
        
        # Display current conversation info
        current_conv = get_current_conversation()
        if current_conv:
            # Show document indicator if present
            if current_conv.get("has_document", False):
                st.info(f"📄 Document loaded: {current_conv['document_name']}")
        
        # Chat container
        chat_container = st.container(height=550)
        
        # Initialize chat if needed
        if not st.session_state.current_conv_id:
            create_new_conversation()
        
        # Display messages
        with chat_container:
            current_conv = get_current_conversation()
            if current_conv and current_conv["messages"]:
                for msg in current_conv["messages"]:
                    with st.chat_message(msg["role"]):
                        st.markdown(msg["content"])
            else:
                st.markdown("""
                <div style='text-align: center; padding: 40px; color: #666;'>
                    <h3>👋 Hello! I'm your AI Assistant</h3>
                    <p>Start a conversation by typing a message below.</p>
                    <p>You can also upload documents for context.</p>
                </div>
                """, unsafe_allow_html=True)
        
        # Chat input
        user_input = st.chat_input("Type your message here...")
        
        if user_input:
            # Display user message
            with st.chat_message("user"):
                st.markdown(user_input)
            
            # Generate and display assistant response
            with st.chat_message("assistant"):
                if MISTRAL_API_KEY:
                    with st.spinner("Thinking..."):
                        response = generate_response(user_input)
                        st.markdown(response)
                else:
                    st.error("API key not configured. Please add MISTRAL_API_KEY in settings.")
            
            # Rerun to update UI
            st.rerun()
        
        # Clear chat button
        if current_conv and current_conv["messages"]:
            if st.button("Clear Chat", use_container_width=True, type="secondary"):
                current_conv["messages"] = []
                st.rerun()

# --------------------------- Run App ---------------------------
if __name__ == "__main__":
    main()