app.py

import streamlit as st
import google.generativeai as genai
from PyPDF2 import PdfReader
import os
import re
import json
import pickle
import hashlib
from datetime import datetime
from pathlib import Path
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_huggingface.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain.schema import Document
import tempfile
import warnings
import numpy as np
import shutil
import time
warnings.filterwarnings('ignore')

# Configure page
st.set_page_config(
    page_title="Ashok 2.0 - AI Problem Solving Assistant",
    page_icon="🧠",
    layout="centered",
    initial_sidebar_state="collapsed"
)

# World-class minimal UI styling
st.markdown("""
<style>
    /* Import premium font */
    @import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap');
    
    /* Global Reset & Base */
    .stApp {
        background: linear-gradient(135deg, #f8fafc 0%, #f1f5f9 100%);
        font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif;
    }
    
    .main .block-container {
        padding-top: 2rem;
        padding-bottom: 2rem;
        max-width: 800px;
    }
    
    /* Hide Streamlit UI elements */
    #MainMenu {visibility: hidden;}
    footer {visibility: hidden;}
    header {visibility: hidden;}
    .stDeployButton {display: none;}
    
    /* Main Heading */
    .main-title {
    font-size: 4rem;
    font-weight: 700;
    text-align: center;
    color: #1f2937;  /* Professional blue */
    margin: 2rem 0 3rem 0;
    letter-spacing: -0.02em;
    line-height: 1.1;
    text-shadow: 2px 2px 4px rgba(0,0,0,0.1);
    }
    
    /* API Key Setup Container */
    .api-setup-container {
        background: rgba(255, 255, 255, 0.9);
        backdrop-filter: blur(10px);
        border: 1px solid rgba(229, 231, 235, 0.6);
        border-radius: 24px;
        padding: 3rem;
        margin: 2rem auto;
        max-width: 500px;
        box-shadow: 0 20px 25px -5px rgba(0, 0, 0, 0.05), 0 10px 10px -5px rgba(0, 0, 0, 0.04);
        text-align: center;
    }
    
    .api-setup-title {
        font-size: 1.5rem;
        font-weight: 600;
        color: #1f2937;
        margin-bottom: 1rem;
    }
    
    .api-setup-subtitle {
        color: #6b7280;
        margin-bottom: 2rem;
        line-height: 1.6;
    }
    
    /* API Key Input Styling - Fixed */
    .stTextInput > div > div > input {
    background: rgba(255, 255, 255, 0.9);
    border: 2px solid rgba(229, 231, 235, 0.6);
    border-radius: 16px;
    padding: 1rem 1.5rem;
    font-size: 1rem;
    transition: all 0.3s ease;
    backdrop-filter: blur(5px);
    color: #1f2937 !important;
    }

    .stTextInput > div > div > input::placeholder {
    color: #9ca3af !important;
    opacity: 1;
    }

    .stTextInput > div > div > input:focus {
    border-color: #3b82f6;
    box-shadow: 0 0 0 4px rgba(59, 130, 246, 0.1);
    outline: none;
    background: rgba(255, 255, 255, 0.95);
    color: #1f2937 !important;
    }

    .stTextInput label {
    font-weight: 500;
    color: #374151 !important;
    margin-bottom: 0.5rem;
    }
    
    /* Chat Interface */
    .chat-container {
        background: rgba(255, 255, 255, 0.7);
        backdrop-filter: blur(10px);
        border: 1px solid rgba(229, 231, 235, 0.4);
        border-radius: 24px;
        padding: 2rem;
        margin: 2rem 0;
        box-shadow: 0 10px 15px -3px rgba(0, 0, 0, 0.03), 0 4px 6px -2px rgba(0, 0, 0, 0.02);
    }
    
    /* Chat Messages */
    .stChatMessage {
        background: rgba(255, 255, 255, 0.9) !important;
        border: 1px solid rgba(0, 0, 0, 0.08) !important;
        border-radius: 16px !important;
        margin-bottom: 1rem !important;
        padding: 1.5rem !important;
        box-shadow: 0 2px 4px rgba(0, 0, 0, 0.02) !important;
        backdrop-filter: blur(5px) !important;
    }
    
    .stChatMessage[data-testid="chat-message-user"] {
        background: rgba(248, 250, 252, 0.9) !important;
        border-left: 3px solid #3b82f6 !important;
    }
    
    .stChatMessage[data-testid="chat-message-assistant"] {
        background: rgba(255, 255, 255, 0.9) !important;
        border-left: 3px solid #10b981 !important;
    }
    
    .stChatMessage .stMarkdown {
        color: #1f2937 !important;
        line-height: 1.6;
    }
    
    .stChatMessage .stMarkdown p {
        color: #1f2937 !important;
        margin-bottom: 0.5rem;
    }
    
    /* Chat Input */
    .stChatInput > div {
        background: rgba(255, 255, 255, 0.9);
        border: 1px solid rgba(229, 231, 235, 0.6);
        border-radius: 20px;
        backdrop-filter: blur(10px);
    }
    
    .stChatInput input {
        color: #1f2937 !important;
        background: transparent !important;
        border: none !important;
        padding: 1rem 1.5rem !important;
    }
    
    .stChatInput input::placeholder {
        color: #ffffff !important;
    }
    
    /* Quick Actions */
    .quick-actions-container {
        background: rgba(255, 255, 255, 0.6);
        backdrop-filter: blur(10px);
        border: 1px solid rgba(229, 231, 235, 0.4);
        border-radius: 20px;
        padding: 2rem;
        margin: 2rem 0;
    }
    
    .quick-actions-title {
        font-size: 1.25rem;
        font-weight: 600;
        color: #1f2937;
        text-align: center;
        margin-bottom: 1.5rem;
    }
    
    .quick-action-grid {
        display: grid;
        grid-template-columns: repeat(auto-fit, minmax(280px, 1fr));
        gap: 1rem;
        margin-top: 1rem;
    }
    
    .quick-action-item {
        background: rgba(255, 255, 255, 0.8);
        border: 1px solid rgba(229, 231, 235, 0.5);
        border-radius: 16px;
        padding: 1.5rem;
        cursor: pointer;
        transition: all 0.3s ease;
        text-align: left;
        backdrop-filter: blur(5px);
    }
    
    .quick-action-item:hover {
        background: rgba(248, 250, 252, 0.9);
        border-color: #3b82f6;
        transform: translateY(-2px);
        box-shadow: 0 8px 25px rgba(59, 130, 246, 0.1);
    }
    
    .quick-action-icon {
        font-size: 1.5rem;
        margin-bottom: 0.5rem;
    }
    
    .quick-action-title {
        font-weight: 600;
        color: #1f2937;
        margin-bottom: 0.5rem;
        font-size: 1rem;
    }
    
    .quick-action-desc {
        color: #6b7280;
        font-size: 0.9rem;
        line-height: 1.5;
    }
    
    /* Control Panel */
    .control-panel {
        display: flex;
        justify-content: center;
        gap: 1rem;
        margin: 2rem 0;
    }
    
    .control-button {
        background: rgba(255, 255, 255, 0.8);
        border: 1px solid rgba(229, 231, 235, 0.6);
        border-radius: 12px;
        padding: 0.75rem 1.5rem;
        color: #374151;
        font-weight: 500;
        cursor: pointer;
        transition: all 0.3s ease;
        backdrop-filter: blur(5px);
    }
    
    .control-button:hover {
        background: rgba(248, 250, 252, 0.9);
        border-color: #3b82f6;
        color: #1f2937;
        transform: translateY(-1px);
    }
    
    /* Status Indicators */
    .status-success {
        background: linear-gradient(135deg, rgba(16, 185, 129, 0.1) 0%, rgba(5, 150, 105, 0.1) 100%);
        border: 1px solid rgba(16, 185, 129, 0.2);
        border-radius: 16px;
        padding: 1rem 1.5rem;
        margin: 1rem 0;
        color: #065f46;
        font-weight: 500;
        text-align: center;
        backdrop-filter: blur(5px);
    }
    
    .status-warning {
        background: linear-gradient(135deg, rgba(245, 158, 11, 0.1) 0%, rgba(217, 119, 6, 0.1) 100%);
        border: 1px solid rgba(245, 158, 11, 0.2);
        border-radius: 16px;
        padding: 1rem 1.5rem;
        margin: 1rem 0;
        color: #92400e;
        font-weight: 500;
        text-align: center;
        backdrop-filter: blur(5px);
    }
    
    /* Learning Indicator */
    .learning-indicator {
        background: linear-gradient(135deg, rgba(16, 185, 129, 0.9) 0%, rgba(5, 150, 105, 0.9) 100%);
        color: white;
        padding: 1rem 1.5rem;
        border-radius: 16px;
        margin: 1rem 0;
        font-weight: 500;
        text-align: center;
        box-shadow: 0 4px 6px rgba(16, 185, 129, 0.2);
        backdrop-filter: blur(10px);
    }
    
    /* Typing Indicator */
    .typing-indicator {
        background: rgba(248, 250, 252, 0.9);
        border: 1px solid rgba(229, 231, 235, 0.6);
        border-radius: 16px;
        padding: 1rem 1.5rem;
        margin: 1rem 0;
        display: flex;
        align-items: center;
        gap: 0.75rem;
        color: #6b7280;
        backdrop-filter: blur(5px);
    }
    
    .typing-dots {
        display: flex;
        gap: 4px;
    }
    
    .typing-dot {
        width: 6px;
        height: 6px;
        border-radius: 50%;
        background-color: #9ca3af;
        animation: typing 1.4s infinite ease-in-out;
    }
    
    .typing-dot:nth-child(1) { animation-delay: -0.32s; }
    .typing-dot:nth-child(2) { animation-delay: -0.16s; }
    
    @keyframes typing {
        0%, 80%, 100% { transform: scale(0.8); opacity: 0.4; }
        40% { transform: scale(1); opacity: 1; }
    }
    
    /* Button Styling */
    .stButton > button {
        background: linear-gradient(135deg, rgba(59, 130, 246, 0.9) 0%, rgba(37, 99, 235, 0.9) 100%);
        color: white !important;
        border: none;
        border-radius: 12px;
        padding: 0.75rem 2rem;
        font-weight: 500;
        transition: all 0.3s ease;
        backdrop-filter: blur(10px);
        box-shadow: 0 4px 6px rgba(59, 130, 246, 0.2);
    }
    
    .stButton > button:hover {
        background: linear-gradient(135deg, rgba(37, 99, 235, 0.9) 0%, rgba(29, 78, 216, 0.9) 100%);
        transform: translateY(-1px);
        box-shadow: 0 6px 12px rgba(59, 130, 246, 0.25);
    }
    
    /* API Key Guide */
    .api-guide {
        background: rgba(255, 255, 255, 0.6);
        border: 1px solid rgba(229, 231, 235, 0.4);
        border-radius: 16px;
        padding: 1.5rem;
        margin: 1.5rem 0;
        backdrop-filter: blur(5px);
    }
    
    .api-guide-title {
        font-weight: 600;
        color: #1f2937;
        margin-bottom: 1rem;
    }
    
    .api-guide-text {
        color: #6b7280;
        line-height: 1.6;
        margin-bottom: 0.5rem;
    }
    
    .api-link {
        color: #3b82f6;
        text-decoration: none;
        font-weight: 500;
        transition: color 0.2s ease;
    }
    
    .api-link:hover {
        color: #1d4ed8;
    }
    
    /* Responsive Design */
    @media (max-width: 768px) {
        .main-title {
            font-size: 3rem;
        }
        
        .api-setup-container {
            padding: 2rem;
            margin: 1rem;
        }
        
        .quick-action-grid {
            grid-template-columns: 1fr;
        }
        
        .control-panel {
            flex-direction: column;
            align-items: center;
        }
    }
    .stChatFloatingInputContainer{
    background-color: transparent !important;
    }
</style>
""", unsafe_allow_html=True)

