import os
import gradio as gr
import PyPDF2
import docx2txt
import logging
import json
from typing import List, Dict, Tuple
import uuid
from datetime import datetime

import nltk
from nltk.corpus import stopwords
from nltk.tokenize import word_tokenize

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

# Try to import Pinecone - fallback to in-memory if not available
PINECONE_AVAILABLE = False
try:
    from pinecone import Pinecone, ServerlessSpec
    PINECONE_AVAILABLE = True
    print("✅ Pinecone library loaded successfully")
except ImportError:
    print("⚠️ Pinecone not available - using in-memory storage")

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

# Download NLTK data at startup
try:
    nltk.download('punkt', quiet=True)
    nltk.download('stopwords', quiet=True) 
    nltk.download('punkt_tab', quiet=True)
except:
    pass

# ----------------------------------------------------------------------------
# 1) Check for Pinecone API Key
# ----------------------------------------------------------------------------

def check_pinecone_setup():
    """Check if Pinecone is properly configured."""
    api_key = os.getenv("PINECONE_API_KEY")
    
    if not api_key:
        return False, "No PINECONE_API_KEY found in environment variables"
    
    if not PINECONE_AVAILABLE:
        return False, "Pinecone library not installed"
    
    try:
        # Test connection
        pc = Pinecone(api_key=api_key)
        pc.list_indexes()  # Simple test call
        return True, "Pinecone connection successful"
    except Exception as e:
        return False, f"Pinecone connection failed: {str(e)}"

# ----------------------------------------------------------------------------
# 2) Pinecone Database Class
# ----------------------------------------------------------------------------

class PineconeResumeDB:
    def __init__(self, api_key: str):
        """Initialize Pinecone connection and index."""
        self.pc = Pinecone(api_key=api_key)
        self.index_name = "resume-vectors"
        self.dimension = 1000  # TF-IDF vector dimension
        
        # Create index if it doesn't exist
        existing_indexes = [index.name for index in self.pc.list_indexes()]
        if self.index_name not in existing_indexes:
            print(f"Creating new Pinecone index: {self.index_name}")
            self.pc.create_index(
                name=self.index_name,
                dimension=self.dimension,
                metric="cosine",
                spec=ServerlessSpec(
                    cloud="aws",
                    region="us-east-1"
                )
            )
        
        self.index = self.pc.Index(self.index_name)
        self.vectorizer = TfidfVectorizer(max_features=self.dimension)
        self.is_fitted = False
        
    def add_resume(self, resume_id: str, text_content: str, metadata: Dict):
        """Store a resume vector in Pinecone with metadata."""
        try:
            processed_text = preprocess_text(text_content)
            
            # For demo purposes, we'll use a simple approach
            # In production, you'd want a more sophisticated vectorization strategy
            if not self.is_fitted:
                # Fit on first document (not ideal, but works for demo)
                vector = self.vectorizer.fit_transform([processed_text]).toarray()[0]
                self.is_fitted = True
            else:
                # For subsequent documents, we'll refit (not ideal for production)
                # In production, you'd maintain a consistent vocabulary
                vector = self.vectorizer.fit_transform([processed_text]).toarray()[0]
            
            # Prepare metadata
            full_metadata = {
                **metadata,
                "text_preview": text_content[:500],
                "upload_date": datetime.now().isoformat(),
                "vector_model": "tfidf"
            }
            
            # Upsert to Pinecone
            self.index.upsert(
                vectors=[{
                    "id": resume_id,
                    "values": vector.tolist(),
                    "metadata": full_metadata
                }]
            )
            
            logging.info(f"Successfully stored resume {resume_id} in Pinecone")
            return True
            
        except Exception as e:
            logging.error(f"Error storing resume in Pinecone: {e}")
            return False
    
    def search_resumes(self, job_description: str, top_k: int = 10) -> List[Tuple[str, float, Dict]]:
        """Search for most similar resumes to a job description."""
        try:
            processed_jd = preprocess_text(job_description)
            
            if not self.is_fitted:
                return []
                
            jd_vector = self.vectorizer.transform([processed_jd]).toarray()[0]
            
            # Search Pinecone
            results = self.index.query(
                vector=jd_vector.tolist(),
                top_k=top_k,
                include_metadata=True,
                include_values=False
            )
            
            # Format results
            formatted_results = []
            for match in results.matches:
                formatted_results.append((
                    match.id,
                    match.score,
                    match.metadata
                ))
            
            return formatted_results
            
        except Exception as e:
            logging.error(f"Error searching Pinecone: {e}")
            return []
    
    def get_resume_count(self) -> int:
        """Get total number of resumes in the database."""
        try:
            stats = self.index.describe_index_stats()
            return stats.total_vector_count
        except:
            return 0
    
    def clear_database(self):
        """Clear all vectors from the database."""
        try:
            # Delete all vectors (be careful with this in production!)
            self.index.delete(delete_all=True)
            return True
        except Exception as e:
            logging.error(f"Error clearing database: {e}")
            return False

