#!/usr/bin/env python3
"""
FAISS Index Building Script for Julien Serbanescu RAG System

This script processes various document types and creates a FAISS vector index:
- Web content (portfolio website)
- Repository READMEs
- PDF documents
- Audio transcriptions

The script creates:
- index.faiss: FAISS vector index
- index.pkl: LangChain docstore and index mapping
- metadata.pkl: Separate metadata list
- document_lookup.txt: Human-readable document index
"""

import os
import sys
import pickle
import faiss
import numpy as np
import cohere
from dotenv import load_dotenv
from langchain_community.docstore.document import Document
from langchain_community.docstore.in_memory import InMemoryDocstore
from langchain_text_splitters import RecursiveCharacterTextSplitter
import requests
from bs4 import BeautifulSoup
import traceback

# Load environment variables
load_dotenv()
cohere_api_key = os.getenv("COHEREAPIKEY")

if not cohere_api_key:
    raise ValueError("COHEREAPIKEY not found in environment variables")

# Initialize Cohere client
co = cohere.Client(cohere_api_key)

class CohereEmbeddingsForIndexing:
    """Custom Cohere embeddings class for indexing documents"""
    def __init__(self, client):
        self.client = client
        self.embed_dim = self._get_embed_dim()

    def _get_embed_dim(self):
        try:
            response = self.client.embed(
                texts=["test"], model="embed-english-v3.0", input_type="search_document"
            )
            return len(response.embeddings[0])
        except Exception as e:
            print(f"Warning: Could not determine embedding dimension automatically: {e}. Defaulting to 4096.")
            return 4096

    def embed_documents(self, texts):
        """Embed a list of documents"""
        try:
            # Process in batches to avoid API limits
            batch_size = 96  # Cohere's limit
            all_embeddings = []
            
            for i in range(0, len(texts), batch_size):
                batch = texts[i:i + batch_size]
                response = self.client.embed(
                    texts=batch,
                    model="embed-english-v3.0",
                    input_type="search_document"
                )
                if hasattr(response, 'embeddings') and len(response.embeddings) > 0:
                    all_embeddings.extend(response.embeddings)
                else:
                    print(f"Warning: No embeddings found for batch {i//batch_size + 1}")
                    # Add zero vectors as fallback
                    all_embeddings.extend([np.zeros(self.embed_dim).tolist() for _ in batch])
            
            return [np.array(emb).astype('float32') for emb in all_embeddings]
        except Exception as e:
            print(f"Error embedding documents: {e}")
            # Return zero vectors as fallback
            return [np.zeros(self.embed_dim, dtype=np.float32) for _ in texts]

def scrape_linkedin_profile(profile_url):
    """Scrape LinkedIn profile (public content only)
    
    Note: LinkedIn has strict anti-scraping measures. This may not work reliably
    without authentication. For best results, consider using LinkedIn's official API
    or manually exporting your profile data.
    """
    try:
        # LinkedIn requires proper headers and may block simple scrapers
        headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
            'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8',
            'Accept-Language': 'en-US,en;q=0.5',
            'Accept-Encoding': 'gzip, deflate',
            'Connection': 'keep-alive',
        }
        
        response = requests.get(profile_url, timeout=15, headers=headers, allow_redirects=True)
        
        # LinkedIn often redirects to login page for unauthenticated requests
        if 'login' in response.url.lower() or 'authwall' in response.url.lower():
            print(f"   Warning: LinkedIn requires authentication. Consider exporting your profile manually.")
            return ""
        
        response.raise_for_status()
        soup = BeautifulSoup(response.content, 'html.parser')
        
        # Remove script and style elements
        for script in soup(["script", "style", "nav", "footer", "header"]):
            script.decompose()
        
        # Try to extract profile content
        # LinkedIn structure may vary, try common selectors
        content_parts = []
        
        # Try to find main profile sections
        selectors = [
            'main', 
            '.core-rail',
            '.profile-section',
            '.pv-profile-section',
            '.ph5',
            'article'
        ]
        
        for selector in selectors:
            elements = soup.select(selector)
            if elements:
                content_parts.extend([elem.get_text() for elem in elements])
                break
        
        # If no specific sections found, get all text
        if not content_parts:
            content_parts = [soup.get_text()]
        
        # Clean up whitespace
        full_text = ' '.join(content_parts)
        lines = (line.strip() for line in full_text.splitlines())
        chunks = (phrase.strip() for line in lines for phrase in line.split("  "))
        cleaned_text = ' '.join(chunk for chunk in chunks if chunk)
        
        return cleaned_text if cleaned_text else ""
        
    except requests.exceptions.RequestException as e:
        print(f"   Error accessing LinkedIn: {e}")
        print(f"   Note: LinkedIn may block automated access. Consider:")
        print(f"   1. Using LinkedIn's official API")
        print(f"   2. Manually exporting your profile data")
        print(f"   3. Using a browser automation tool like Selenium (more complex)")
        return ""
    except Exception as e:
        print(f"   Error scraping LinkedIn profile: {e}")
        return ""