class PersistentHFKnowledgeBase:
    """Persistent Knowledge Base with silent initialization"""
    
    def __init__(self):
        # Create persistent directories
        self.data_dir = Path("./persistent_data")
        self.data_dir.mkdir(exist_ok=True)
        
        # File paths for persistence
        self.vectorstore_path = self.data_dir / "vectorstore"
        self.metadata_path = self.data_dir / "metadata.json"
        self.conversations_path = self.data_dir / "conversations.json"
        self.stats_path = self.data_dir / "stats.json"
        self.init_flag_path = self.data_dir / "initialized.flag"
        
        # Initialize components
        self.embeddings = None
        self.vectorstore = None
        self.metadata = {}
        self.conversations = []
        self.stats = {}
        
        # Initialize system silently
        self.initialize_system()
    
    def initialize_system(self):
        """Initialize the complete system with silent book processing"""
        try:
            # Initialize embeddings first
            self.init_embeddings()
            
            # Check if system was already initialized
            if self.init_flag_path.exists():
                # Load existing knowledge base silently
                self.load_existing_knowledge()
            else:
                # First time initialization - do it silently
                self.first_time_initialization()
                    
        except Exception as e:
            # Silent fallback initialization
            self.fallback_initialization()
    
    def init_embeddings(self):
        """Initialize embeddings model with silent caching"""
        if self.embeddings is None:
            try:
                cache_dir = self.data_dir / "embeddings_cache"
                cache_dir.mkdir(exist_ok=True)
                
                self.embeddings = HuggingFaceEmbeddings(
                    model_name="sentence-transformers/all-MiniLM-L6-v2",
                    cache_folder=str(cache_dir)
                )
            except Exception as e:
                return False
        return True
    
    def first_time_initialization(self):
        """Complete first-time setup with silent book processing"""
        try:
            # Initialize metadata
            self.metadata = {
                'version': '2.0-minimal',
                'created_at': datetime.now().isoformat(),
                'last_updated': datetime.now().isoformat(),
                'total_documents': 0,
                'book_processed': False,
                'book_info': {},
                'initialization_complete': False
            }
            
            # Initialize stats
            self.stats = {
                'total_queries': 0,
                'learning_sessions': 0,
                'book_chunks': 0,
                'conversation_chunks': 0,
                'silly_questions_blocked': 0
            }
            
            # Initialize conversations
            self.conversations = []
            
            # Process book if available (silently)
            book_processed = self.process_startup_book()
            
            # Create default knowledge if no book
            if not book_processed:
                self.create_default_knowledge()
            
            # Mark as initialized
            self.metadata['initialization_complete'] = True
            self.save_all_data()
            
            # Create initialization flag
            with open(self.init_flag_path, 'w') as f:
                f.write(f"Initialized on {datetime.now().isoformat()}")
            
        except Exception as e:
            self.fallback_initialization()
    
    def process_startup_book(self):
        """Process the book included with the deployment (silently)"""
        book_paths = [
            "book.pdf",
            "problem_solving_book.pdf", 
            "default_book.pdf",
            "ashok_book.pdf"
        ]
        
        for book_path in book_paths:
            if Path(book_path).exists():
                success = self.process_book_file(Path(book_path))
                if success:
                    return True
        
        return False
    
    def process_book_file(self, book_path):
        """Process a specific book file (silently)"""
        try:
            # Extract text from PDF
            reader = PdfReader(str(book_path))
            page_texts = []
            
            for page_num, page in enumerate(reader.pages, 1):
                page_text = page.extract_text()
                if page_text.strip():
                    page_texts.append({
                        'page': page_num,
                        'text': page_text,
                        'word_count': len(page_text.split())
                    })
            
            if not page_texts:
                return False
            
            # Create book info
            book_info = {
                'title': book_path.name,
                'path': str(book_path),
                'pages': len(page_texts),
                'processed_at': datetime.now().isoformat(),
                'source': 'deployment_book'
            }
            
            # Process content
            success, message = self.process_book_content("", page_texts, book_info)
            
            return success
                
        except Exception as e:
            return False
    
    def create_default_knowledge(self):
        """Create comprehensive default knowledge base"""
        default_knowledge = [
            {
                "content": "Problem-solving methodology: 1) Problem Definition - Clearly articulate what needs to be solved, 2) Information Gathering - Collect relevant data and context, 3) Root Cause Analysis - Identify underlying causes, not just symptoms, 4) Solution Generation - Brainstorm multiple potential solutions, 5) Solution Evaluation - Assess feasibility, impact, and resources, 6) Implementation Planning - Create detailed action steps, 7) Execution and Monitoring - Implement and track progress, 8) Review and Learning - Evaluate outcomes and extract lessons.",
                "metadata": {
                    "source": "core_knowledge",
                    "type": "framework",
                    "topic": "problem_solving_process",
                    "chapter": "Core Problem-Solving Framework"
                }
            },
            {
                "content": "Decision-making best practices: Use the DECIDE framework - D: Define the problem clearly, E: Establish criteria for solutions, C: Consider alternatives systematically, I: Identify best alternatives using criteria, D: Develop and implement action plan, E: Evaluate and monitor solution effectiveness. Always consider stakeholder impact, resource constraints, time limitations, and potential risks.",
                "metadata": {
                    "source": "core_knowledge",
                    "type": "framework",
                    "topic": "decision_making",
                    "chapter": "Decision-Making Framework"
                }
            },
            {
                "content": "Conflict resolution strategies: 1) Active Listening - Understand all perspectives without judgment, 2) Identify Interests - Focus on underlying needs, not stated positions, 3) Find Common Ground - Identify shared goals and values, 4) Generate Options - Create win-win solutions collaboratively, 5) Use Objective Criteria - Apply fair standards for evaluation, 6) Separate People from Problems - Address issues, not personalities, 7) Maintain Relationships - Preserve working relationships while solving problems.",
                "metadata": {
                    "source": "core_knowledge",
                    "type": "strategy",
                    "topic": "conflict_resolution",
                    "chapter": "Conflict Resolution Techniques"
                }
            },
            {
                "content": "Critical thinking skills development: Analysis (breaking complex information into components), Evaluation (assessing credibility and logical strength), Inference (drawing reasonable conclusions), Interpretation (understanding meaning and significance), Explanation (articulating reasoning clearly), Self-regulation (monitoring and correcting one's thinking). Practice questioning assumptions, considering multiple perspectives, examining evidence quality, and recognizing logical fallacies.",
                "metadata": {
                    "source": "core_knowledge",
                    "type": "skills",
                    "topic": "critical_thinking",
                    "chapter": "Critical Thinking Development"
                }
            },
            {
                "content": "Team problem-solving dynamics: Establish psychological safety for open communication, define roles and responsibilities clearly, use structured problem-solving processes, encourage diverse perspectives, facilitate effective meetings, manage conflicts constructively, ensure equal participation, document decisions and action items, follow up on commitments, celebrate successes and learn from failures.",
                "metadata": {
                    "source": "core_knowledge",
                    "type": "team_dynamics",
                    "topic": "team_problem_solving",
                    "chapter": "Team Collaboration for Problem Solving"
                }
            }
        ]
        
        # Create documents
        documents = []
        for item in default_knowledge:
            doc = Document(
                page_content=item["content"],
                metadata=item["metadata"]
            )
            documents.append(doc)
        
        # Create vectorstore
        if documents and self.embeddings:
            self.vectorstore = FAISS.from_documents(documents, self.embeddings)
            self.stats['book_chunks'] = len(documents)
            self.metadata['total_documents'] = len(documents)
            self.metadata['book_processed'] = True
            self.metadata['book_info'] = {
                'title': 'Core Problem-Solving Knowledge',
                'type': 'built_in',
                'chunks': len(documents)
            }
            return True
        return False
    
    def load_existing_knowledge(self):
        """Load existing knowledge base from persistent storage"""
        try:
            # Load metadata
            if self.metadata_path.exists():
                with open(self.metadata_path, 'r') as f:
                    self.metadata = json.load(f)
            
            # Load stats
            if self.stats_path.exists():
                with open(self.stats_path, 'r') as f:
                    self.stats = json.load(f)
            
            # Load conversations
            if self.conversations_path.exists():
                with open(self.conversations_path, 'r') as f:
                    self.conversations = json.load(f)
            
            # Load vectorstore
            if self.vectorstore_path.exists() and self.embeddings:
                self.vectorstore = FAISS.load_local(
                    str(self.vectorstore_path), 
                    self.embeddings,
                    allow_dangerous_deserialization=True
                )
                return True
        except Exception as e:
            return False
        return False
    
    def save_all_data(self):
        """Save all knowledge base data to persistent storage"""
        try:
            # Save metadata
            self.metadata['last_updated'] = datetime.now().isoformat()
            with open(self.metadata_path, 'w') as f:
                json.dump(self.metadata, f, indent=2)
            
            # Save stats
            with open(self.stats_path, 'w') as f:
                json.dump(self.stats, f, indent=2)
            
            # Save conversations
            with open(self.conversations_path, 'w') as f:
                json.dump(self.conversations, f, indent=2)
            
            # Save vectorstore
            if self.vectorstore:
                self.vectorstore.save_local(str(self.vectorstore_path))
            
            return True
        except Exception as e:
            return False
    
    def fallback_initialization(self):
        """Fallback initialization if main process fails"""
        self.create_default_knowledge()
        self.metadata = {'fallback': True, 'created_at': datetime.now().isoformat()}
        self.stats = {'total_queries': 0, 'learning_sessions': 0, 'book_chunks': 0, 'conversation_chunks': 0, 'silly_questions_blocked': 0}
        self.conversations = []
    
    def process_book_content(self, text, page_texts, book_info):
        """Process book content and add to knowledge base"""
        try:
            # Text splitter
            text_splitter = RecursiveCharacterTextSplitter(
                chunk_size=800,
                chunk_overlap=150,
                length_function=len,
                separators=["\n\n", "\n", ".", "!", "?", ",", " ", ""]
            )
            
            # Create documents
            documents = []
            for page_info in page_texts:
                page_text = page_info['text']
                chapter_title = self._extract_chapter_title(page_text.split('\n'))
                
                doc = Document(
                    page_content=page_text,
                    metadata={
                        "source": "book",
                        "type": "book_content",
                        "page": page_info['page'],
                        "chapter": chapter_title,
                        "word_count": page_info['word_count'],
                        "book_title": book_info.get('title', 'Problem Solving Book'),
                        "processed_at": datetime.now().isoformat()
                    }
                )
                documents.append(doc)
            
            # Split into chunks
            chunks = text_splitter.split_documents(documents)
            
            # Add to vectorstore
            if self.vectorstore is None:
                self.vectorstore = FAISS.from_documents(chunks, self.embeddings)
            else:
                new_vectorstore = FAISS.from_documents(chunks, self.embeddings)
                self.vectorstore.merge_from(new_vectorstore)
            
            # Update metadata
            self.metadata['book_processed'] = True
            self.metadata['book_info'] = book_info
            self.metadata['total_documents'] += len(chunks)
            self.stats['book_chunks'] += len(chunks)
            
            return True, f"Successfully processed {len(chunks)} chunks from {len(page_texts)} pages!"
            
        except Exception as e:
            return False, f"Error processing book: {str(e)}"
    
    def add_conversation_to_knowledge(self, question, answer):
        """Add conversation to persistent knowledge base (auto-save)"""
        if len(question.strip()) < 10 or len(answer.strip()) < 20:
            return False
        
        try:
            conversation_text = f"Question: {question}\n\nAnswer: {answer}"
            
            doc = Document(
                page_content=conversation_text,
                metadata={
                    "source": "learned_conversation",
                    "type": "qa_pair",
                    "question": question,
                    "answer_preview": answer[:200] + "..." if len(answer) > 200 else answer,
                    "conversation_id": hashlib.md5(conversation_text.encode()).hexdigest()[:8],
                    "added_at": datetime.now().isoformat(),
                    "quality_score": self._calculate_quality_score(question, answer)
                }
            )
            
            # Add to vectorstore
            if self.vectorstore is None:
                self.vectorstore = FAISS.from_documents([doc], self.embeddings)
            else:
                new_vectorstore = FAISS.from_documents([doc], self.embeddings)
                self.vectorstore.merge_from(new_vectorstore)
            
            # Store conversation
            self.conversations.append({
                'question': question,
                'answer': answer,
                'timestamp': datetime.now().isoformat(),
                'learned': True
            })
            
            # Update stats
            self.stats['conversation_chunks'] += 1
            self.stats['learning_sessions'] += 1
            self.metadata['total_documents'] += 1
            
            # Auto-save to persistent storage
            self.save_all_data()
            
            return True
            
        except Exception as e:
            return False
    
    def search_knowledge_base(self, query, k=5):
        """Search the persistent knowledge base"""
        if self.vectorstore is None:
            return []
        
        try:
            self.stats['total_queries'] += 1
            docs = self.vectorstore.similarity_search_with_score(query, k=k)
            
            results = []
            for doc, score in docs:
                result = {
                    'content': doc.page_content,
                    'source': doc.metadata.get('source', 'unknown'),
                    'type': doc.metadata.get('type', 'unknown'),
                    'page': doc.metadata.get('page', 'N/A'),
                    'chapter': doc.metadata.get('chapter', 'Unknown Section'),
                    'similarity_score': float(score),
                    'metadata': doc.metadata
                }
                results.append(result)
            
            return results
            
        except Exception as e:
            return []
    
    def _extract_chapter_title(self, lines):
        """Extract chapter title from text lines"""
        for line in lines[:10]:
            line = line.strip()
            if line and len(line) < 100:
                if re.match(r'(chapter|section|part|unit)\s+\d+', line.lower()):
                    return line
                if line.isupper() or line.istitle():
                    return line
        return "General Content"
    
    def _calculate_quality_score(self, question, answer):
        """Calculate conversation quality score"""
        score = 0
        
        # Question quality
        if len(question.split()) >= 5: score += 1
        if any(word in question.lower() for word in ['how', 'what', 'why', 'strategy', 'problem']): score += 1
        if '?' in question: score += 1
        
        # Answer quality  
        if len(answer.split()) >= 20: score += 1
        if any(word in answer.lower() for word in ['approach', 'solution', 'method', 'step']): score += 1
        
        return min(score, 5)