# ----------------------------------------------------------------------------
# 3) In-Memory Fallback Database
# ----------------------------------------------------------------------------

class InMemoryResumeDB:
    def __init__(self):
        """Initialize in-memory storage for fallback."""
        self.resumes = {}
        self.vectorizer = None
        self.resume_vectors = {}
        
    def add_resume(self, resume_id: str, text_content: str, metadata: Dict):
        """Store a resume in memory."""
        try:
            processed_text = preprocess_text(text_content)
            
            self.resumes[resume_id] = {
                'text': processed_text,
                'raw_text': text_content,
                'metadata': {
                    **metadata,
                    'upload_date': datetime.now().isoformat(),
                    'text_length': len(text_content)
                }
            }
            
            self._update_vectors()
            return True
            
        except Exception as e:
            logging.error(f"Error storing resume: {e}")
            return False
    
    def _update_vectors(self):
        """Update TF-IDF vectors for all resumes."""
        if not self.resumes:
            return
            
        try:
            texts = [resume['text'] for resume in self.resumes.values()]
            resume_ids = list(self.resumes.keys())
            
            self.vectorizer = TfidfVectorizer(max_features=1000)
            vectors = self.vectorizer.fit_transform(texts)
            
            for i, resume_id in enumerate(resume_ids):
                self.resume_vectors[resume_id] = vectors[i]
                
        except Exception as e:
            logging.error(f"Error updating vectors: {e}")
    
    def search_resumes(self, job_description: str, top_k: int = 10) -> List[Tuple[str, float, Dict]]:
        """Search for most similar resumes."""
        try:
            if not self.resumes or not self.vectorizer:
                return []
                
            processed_jd = preprocess_text(job_description)
            jd_vector = self.vectorizer.transform([processed_jd])
            
            results = []
            for resume_id, resume_vector in self.resume_vectors.items():
                similarity = cosine_similarity(jd_vector, resume_vector)[0][0]
                results.append((resume_id, similarity, self.resumes[resume_id]['metadata']))
            
            results.sort(key=lambda x: x[1], reverse=True)
            return results[:top_k]
            
        except Exception as e:
            logging.error(f"Error searching resumes: {e}")
            return []
    
    def get_resume_count(self) -> int:
        return len(self.resumes)
    
    def clear_database(self):
        self.resumes.clear()
        self.resume_vectors.clear()
        self.vectorizer = None
        return True

# ----------------------------------------------------------------------------
# 4) Initialize Database
# ----------------------------------------------------------------------------

# Check Pinecone setup and initialize appropriate database
pinecone_available, pinecone_status = check_pinecone_setup()
print(f"Pinecone status: {pinecone_status}")

