import logging
from pathlib import Path
from PyPDF2 import PdfReader
from sentence_transformers import SentenceTransformer
import faiss
import json
import concurrent.futures
import numpy as np
from typing import List, Tuple
from datetime import datetime

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('script_analysis.log'),
        logging.StreamHandler()
    ]
)


class EmbeddingManager:
    def __init__(self, output_dir: Path, max_workers: int = 6):
        self.model = SentenceTransformer('all-MiniLM-L6-v2')
        self.chunk_size = 512
        self.output_dir = output_dir
        (self.output_dir / 'embeddings').mkdir(exist_ok=True)
        self.max_workers = max_workers

    def extract_text(self, pdf_path: str) -> str:
        with open(pdf_path, 'rb') as file:
            reader = PdfReader(file)
            text = ''
            for page in reader.pages:
                text += page.extract_text() + '\n'
            return text

    def create_chunks(self, text: str) -> list:
        words = text.split()
        chunks = []
        current_chunk = []
        current_size = 0

        for word in words:
            current_size += len(word) + 1
            if current_size > self.chunk_size:
                chunks.append(' '.join(current_chunk))
                current_chunk = [word]
                current_size = len(word)
            else:
                current_chunk.append(word)

        if current_chunk:
            chunks.append(' '.join(current_chunk))

        return chunks
    
    def process_chunk_batch(self, chunks: List[str]) -> np.ndarray:
        """Process a batch of chunks and return their embeddings."""
        try:
            return self.model.encode(chunks)
        except Exception as e:
            logging.error(f"Error encoding chunk batch: {e}")
            raise
    
    def save_embeddings(self, chunks: list,file_name:str):
        try:
            embeddings_dir = self.output_dir / 'embeddings' / file_name
            embeddings_dir.mkdir(parents=True, exist_ok=True)
            
            # Split chunks into batches for parallel processing
            batch_size = len(chunks) // self.max_workers
            batches = [chunks[i:i + batch_size] for i in range(0, len(chunks), batch_size)]
            
            # Process batches in parallel
            embeddings_list = []
            with concurrent.futures.ProcessPoolExecutor(max_workers=self.max_workers) as executor:
                future_to_batch = {executor.submit(self.process_chunk_batch, batch): batch 
                                 for batch in batches}
                
                for future in concurrent.futures.as_completed(future_to_batch):
                    batch_embeddings = future.result()
                    embeddings_list.append(batch_embeddings)
            
            # Combine all embeddings
            embeddings = np.vstack(embeddings_list)
            
            # Create and save FAISS index
            index = faiss.IndexFlatL2(embeddings.shape[1])
            index.add(embeddings.astype('float32'))
            faiss.write_index(index, str(embeddings_dir / 'script.index'))
            
            # Save metadata and chunks
            metadata = {
                'file_name': file_name,
                'timestamp': datetime.now().isoformat(),
                'num_chunks': len(chunks),
                'chunk_size': self.chunk_size,
                'model': 'all-MiniLM-L6-v2',
                'embedding_dimension': embeddings.shape[1],
                'num_workers': self.max_workers
            }
            
            with open(embeddings_dir / 'metadata.json', 'w') as f:
                json.dump(metadata, f, indent=4)
                
            with open(embeddings_dir / 'chunks.json', 'w') as f:
                json.dump(chunks, f, indent=4)
                
            logging.info(f"Saved embeddings and metadata to {embeddings_dir}")
            return embeddings_dir
            
        except Exception as e:
            logging.error(f"Error saving embeddings: {e}")
            raise

    def process_script(self, data: str,filename):
        chunks = self.create_chunks(data)
        logging.info(f"Created {len(chunks)} chunks from script")
        embeddings_dir = self.save_embeddings(chunks,file_name=filename)
        return chunks, embeddings_dir