class AshokMinimalChatbot:
    def __init__(self):
        self.knowledge_base = PersistentHFKnowledgeBase()
    
    def is_silly_question(self, question):
        """Detect silly or irrelevant questions"""
        question_lower = question.lower().strip()
        
        if len(question_lower) < 3:
            return True
        
        # Greeting patterns
        greeting_patterns = [
            r'\b(hello|hi|hey|salam|namaste|adab)\b',
            r'\b(good morning|evening|afternoon)\b',
            r'\b(how are you|kaise ho|kya hal)\b',
            r'\b(what.*your name|who are you)\b'
        ]
        
        # Silly keywords
        silly_keywords = [
            'stupid', 'dumb', 'joke', 'funny', 'lol', 'weather', 'movie',
            'song', 'game', 'gossip', 'love', 'dating', 'facebook',
            'instagram', 'politics', 'religion', 'age', 'appearance'
        ]
        
        # Problem-solving keywords
        good_keywords = [
            'problem', 'solve', 'solution', 'strategy', 'approach',
            'method', 'challenge', 'decision', 'plan', 'analyze',
            'conflict', 'team', 'work', 'project', 'manage'
        ]
        
        # Check patterns
        for pattern in greeting_patterns:
            if re.search(pattern, question_lower):
                return True
        
        # Score keywords
        good_score = sum(1 for keyword in good_keywords if keyword in question_lower)
        if good_score >= 2:
            return False
        
        silly_score = sum(1 for keyword in silly_keywords if keyword in question_lower)
        if silly_score >= 1:
            return True
        
        # Check structure
        question_words = ['what', 'how', 'why', 'when', 'where', 'can', 'should']
        has_question_word = any(word in question_lower.split() for word in question_words)
        word_count = len(question_lower.split())
        
        if word_count < 4 and not has_question_word:
            return True
        
        return False
    
    def generate_response(self, question, api_key):
        """Generate response using Gemini"""
        try:
            # Check for silly questions
            if self.is_silly_question(question):
                self.knowledge_base.stats['silly_questions_blocked'] += 1
                
                silly_responses = [
                    "Don't waste your time. Ask me something related to problem solving yaar!",
                    "I'm here to help with problem solving, not for chit-chat. Be serious!",
                    "Focus on real problems that need solving, samjha? Ask about strategies and approaches!",
                    "This is a problem-solving platform. Ask me about challenges, decisions, ya conflict resolution!",
                    "Tumhara dimagh kahan hai? Ask meaningful questions about problem-solving techniques."
                ]
                
                import random
                return random.choice(silly_responses), False
            
            # Configure Gemini
            genai.configure(api_key=api_key)
            model = genai.GenerativeModel('gemini-2.0-flash')
            
            # Search knowledge base
            relevant_results = self.knowledge_base.search_knowledge_base(question, k=5)
            
            # Build context
            context = ""
            references = []
            
            if relevant_results:
                context = "=== RELEVANT KNOWLEDGE ===\n\n"
                for i, result in enumerate(relevant_results, 1):
                    source_type = result['type']
                    if source_type == 'book_content':
                        context += f"**Reference {i}** (Chapter: {result['chapter']}, Page: {result['page']}):\n"
                    elif source_type == 'qa_pair':
                        context += f"**Learning {i}** (From past conversations):\n"
                    else:
                        context += f"**Framework {i}** (Core knowledge):\n"
                    
                    context += f"{result['content']}\n\n"
                    references.append(result)
            
            # Create prompt
            prompt = f"""
            You are Ashok, a problem-solving expert with Pakistani/Indian conversational style.
            
            Your characteristics:
            1. Mix English with Urdu naturally: "acha", "bilkul", "samjha", "dekho"
            2. Enthusiastic responses: "Excellent question bache!" or "Bahut acha sawal!"
            3. Reference knowledge sources when available
            4. Provide practical, actionable advice
            5. Encouraging and professional tone
            
            {context}
            
            User Question: {question}
            
            Provide a comprehensive, practical response using your characteristic style.
            Reference the knowledge sources when relevant and give actionable steps.
            """
            
            response = model.generate_content(prompt)
            final_response = response.text
            
            # Add clean references section
            if references:
                final_response += "\n\n**Knowledge Sources:**\n"
                for ref in references:
                    if ref['type'] == 'book_content':
                        final_response += f"• Book: {ref['chapter']} (Page {ref['page']})\n"
                    elif ref['type'] == 'qa_pair':
                        final_response += f"• Previous Learning\n"
                    else:
                        final_response += f"• Core Framework: {ref['metadata'].get('topic', 'Problem-solving')}\n"
            
            return final_response, True
            
        except Exception as e:
            return f"Sorry bache, I encountered an error: {str(e)}. Please check your API key and try again!", False