if pinecone_available:
    api_key = os.getenv("PINECONE_API_KEY")
    resume_db = PineconeResumeDB(api_key=api_key)
    DB_TYPE = "Pinecone Vector Database"
else:
    resume_db = InMemoryResumeDB()
    DB_TYPE = "In-Memory Storage (Demo Mode)"

print(f"Using database type: {DB_TYPE}")

# ----------------------------------------------------------------------------
# 5) Utility Functions (same as before)
# ----------------------------------------------------------------------------

def extract_text_from_pdf(file_obj):
    """Extract all text from a PDF file object."""
    text_content = []
    try:
        if hasattr(file_obj, 'read'):
            pdf_reader = PyPDF2.PdfReader(file_obj)
        else:
            with open(file_obj, 'rb') as f:
                pdf_reader = PyPDF2.PdfReader(f)
        
        for page in pdf_reader.pages:
            page_text = page.extract_text()
            if page_text:
                text_content.append(page_text)
        return "\n".join(text_content)
    except Exception as e:
        logging.error(f"Error reading PDF: {e}")
        return f"Error reading PDF: {str(e)}"

def extract_text_from_docx(file_path):
    """Extract all text from a DOCX file."""
    try:
        return docx2txt.process(file_path)
    except Exception as e:
        logging.error(f"Error reading DOCX: {e}")
        return f"Error reading DOCX: {str(e)}"

def extract_text_from_txt(file_obj):
    """Extract all text from a TXT file."""
    try:
        if hasattr(file_obj, 'read'):
            content = file_obj.read()
            if isinstance(content, bytes):
                return content.decode("utf-8", errors="ignore")
            return content
        else:
            with open(file_obj, 'r', encoding='utf-8', errors='ignore') as f:
                return f.read()
    except Exception as e:
        logging.error(f"Error reading TXT: {e}")
        return f"Error reading TXT: {str(e)}"

def preprocess_text(text):
    """Clean and preprocess text."""
    try:
        text = str(text).lower()
        tokens = word_tokenize(text)
        stop_words = set(stopwords.words('english'))
        filtered_tokens = [t for t in tokens if t.isalpha() and t not in stop_words and len(t) > 2]
        return " ".join(filtered_tokens)
    except Exception as e:
        logging.error(f"Error preprocessing text: {e}")
        return str(text).lower()

# ----------------------------------------------------------------------------
# 6) Gradio Functions (same as before, but with database type info)
# ----------------------------------------------------------------------------

def add_resumes_to_db(cv_files, position_title="General"):
    """Add uploaded resumes to database."""
    if not cv_files:
        return "❌ No files uploaded."
    
    added_count = 0
    results = [f"🗄️ Using: {DB_TYPE}\n"]
    
    for uploaded_file in cv_files:
        try:
            filename = uploaded_file.name if hasattr(uploaded_file, 'name') else str(uploaded_file)
            filename = os.path.basename(filename)
            file_ext = os.path.splitext(filename)[1].lower()
            
            # Extract text based on file type
            if file_ext == ".pdf":
                file_content = extract_text_from_pdf(uploaded_file)
            elif file_ext == ".txt":
                file_content = extract_text_from_txt(uploaded_file)
            elif file_ext == ".docx":
                file_content = extract_text_from_docx(uploaded_file)
            else:
                results.append(f"❌ {filename}: Unsupported file type")
                continue
            
            if "Error reading" in file_content:
                results.append(f"❌ {filename}: {file_content}")
                continue
                
            if len(file_content.strip()) < 50:
                results.append(f"❌ {filename}: File appears to be empty or too short")
                continue
            
            # Generate unique ID
            resume_id = f"{filename}_{uuid.uuid4().hex[:8]}"
            
            # Metadata
            metadata = {
                "filename": filename,
                "file_type": file_ext,
                "position_applied": position_title,
            }
            
            # Store in database
            if resume_db.add_resume(resume_id, file_content, metadata):
                added_count += 1
                results.append(f"✅ {filename}: Successfully added to database")
            else:
                results.append(f"❌ {filename}: Failed to add to database")
                
        except Exception as e:
            results.append(f"❌ Error processing file: {str(e)}")
    
    summary = f"\n📊 Summary: {added_count} out of {len(cv_files)} resumes added successfully."
    summary += f"\n📁 Total resumes in database: {resume_db.get_resume_count()}"
    
    return "\n".join(results) + summary