def scrape_website(url, max_pages=10):
    """Scrape content from a website, including linked pages"""
    try:
        scraped_urls = set()
        all_content = []
        
        def scrape_page(page_url):
            """Scrape a single page"""
            if page_url in scraped_urls or len(scraped_urls) >= max_pages:
                return ""
            
            try:
                response = requests.get(page_url, timeout=10, headers={
                    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36'
                })
                response.raise_for_status()
                soup = BeautifulSoup(response.content, 'html.parser')
                
                scraped_urls.add(page_url)
                
                # Remove script and style elements
                for script in soup(["script", "style", "nav", "footer", "header"]):
                    script.decompose()
                
                # Extract main content areas (common portfolio patterns)
                main_content = ""
                
                # Try to find main content containers
                for selector in ['main', 'article', '.content', '#content', '.main-content', '.portfolio-content']:
                    elements = soup.select(selector)
                    if elements:
                        main_content += "\n".join([elem.get_text() for elem in elements])
                        break
                
                # If no main content found, get all text
                if not main_content:
                    main_content = soup.get_text()
                
                # Clean up whitespace
                lines = (line.strip() for line in main_content.splitlines())
                chunks = (phrase.strip() for line in lines for phrase in line.split("  "))
                cleaned_text = ' '.join(chunk for chunk in chunks if chunk)
                
                # Extract links for additional pages (same domain only)
                if len(scraped_urls) < max_pages:
                    base_url = '/'.join(page_url.split('/')[:3])
                    links = soup.find_all('a', href=True)
                    for link in links[:5]:  # Limit to first 5 links per page
                        href = link['href']
                        if href.startswith('/'):
                            full_url = base_url + href
                        elif href.startswith('http') and base_url in href:
                            full_url = href
                        else:
                            continue
                        
                        if full_url not in scraped_urls:
                            additional_content = scrape_page(full_url)
                            if additional_content:
                                cleaned_text += f"\n\n--- Content from {full_url} ---\n{additional_content}"
                
                return cleaned_text
            except Exception as e:
                print(f"Error scraping page {page_url}: {e}")
                return ""
        
        # Start scraping from the main URL
        content = scrape_page(url)
        if content:
            all_content.append(content)
        
        return "\n\n".join(all_content) if all_content else ""
    except Exception as e:
        print(f"Error scraping {url}: {e}")
        return ""

def load_pdf_documents(pdf_dir):
    """Load and process PDF documents"""
    documents = []
    
    if not os.path.exists(pdf_dir):
        print(f"PDF directory {pdf_dir} not found, skipping...")
        return documents
    
    pdf_files = [f for f in os.listdir(pdf_dir) if f.endswith('.pdf')]
    
    for pdf_file in pdf_files:
        try:
            pdf_path = os.path.join(pdf_dir, pdf_file)
            print(f"Processing PDF: {pdf_file}")
            
            # Try to extract text using PyPDF2 if available, otherwise use placeholder
            try:
                import PyPDF2
                with open(pdf_path, 'rb') as f:
                    pdf_reader = PyPDF2.PdfReader(f)
                    content = ""
                    for page_num, page in enumerate(pdf_reader.pages):
                        page_text = page.extract_text()
                        if page_text:
                            content += f"\n--- Page {page_num + 1} ---\n{page_text}\n"
                    
                    if not content.strip():
                        content = f"PDF content from {pdf_file} (text extraction failed)"
            except ImportError:
                print(f"PyPDF2 not available, using placeholder for {pdf_file}")
                content = f"PDF content from {pdf_file} (install PyPDF2 for text extraction)"
            except Exception as e:
                print(f"Error extracting text from {pdf_file}: {e}")
                content = f"PDF content from {pdf_file} (extraction error: {str(e)})"
            
            # Create document
            doc = Document(
                page_content=content,
                metadata={
                    'source': f"docs/pdfs/{pdf_file}",
                    'type': 'pdf'
                }
            )
            documents.append(doc)
                
        except Exception as e:
            print(f"Error processing PDF {pdf_file}: {e}")
            continue
    
    return documents