def show_typing_indicator():
    """Show typing indicator"""
    st.markdown("""
    <div class="typing-indicator">
        <span>Ashok is thinking</span>
        <div class="typing-dots">
            <div class="typing-dot"></div>
            <div class="typing-dot"></div>
            <div class="typing-dot"></div>
        </div>
    </div>
    """, unsafe_allow_html=True)

def show_api_setup():
    """Show API key setup interface"""
    st.markdown("""
    <div class="api-setup-container">
        <div class="api-setup-title">Welcome to Ashok 2.0</div>
        <div class="api-setup-subtitle">To get started, please enter your Gemini API key</div>
    </div>
    """, unsafe_allow_html=True)
    
    # API key input
    api_key = st.text_input(
        "Gemini API Key",
        type="password",
        placeholder="Enter your API key here...",
        label_visibility="collapsed"
    )
    
    # API guide
    st.markdown("""
    <div class="api-guide">
        <div class="api-guide-title">How to get your free API key:</div>
        <div class="api-guide-text">1. Visit <a href="https://makersuite.google.com/app/apikey" target="_blank" class="api-link">Google AI Studio</a></div>
        <div class="api-guide-text">2. Sign in with your Google account</div>
        <div class="api-guide-text">3. Click "Create API Key"</div>
        <div class="api-guide-text">4. Copy and paste the key above</div>
    </div>
    """, unsafe_allow_html=True)
    
    return api_key