def search_resumes_in_db(job_description, top_k=10):
    """Search for matching resumes."""
    if not job_description.strip():
        return [], "❌ Please enter a job description."
    
    try:
        results = resume_db.search_resumes(job_description, top_k=int(top_k))
        
        if not results:
            return [], f"❌ No matching resumes found. Add some resumes first! (Using: {DB_TYPE})"
        
        # Format for display
        display_data = []
        for resume_id, score, metadata in results:
            display_data.append([
                metadata.get('filename', 'Unknown'),
                f"{score:.4f}",
                metadata.get('position_applied', 'N/A'),
                metadata.get('upload_date', 'N/A')[:16] if metadata.get('upload_date') else 'N/A',
                f"{metadata.get('text_length', 0)} chars" if 'text_length' in metadata else "N/A"
            ])
        
        status = f"✅ Found {len(results)} matching resumes from {resume_db.get_resume_count()} total resumes. (Using: {DB_TYPE})"
        return display_data, status
        
    except Exception as e:
        logging.error(f"Search error: {e}")
        return [], f"❌ Search failed: {str(e)}"

def get_database_stats():
    """Get current database statistics."""
    count = resume_db.get_resume_count()
    return f"🗄️ Database Type: {DB_TYPE}\n📊 Contains {count} resumes\n🔗 Pinecone Status: {pinecone_status}"

def clear_database():
    """Clear the database."""
    if resume_db.clear_database():
        return f"🗑️ Database cleared successfully! (Using: {DB_TYPE})"
    else:
        return "❌ Failed to clear database"

# ----------------------------------------------------------------------------
# 7) Gradio Interface (same as before)
# ----------------------------------------------------------------------------