def load_audio_transcripts(transcript_dir):
    """Load audio transcript files"""
    documents = []
    
    if not os.path.exists(transcript_dir):
        print(f"Transcript directory {transcript_dir} not found, skipping...")
        return documents
    
    transcript_files = [f for f in os.listdir(transcript_dir) if f.endswith('_transcript.txt')]
    
    for transcript_file in transcript_files:
        try:
            transcript_path = os.path.join(transcript_dir, transcript_file)
            print(f"Processing transcript: {transcript_file}")
            
            with open(transcript_path, 'r', encoding='utf-8') as f:
                content = f.read()
            
            # Create document
            doc = Document(
                page_content=content,
                metadata={
                    'source': f"docs/youtube/{transcript_file}",
                    'type': 'audio_transcription',
                    'original_file': transcript_file.replace('_transcript.txt', '.m4a')
                }
            )
            documents.append(doc)
            
        except Exception as e:
            print(f"Error processing transcript {transcript_file}: {e}")
            continue
    
    return documents

def load_readme_documents(readme_dir):
    """Load README documents from various repositories, filtering out Tpoze-subnet content"""
    documents = []
    
    if not os.path.exists(readme_dir):
        print(f"README directory {readme_dir} not found, skipping...")
        return documents
    
    readme_files = [f for f in os.listdir(readme_dir) if f.endswith('.md')]
    
    for readme_file in readme_files:
        try:
            readme_path = os.path.join(readme_dir, readme_file)
            print(f"Processing README: {readme_file}")
            
            with open(readme_path, 'r', encoding='utf-8') as f:
                content = f.read()
            
            # Filter out content containing "Tpoze-subnet" (case-insensitive)
            if "tpoze-subnet" in content.lower():
                print(f"   Skipping {readme_file} - contains Tpoze-subnet content")
                continue
            
            # Create document
            doc = Document(
                page_content=content,
                metadata={
                    'source': f"docs/readmes/{readme_file}",
                    'type': 'repo'
                }
            )
            documents.append(doc)
            
        except Exception as e:
            print(f"Error processing README {readme_file}: {e}")
            continue
    
    return documents

def load_github_activity(activity_dir):
    """Load GitHub issues and PRs downloaded by download_github_activity.py"""
    documents = []
    if not os.path.exists(activity_dir):
        print(f"  GitHub activity directory {activity_dir} not found, skipping...")
        return documents

    md_files = [f for f in os.listdir(activity_dir) if f.endswith('.md')]
    for md_file in md_files:
        try:
            filepath = os.path.join(activity_dir, md_file)
            with open(filepath, 'r', encoding='utf-8') as f:
                content = f.read()

            doc_type = 'github_pr' if '_PR_' in md_file else 'github_issue'
            doc = Document(
                page_content=content,
                metadata={
                    'source': f"docs/github_activity/{md_file}",
                    'type': doc_type,
                }
            )
            documents.append(doc)
        except Exception as e:
            print(f"  Error loading {md_file}: {e}")
    return documents


def load_publications(papers_dir):
    """Load publication documents written by download_publications.py.

    Covers every source that script merges (OpenAlex, Crossref, PMLR, arXiv,
    manual), not just arXiv.
    """
    documents = []
    if not os.path.exists(papers_dir):
        print(f"  Papers directory {papers_dir} not found, skipping...")
        return documents

    md_files = [f for f in os.listdir(papers_dir) if f.endswith('.md')]
    for md_file in md_files:
        try:
            filepath = os.path.join(papers_dir, md_file)
            with open(filepath, 'r', encoding='utf-8') as f:
                content = f.read()

            doc = Document(
                page_content=content,
                metadata={
                    'source': f"docs/papers/{md_file}",
                    'type': 'publication',
                }
            )
            documents.append(doc)
        except Exception as e:
            print(f"  Error loading {md_file}: {e}")
    return documents