def show_quick_actions():
    """Show enhanced quick action buttons"""
    st.markdown("""
    <div class="quick-actions-container">
        <div class="quick-actions-title">Quick Start Questions</div>
        <div class="quick-action-grid">
    """, unsafe_allow_html=True)
    
    # Quick action data
    actions = [
        {
            "icon": "🎯",
            "title": "Task Prioritization",
            "desc": "Learn how to prioritize when everything seems urgent",
            "question": "How do I prioritize tasks when everything seems urgent and important?"
        },
        {
            "icon": "🤝",
            "title": "Conflict Resolution",
            "desc": "Effective strategies for resolving team conflicts",
            "question": "What's the best approach to resolve conflicts in my team?"
        },
        {
            "icon": "⚡",
            "title": "Decision Making",
            "desc": "Improve your decision-making process",
            "question": "How can I improve my decision-making process for complex problems?"
        },
        {
            "icon": "🎪",
            "title": "Difficult People",
            "desc": "Handle challenging stakeholders professionally",
            "question": "How do I deal with difficult stakeholders effectively?"
        },
        {
            "icon": "🔄",
            "title": "Change Management",
            "desc": "Navigate organizational changes smoothly",
            "question": "How can I help my team adapt to organizational changes?"
        },
        {
            "icon": "💡",
            "title": "Creative Solutions",
            "desc": "Generate innovative solutions to problems",
            "question": "What techniques can I use to think more creatively about problems?"
        }
    ]
    
    # Create columns for actions
    cols = st.columns(2)
    for i, action in enumerate(actions):
        with cols[i % 2]:
            if st.button(
                f"{action['icon']} {action['title']}",
                key=f"action_{i}",
                help=action['desc'],
                use_container_width=True
            ):
                st.session_state.auto_question = action['question']
                st.rerun()
    
    st.markdown("</div></div>", unsafe_allow_html=True)