def create_gradio_interface():
    with gr.Blocks(title="AI Resume Ranking System", theme=gr.themes.Soft()) as demo:
        gr.Markdown(f"""
        # 🎯 AI Resume Ranking System
        
        **Upload resumes and find the best matches for your job openings!**
        
        Currently using: **{DB_TYPE}**
        
        This system uses TF-IDF vectorization and cosine similarity to rank resumes based on job descriptions.
        """)
        
        with gr.Tab("📤 Upload Resumes"):
            gr.Markdown("### Add New Resumes to Database")
            
            with gr.Row():
                with gr.Column():
                    position_input = gr.Textbox(
                        label="Position Title (Optional)", 
                        placeholder="e.g., Software Engineer, Data Scientist",
                        value="General"
                    )
                    
                    upload_files = gr.File(
                        label="Upload Resumes (PDF/DOCX/TXT)", 
                        file_count="multiple",
                        file_types=[".pdf", ".docx", ".txt"]
                    )
                    
                    upload_btn = gr.Button("📁 Add to Database", variant="primary", size="lg")
                
                with gr.Column():
                    upload_status = gr.Textbox(
                        label="Upload Status", 
                        lines=12,
                        placeholder="Upload results will appear here..."
                    )
            
            upload_btn.click(
                fn=add_resumes_to_db,
                inputs=[upload_files, position_input],
                outputs=[upload_status]
            )
        
        with gr.Tab("🔍 Search & Rank"):
            gr.Markdown("### Find Best Resume Matches")
            
            with gr.Row():
                with gr.Column():
                    job_description_input = gr.Textbox(
                        label="Job Description", 
                        placeholder="Enter the job requirements, skills needed, and role description...",
                        lines=8
                    )
                    
                    with gr.Row():
                        top_k_input = gr.Slider(
                            label="Number of Results", 
                            minimum=1, 
                            maximum=20, 
                            value=10, 
                            step=1
                        )
                        
                        search_btn = gr.Button("🔍 Search & Rank", variant="primary", size="lg")
            
            search_results = gr.Dataframe(
                headers=["Resume File", "Similarity Score", "Position", "Upload Date", "Size"],
                label="📊 Ranking Results",
                wrap=True
            )
            
            search_status = gr.Textbox(label="Search Status")
            
            search_btn.click(
                fn=search_resumes_in_db,
                inputs=[job_description_input, top_k_input],
                outputs=[search_results, search_status]
            )
        
        with gr.Tab("📊 Database Info"):
            gr.Markdown("### Database Management")
            
            with gr.Row():
                stats_btn = gr.Button("📊 Refresh Stats", variant="secondary")
                clear_btn = gr.Button("🗑️ Clear Database", variant="stop")
            
            with gr.Row():
                stats_display = gr.Textbox(label="Database Statistics", lines=5)
                clear_status = gr.Textbox(label="Clear Status", lines=3)
            
            stats_btn.click(fn=get_database_stats, outputs=[stats_display])
            clear_btn.click(fn=clear_database, outputs=[clear_status])
            
            # Auto-load stats on page load
            demo.load(fn=get_database_stats, outputs=[stats_display])
        
        with gr.Tab("🔑 Setup"):
            gr.Markdown(f"""
            ### Current Configuration
            
            **Database Type**: {DB_TYPE}  
            **Pinecone Status**: {pinecone_status}
            
            ### To Enable Pinecone (Optional)
            
            If you want persistent storage with Pinecone:
            
            1. **Get API Key**: Sign up at [pinecone.io](https://pinecone.io) and get your API key
            2. **Set Environment Variable**: Add `PINECONE_API_KEY` to your Hugging Face Space secrets
            3. **Restart Space**: The app will automatically detect and use Pinecone
            
            ### Hugging Face Spaces Setup
            
            1. Go to your Space settings
            2. Click on "Variables and secrets"
            3. Add a new secret:
               - Name: `PINECONE_API_KEY`
               - Value: Your actual Pinecone API key
            4. Save and restart your Space
            
            ### Benefits of Pinecone
            
            - ✅ Persistent storage (data survives restarts)
            - ✅ Scalable to millions of resumes  
            - ✅ Advanced vector search capabilities
            - ✅ Real-time updates and deletions
            """)
        
        with gr.Tab("ℹ️ How It Works"):
            gr.Markdown("""
            ### How This System Works
            
            1. **Upload Resumes**: Add PDF, DOCX, or TXT resume files to the database
            
            2. **Text Processing**: 
               - Extracts text from documents
               - Removes common words (stopwords)
               - Converts to lowercase and tokenizes
            
            3. **Vectorization**: 
               - Uses TF-IDF (Term Frequency-Inverse Document Frequency)
               - Creates numerical vectors representing document content
               - Emphasizes rare but relevant terms
            
            4. **Similarity Scoring**:
               - Calculates cosine similarity between job description and resumes
               - Ranks resumes by similarity score (0.0 to 1.0)
               - Higher scores indicate better matches
            
            ### Database Options
            
            - **In-Memory**: Fast, but data is lost when session ends
            - **Pinecone**: Persistent, scalable vector database (requires API key)
            
            ### Tips for Best Results
            
            - **Job Descriptions**: Be specific about required skills, technologies, and experience
            - **Resume Quality**: Well-formatted resumes with clear text work best
            - **Keywords**: Include industry-specific terms and technical skills
            - **File Formats**: PDF and DOCX usually work better than scanned images
            """)
    
    return demo

# ----------------------------------------------------------------------------
# 8) Launch Application
# ----------------------------------------------------------------------------

if __name__ == "__main__":
    app = create_gradio_interface()
    app.launch()