def load_linkedin_export(linkedin_dir):
    """Load LinkedIn data export files.
    
    LinkedIn lets you download your data at: https://www.linkedin.com/mypreferences/d/download-my-data
    Place the exported CSV/JSON files in docs/linkedin/.
    Supports: Profile.csv, Positions.csv, Education.csv, Skills.csv, Projects.csv, and any .txt/.md files.
    """
    documents = []
    if not os.path.exists(linkedin_dir):
        print(f"  LinkedIn directory {linkedin_dir} not found, skipping...")
        print(f"  To add LinkedIn data: export from https://www.linkedin.com/mypreferences/d/download-my-data")
        print(f"  Then place files in {linkedin_dir}/")
        return documents

    import csv
    for filename in os.listdir(linkedin_dir):
        filepath = os.path.join(linkedin_dir, filename)
        try:
            if filename.endswith('.csv'):
                with open(filepath, 'r', encoding='utf-8', errors='replace') as f:
                    reader = csv.DictReader(f)
                    rows = list(reader)
                if not rows:
                    continue
                content = f"# LinkedIn Data: {filename}\n\n"
                for row in rows:
                    content += "\n".join(f"**{k}:** {v}" for k, v in row.items() if v) + "\n---\n"
                doc = Document(
                    page_content=content,
                    metadata={'source': f"docs/linkedin/{filename}", 'type': 'linkedin_export'}
                )
                documents.append(doc)
            elif filename.endswith(('.txt', '.md')):
                with open(filepath, 'r', encoding='utf-8', errors='replace') as f:
                    content = f.read()
                if content.strip():
                    doc = Document(
                        page_content=content,
                        metadata={'source': f"docs/linkedin/{filename}", 'type': 'linkedin_export'}
                    )
                    documents.append(doc)
        except Exception as e:
            print(f"  Error loading {filename}: {e}")

    return documents


