# services/pdf_service.py
from pathlib import Path
from typing import List, Dict, Any, Optional
from PyPDF2 import PdfReader
from langchain.text_splitter import RecursiveCharacterTextSplitter
import faiss
import numpy as np
import asyncio
from concurrent.futures import ThreadPoolExecutor
import logging
from datetime import datetime
from config.config import settings

logger = logging.getLogger(__name__)

class PDFService:
    def __init__(self, model_service):
        self.embedder = model_service.embedder
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=settings.CHUNK_SIZE,
            chunk_overlap=settings.CHUNK_OVERLAP,
            separators=["\n\n", "\n", ".", "!", "?", ",", " ", ""]
        )
        self.index = None
        self.chunks = []
        self.last_update = None
        self.pdf_metadata = {}

    def process_pdf(self, pdf_path: Path) -> List[Dict[str, Any]]:
        """Process a single PDF file - now synchronous"""
        try:
            reader = PdfReader(str(pdf_path))
            chunks = []
            
            # Extract metadata
            metadata = {
                'title': reader.metadata.get('/Title', ''),
                'author': reader.metadata.get('/Author', ''),
                'creation_date': reader.metadata.get('/CreationDate', ''),
                'pages': len(reader.pages),
                'filename': pdf_path.name
            }
            self.pdf_metadata[pdf_path.name] = metadata
            
            # Process each page
            for page_num, page in enumerate(reader.pages):
                text = page.extract_text()
                if not text:
                    continue
                    
                page_chunks = self.text_splitter.split_text(text)
                for i, chunk in enumerate(page_chunks):
                    chunks.append({
                        'text': chunk,
                        'source': pdf_path.name,
                        'page': page_num + 1,
                        'chunk_index': i,
                        'metadata': metadata,
                        'timestamp': datetime.now().isoformat()
                    })
            print("--------------------------- chunks ----------------------------------")
            print("--------------------------- chunks ----------------------------------")
            print(chunks)         
            return chunks
            
        except Exception as e:
            logger.error(f"Error processing PDF {pdf_path}: {e}")
            return []

    async def index_pdfs(self, pdf_folder: Path = settings.PDF_FOLDER) -> None:
        """Index all PDFs in the specified folder"""
        try:
            pdf_files = list(pdf_folder.glob('*.pdf'))
            if not pdf_files:
                logger.warning(f"No PDF files found in {pdf_folder}")
                return
            
            # Process PDFs using thread pool
            loop = asyncio.get_running_loop()
            with ThreadPoolExecutor() as executor:
                chunk_lists = await loop.run_in_executor(
                    executor,
                    lambda: [self.process_pdf(pdf_file) for pdf_file in pdf_files]
                )
            
            # Combine all chunks
            self.chunks = []
            for chunk_list in chunk_lists:
                self.chunks.extend(chunk_list)
            
            if not self.chunks:
                logger.warning("No text chunks extracted from PDFs")
                return
                
            # Create FAISS index
            texts = [chunk['text'] for chunk in self.chunks]
            embeddings = await loop.run_in_executor(
                None,
                lambda: self.embedder.encode(
                    texts,
                    convert_to_tensor=True,
                    show_progress_bar=True
                ).cpu().detach().numpy()
            )
            
            dimension = embeddings.shape[1]
            self.index = faiss.IndexFlatL2(dimension)
            self.index.add(embeddings)
            
            self.last_update = datetime.now()
            
            logger.info(f"Indexed {len(self.chunks)} chunks from {len(pdf_files)} PDFs")
            
        except Exception as e:
            logger.error(f"Error indexing PDFs: {e}")
            raise

    async def search(
        self,
        query: str,
        top_k: int = 5,
        min_score: float = 0.5
    ) -> List[Dict[str, Any]]:
        """Search indexed PDFs with debug logs"""
        print("--------------------------- query ----------------------------------")
        print(query)
        if not self.index or not self.chunks:
            await self.index_pdfs()
    
        try:
            # Create query embedding
            query_embedding = self.embedder.encode([query], convert_to_tensor=True)
            query_embedding_np = query_embedding.cpu().detach().numpy()
            print("Query Embedding Shape:", query_embedding_np.shape)
    
            # Search in FAISS index
            distances, indices = self.index.search(query_embedding_np, top_k)
            print("Distances:", distances)
            print("Indices:", indices)
    
            # Process results
            results = []
            for i, idx in enumerate(indices[0]):
                if idx >= len(self.chunks):
                    continue  # Skip invalid indices
    
                score = 1 - distances[0][i]  # Convert distance to similarity score
                print(f"Chunk Index: {idx}, Distance: {distances[0][i]}, Score: {score}")
                print("----- score < min_score")
                print(score < min_score)
                if score < min_score:
                    print("skipped ---- ")
                    
                    #continue  # Skip low scores
    
                chunk = self.chunks[idx].copy()
                chunk['score'] = score
                print("---- chuck " )
                print(chunk)
                results.append(chunk)
    
            # Sort by score and take top_k
            results.sort(key=lambda x: x['score'], reverse=True)
    
            print("--------------------------- results ----------------------------------")
            print(results)
    
            return results[:top_k]
    
        except Exception as e:
            logger.error(f"Error searching PDFs: {e}")
            raise