# Document processing utilities for different file types

import os
import tempfile
import numpy as np
import torch
from PIL import Image
import fitz  # PyMuPDF
from docx import Document
import streamlit as st
from config import EMBEDDING_DIMENSIONS
from device_setup import clear_gpu_memory


class DocumentProcessor:
    """Handles processing of different document types"""
    
    def __init__(self, clip_model, clip_processor, whisper_model, ocr_reader, text_encoder, device):
        self.clip_model = clip_model
        self.clip_processor = clip_processor
        self.whisper_model = whisper_model
        self.ocr_reader = ocr_reader
        self.text_encoder = text_encoder
        self.device = device
    
    def extract_text_from_pdf(self, pdf_path: str) -> str:
        """extract text  from PDF"""
        try:
            doc = fitz.open(pdf_path)
            text = ""
            for page in doc:
                text += page.get_text()
            doc.close()
            return text.strip()
        except Exception as e:
            st.error(f"Error extracting PDF text: {e}")
            return ""

    def extract_text_from_docx(self, docx_path: str) -> str:
        """extract text content from DOCX"""
        try:
            doc = Document(docx_path)
            text = ""
            for paragraph in doc.paragraphs:
                text += paragraph.text + "\n"
            return text.strip()
        except Exception as e:
            st.error(f"Error extracting DOCX text: {e}")
            return ""

    def extract_text_from_image(self, image_path: str) -> str:
        """extract text from image using OCR"""
        try:
            results = self.ocr_reader.readtext(image_path)
            text = " ".join([result[1] for result in results])
            return text.strip()
        except Exception as e:
            st.error(f"Error extracting text from image: {e}")
            return ""

    def get_image_embedding(self, image_path: str) -> np.ndarray:
        """generate CLIP embedding for image"""
        try:
            image = Image.open(image_path).convert("RGB")
            inputs = self.clip_processor(images=image, return_tensors="pt", padding=True)
            inputs = {k: v.to(self.device) for k, v in inputs.items()}

            with torch.no_grad():
                image_features = self.clip_model.get_image_features(**inputs)

            embedding = image_features.cpu().numpy().flatten()
            clear_gpu_memory()
            return embedding
        except Exception as e:
            st.error(f"Error getting image embedding: {e}")
            # Return zero embedding with same dimension
            clip_dim = EMBEDDING_DIMENSIONS.get("clip_base", 512)
            return np.zeros(clip_dim)

    def transcribe_audio(self, audio_path: str) -> str:
        """transcribe audio file to text using Whisper"""
        try:
            result = self.whisper_model.transcribe(audio_path)
            return result["text"].strip()
        except Exception as e:
            st.error(f"Error transcribing audio: {e}")
            return ""

    def get_text_embedding(self, text: str) -> np.ndarray:
        """generate text embedding using sentence transformer"""
        try:
            if not text.strip():
                return np.zeros(EMBEDDING_DIMENSIONS["text"])
            
            with torch.no_grad():
                embedding = self.text_encoder.encode([text], convert_to_numpy=True)
            return embedding[0]
        except Exception as e:
            st.error(f"Error getting text embedding: {e}")
            return np.zeros(EMBEDDING_DIMENSIONS["text"])

    def process_uploaded_file(self, uploaded_file, file_type: str, data_dir) -> dict:
        """
        Process uploaded file and return document dictionary
        Returns None if processing fails
        """
        try:
            # save uploaded file to temporary location
            with tempfile.NamedTemporaryFile(delete=False, suffix=f".{file_type}") as tmp_file:
                tmp_file.write(uploaded_file.read())
                tmp_path = tmp_file.name

            # process based on file type
            if file_type == 'pdf':
                content = self.extract_text_from_pdf(tmp_path)
                embedding = self.get_text_embedding(content)
                embedding_type = 'text'
                
            elif file_type == 'docx':
                content = self.extract_text_from_docx(tmp_path)
                embedding = self.get_text_embedding(content)
                embedding_type = 'text'
                
            elif file_type in ['png', 'jpg', 'jpeg']:
                content = self.extract_text_from_image(tmp_path)
                embedding = self.get_image_embedding(tmp_path)
                embedding_type = 'image'
                
            elif file_type in ['mp3', 'wav', 'ogg', 'm4a']:
                content = self.transcribe_audio(tmp_path)
                embedding = self.get_text_embedding(content)
                embedding_type = 'text'
            else:
                os.unlink(tmp_path)
                return None

            # clean up temp file
            os.unlink(tmp_path)
            
            return {
                'content': content,
                'embedding': embedding,
                'embedding_type': embedding_type
            }
            
        except Exception as e:
            st.error(f"Error processing file {uploaded_file.name}: {e}")
            if 'tmp_path' in locals() and os.path.exists(tmp_path):
                os.unlink(tmp_path)
            return None