def main():
    """Main function to build the FAISS index"""
    print("=== Building FAISS Index for Julien Serbanescu RAG System ===")
    
    # Create output directory
    output_dir = "docs/faiss"
    os.makedirs(output_dir, exist_ok=True)
    
    # Initialize embedding function
    embedding_function = CohereEmbeddingsForIndexing(co)
    print(f"Embedding dimension: {embedding_function.embed_dim}")
    
    # Collect all documents
    all_documents = []
    
    # 1. Scrape portfolio website
    print("\n1. Scraping portfolio website...")
    portfolio_url = "https://julien-ser.github.io/JulienSerbanescu/"
    portfolio_content = scrape_website(portfolio_url, max_pages=5)
    if portfolio_content:
        doc = Document(
            page_content=portfolio_content,
            metadata={
                'source': portfolio_url,
                'type': 'portfolio_website'
            }
        )
        all_documents.append(doc)
        print(f"   Scraped {len(portfolio_content)} characters from portfolio website")
    else:
        print("   Warning: Could not scrape portfolio website content")
    
    # 1b. Load LinkedIn export data
    print("\n1b. Loading LinkedIn export data...")
    linkedin_docs = load_linkedin_export("docs/linkedin")
    all_documents.extend(linkedin_docs)
    print(f"   Loaded {len(linkedin_docs)} LinkedIn documents")
    
    # 2. Load PDF documents
    print("\n2. Loading PDF documents...")
    pdf_dir = "docs/pdfs"
    pdf_docs = load_pdf_documents(pdf_dir)
    all_documents.extend(pdf_docs)
    print(f"   Loaded {len(pdf_docs)} PDF pages")
    
    # 3. Load audio transcripts
    print("\n3. Loading audio transcripts...")
    transcript_dir = "docs/youtube"
    transcript_docs = load_audio_transcripts(transcript_dir)
    all_documents.extend(transcript_docs)
    print(f"   Loaded {len(transcript_docs)} audio transcripts")
    
    # 4. Load README documents
    print("\n4. Loading README documents...")
    readme_dir = "docs/readmes"
    readme_docs = load_readme_documents(readme_dir)
    all_documents.extend(readme_docs)
    print(f"   Loaded {len(readme_docs)} README documents")

    # 5. Load GitHub issues & PRs
    print("\n5. Loading GitHub issues & PRs...")
    activity_docs = load_github_activity("docs/github_activity")
    all_documents.extend(activity_docs)
    print(f"   Loaded {len(activity_docs)} GitHub issues/PRs")

    # 6. Load publications (all sources, not just arXiv)
    print("\n6. Loading publications...")
    paper_docs = load_publications("docs/papers")
    all_documents.extend(paper_docs)
    print(f"   Loaded {len(paper_docs)} publications")
    
    print(f"\nTotal documents collected: {len(all_documents)}")
    
    if not all_documents:
        print("No documents found! Please check your document directories.")
        return
    
    # 7. Split documents into chunks
    print("\n7. Splitting documents into chunks...")
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200,
        length_function=len,
    )
    
    split_documents = text_splitter.split_documents(all_documents)
    print(f"   Created {len(split_documents)} document chunks")
    
    # 8. Create embeddings
    print("\n8. Creating embeddings...")
    texts = [doc.page_content for doc in split_documents]
    embeddings = embedding_function.embed_documents(texts)
    
    # Convert to numpy array
    embedding_matrix = np.array(embeddings).astype('float32')
    print(f"   Created embedding matrix with shape: {embedding_matrix.shape}")
    
    # 9. Build FAISS index
    print("\n9. Building FAISS index...")
    dimension = embedding_matrix.shape[1]
    index = faiss.IndexFlatIP(dimension)  # Inner product for cosine similarity
    
    # Normalize embeddings for cosine similarity
    faiss.normalize_L2(embedding_matrix)
    index.add(embedding_matrix)
    
    print(f"   FAISS index built with {index.ntotal} vectors")
    
    # 10. Create docstore and mapping
    print("\n10. Creating docstore and mapping...")
    docstore = InMemoryDocstore()
    index_to_docstore_id = {}
    
    for i, doc in enumerate(split_documents):
        doc_id = f"doc_{i}"
        docstore.add({doc_id: doc})
        index_to_docstore_id[i] = doc_id
    
    # 11. Save everything
    print("\n11. Saving index files...")
    
    # Save FAISS index
    faiss_index_path = os.path.join(output_dir, "index.faiss")
    faiss.write_index(index, faiss_index_path)
    print(f"   Saved FAISS index to: {faiss_index_path}")
    
    # Save docstore and mapping
    pkl_path = os.path.join(output_dir, "index.pkl")
    with open(pkl_path, 'wb') as f:
        pickle.dump((docstore, index_to_docstore_id), f)
    print(f"   Saved docstore to: {pkl_path}")
    
    # Save separate metadata list
    metadata_list = [doc.metadata for doc in split_documents]
    metadata_path = os.path.join(output_dir, "metadata.pkl")
    with open(metadata_path, 'wb') as f:
        pickle.dump(metadata_list, f)
    print(f"   Saved metadata to: {metadata_path}")
    
    # 12. Create human-readable document lookup
    print("\n12. Creating document lookup file...")
    lookup_path = os.path.join(output_dir, "document_lookup.txt")
    with open(lookup_path, 'w', encoding='utf-8') as f:
        for i, doc in enumerate(split_documents):
            f.write(f"Document {i}:\n")
            f.write(f"Source: {doc.metadata.get('source', 'Unknown')}\n")
            f.write(f"Type: {doc.metadata.get('type', 'Unknown')}\n")
            f.write(f"Content Preview: {doc.page_content[:100]}...\n")
            f.write("-" * 80 + "\n")
    
    print(f"   Saved document lookup to: {lookup_path}")
    
    print(f"\n=== Index building completed successfully! ===")
    print(f"Total documents indexed: {len(split_documents)}")
    print(f"Index files saved to: {output_dir}/")
    print("\nYou can now run the query system with: python queryrun.py")

if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        print(f"Error during index building: {e}")
        traceback.print_exc()
        sys.exit(1)