def test_api_key(api_key):
    """Test if API key is valid"""
    try:
        genai.configure(api_key=api_key)
        model = genai.GenerativeModel('gemini-2.0-flash')
        test_response = model.generate_content("Hello")
        return True
    except Exception as e:
        return False

# Global instance for the app
@st.cache_resource
def get_chatbot():
    """Get cached chatbot instance"""
    return AshokMinimalChatbot()

def main():
    # Get cached chatbot instance
    chatbot = get_chatbot()
    
    # Initialize session state
    if 'messages' not in st.session_state:
        st.session_state.messages = []
    
    if 'auto_question' not in st.session_state:
        st.session_state.auto_question = None
    
    if 'api_key_valid' not in st.session_state:
        st.session_state.api_key_valid = False
    
    # Main title
    st.markdown('<h1 class="main-title">ASHOK 2.0</h1>', unsafe_allow_html=True)
    
    # API Key Setup Phase
    if not st.session_state.api_key_valid:
        api_key = show_api_setup()
        
        if api_key:
            if test_api_key(api_key):
                st.session_state.api_key = api_key
                st.session_state.api_key_valid = True
                st.markdown("""
                <div class="status-success">
                    API key configured successfully! Ready to chat.
                </div>
                """, unsafe_allow_html=True)
                time.sleep(1)
                st.rerun()
            else:
                st.markdown("""
                <div class="status-warning">
                    Invalid API key. Please check and try again.
                </div>
                """, unsafe_allow_html=True)
    
    # Main Chat Interface Phase
    else:
        api_key = st.session_state.api_key
        
        # Handle auto questions from quick actions
        if st.session_state.auto_question:
            # Add user message
            st.session_state.messages.append({"role": "user", "content": st.session_state.auto_question})
            
            # Generate response
            with st.spinner("Processing your question..."):
                response, is_helpful = chatbot.generate_response(st.session_state.auto_question, api_key)
                st.session_state.messages.append({"role": "assistant", "content": response})
                
                # Auto-learn from helpful conversations
                if is_helpful and not chatbot.is_silly_question(st.session_state.auto_question):
                    chatbot.knowledge_base.add_conversation_to_knowledge(st.session_state.auto_question, response)
            
            # Clear auto question
            st.session_state.auto_question = None
            st.rerun()
        
        # Display chat history
        if st.session_state.messages:
            st.markdown('<div class="chat-container">', unsafe_allow_html=True)
            for message in st.session_state.messages:
                with st.chat_message(message["role"]):
                    st.markdown(message["content"])
            st.markdown('</div>', unsafe_allow_html=True)
        
        # Chat input
        if prompt := st.chat_input("Ask about problem-solving strategies...", key="main_chat"):
            # Add user message
            st.session_state.messages.append({"role": "user", "content": prompt})
            with st.chat_message("user"):
                st.markdown(prompt)
            
            # Generate response with typing indicator
            with st.chat_message("assistant"):
                # Show typing indicator
                typing_placeholder = st.empty()
                with typing_placeholder:
                    show_typing_indicator()
                
                # Simulate typing delay
                time.sleep(1)
                
                # Clear typing indicator and show response
                typing_placeholder.empty()
                
                with st.spinner("Processing..."):
                    response, is_helpful = chatbot.generate_response(prompt, api_key)
                    st.markdown(response)
                    
                    # Add to chat history
                    st.session_state.messages.append({"role": "assistant", "content": response})
                    
                    # Auto-learn from helpful conversations
                    if is_helpful and not chatbot.is_silly_question(prompt):
                        learned = chatbot.knowledge_base.add_conversation_to_knowledge(prompt, response)
                        if learned:
                            st.markdown(
                                '<div class="learning-indicator">Knowledge updated - This conversation has been learned!</div>',
                                unsafe_allow_html=True
                            )
        
        # Show quick actions if no messages yet
        if not st.session_state.messages:
            show_quick_actions()
        
        # Control panel
        if st.session_state.messages:
            st.markdown("""
            <div class="control-panel">
            """, unsafe_allow_html=True)
            
            if st.button("Clear Chat", key="clear_chat"):
                st.session_state.messages = []
                st.rerun()
            
            st.markdown("</div>", unsafe_allow_html=True)

if __name__ == "__main__":
    main()