import os
import faiss
import pickle
import numpy as np
from typing import List, Dict
from docarray import Document, DocumentArray
from jina import Executor, requests
from sentence_transformers import SentenceTransformer
from transformers import BlipProcessor, BlipForConditionalGeneration, AutoTokenizer, AutoModelForSeq2SeqLM
from pdfminer.high_level import extract_text
import fitz  # PyMuPDF - better than pdf2image
from PIL import Image
import traceback
import torch
import re


class MultimodalRAGExecutor(Executor):
    def __init__(
        self,
        llm_model_name: str = "google/flan-t5-large",  # FREE and reliable
        embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2",  # FREE embeddings
        image_model: str = "Salesforce/blip-image-captioning-base",  # FREE vision model
        index_file: str = "faiss_index.bin",
        metadata_file: str = "metadata.pkl",
        dim: int = 384,  # MiniLM embedding dimension
        use_huggingface_api: bool = False,
        **kwargs,
    ):
        super().__init__(**kwargs)

        self.llm_model_name = llm_model_name
        self.embedding_model = embedding_model
        self.image_model = image_model
        self.index_file = index_file
        self.metadata_file = metadata_file
        self.dim = dim
        self.use_huggingface_api = use_huggingface_api

        # Initialize FREE embedding model
        print(f"📥 Loading embedding model: {embedding_model}...")
        self.embedder = SentenceTransformer(embedding_model)
        print("✅ Embedding model loaded!")

        # Initialize FREE image captioning model
        print(f"📥 Loading image model: {image_model}...")
        try:
            self.blip_processor = BlipProcessor.from_pretrained(image_model)
            self.blip_model = BlipForConditionalGeneration.from_pretrained(image_model)
            print("✅ Image model loaded!")
        except Exception as e:
            print(f"⚠️ Warning: Could not load BLIP model: {e}")
            self.blip_processor = None
            self.blip_model = None

        # Initialize LLM
        self._initialize_llm()

        # Load or create FAISS index
        self._load_or_create_index()

    def _initialize_llm(self):
        """Initialize LLM (API or local)"""
        if self.use_huggingface_api:
            hf_token = os.getenv("HUGGINGFACE_API_KEY")
            if hf_token:
                try:
                    from huggingface_hub import InferenceClient
                    self.llm_client = InferenceClient(token=hf_token)
                    self.llm_mode = "api"
                    print("✅ Using HuggingFace API for LLM!")
                    return
                except Exception as e:
                    print(f"⚠️ HF API initialization failed: {e}")
        
        # Load local LLM (FLAN-T5 works best)
        print(f"📥 Loading local LLM: {self.llm_model_name}...")
        try:
            self.tokenizer = AutoTokenizer.from_pretrained(self.llm_model_name)
            self.llm_model = AutoModelForSeq2SeqLM.from_pretrained(self.llm_model_name)
            self.llm_mode = "local"
            print("✅ Local LLM loaded!")
        except Exception as e:
            print(f"⚠️ LLM loading failed: {e}")
            self.llm_mode = "extractive"
            print("⚠️ Falling back to extractive QA")

    def _load_or_create_index(self):
        """Load existing index or create new one"""
        if os.path.exists(self.index_file) and os.path.exists(self.metadata_file):
            try:
                self.index = faiss.read_index(self.index_file)
                with open(self.metadata_file, "rb") as f:
                    self.metadata = pickle.load(f)
                print(f"✅ Loaded existing index with {self.index.ntotal} vectors")
            except Exception as e:
                print(f"⚠️ Could not load index: {e}. Creating new one.")
                self.index = faiss.IndexFlatL2(self.dim)
                self.metadata = []
        else:
            self.index = faiss.IndexFlatL2(self.dim)
            self.metadata = []
            print("✅ Created new empty index")

    def _get_embedding(self, text: str) -> np.ndarray:
        """Generate embeddings using sentence-transformers"""
        embedding = self.embedder.encode(text, convert_to_numpy=True)
        return embedding.astype(np.float32)

    def _analyze_image(self, image: Image.Image, context: str = "") -> str:
        """Analyze image using BLIP model"""
        if not self.blip_processor or not self.blip_model:
            return "Image analysis unavailable"
        
        try:
            # Generate caption
            inputs = self.blip_processor(image, return_tensors="pt")
            out = self.blip_model.generate(**inputs, max_length=100)
            caption = self.blip_processor.decode(out[0], skip_special_tokens=True)
            
            # Generate detailed description
            text_prompt = "a detailed description of"
            inputs = self.blip_processor(image, text_prompt, return_tensors="pt")
            out = self.blip_model.generate(**inputs, max_length=150)
            detailed = self.blip_processor.decode(out[0], skip_special_tokens=True)
            
            description = f"Caption: {caption}. Details: {detailed}"
            if context:
                description += f". Context: {context[:100]}"
            
            return description
        except Exception as e:
            print(f"Image analysis error: {e}")
            return "Image analysis failed"

    def _extract_text_from_pdf(self, pdf_path: str) -> List[str]:
        """Extract text from PDF using PyMuPDF"""
        texts = []
        try:
            doc = fitz.open(pdf_path)
            for page_num, page in enumerate(doc, start=1):
                text = page.get_text("text")
                if text and text.strip():
                    texts.append(f"Page {page_num}:\n{text.strip()}")
            doc.close()
        except Exception as e:
            print(f"Text extraction error: {e}")
        return texts

    def _extract_images_from_pdf(self, pdf_path: str) -> List[Dict]:
        """Extract images from PDF using PyMuPDF"""
        images_data = []
        try:
            doc = fitz.open(pdf_path)
            for page_num, page in enumerate(doc, start=1):
                image_list = page.get_images(full=True)
                
                for img_idx, img in enumerate(image_list):
                    try:
                        xref = img[0]
                        base_image = doc.extract_image(xref)
                        image_bytes = base_image["image"]
                        
                        # Convert to PIL Image
                        import io
                        pil_image = Image.open(io.BytesIO(image_bytes)).convert('RGB')
                        
                        # Filter small images
                        width, height = pil_image.size
                        if width >= 100 and height >= 100:
                            images_data.append({
                                'image': pil_image,
                                'page': page_num,
                                'index': img_idx
                            })
                    except Exception as e:
                        print(f"Error extracting image on page {page_num}: {e}")
                        continue
            
            doc.close()
        except Exception as e:
            print(f"Image extraction error: {e}")
        
        return images_data

    def _generate_answer_local(self, prompt: str) -> str:
        """Generate answer using local FLAN-T5 model"""
        try:
            inputs = self.tokenizer(
                prompt, 
                return_tensors="pt", 
                max_length=512, 
                truncation=True
            )
            
            outputs = self.llm_model.generate(
                inputs.input_ids,
                max_new_tokens=200,
                temperature=0.3,
                do_sample=False,
                num_beams=4
            )
            
            answer = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            return answer.strip()
        except Exception as e:
            print(f"Local generation error: {e}")
            return ""

    def _generate_answer_api(self, prompt: str) -> str:
        """Generate answer using HuggingFace API"""
        try:
            response = self.llm_client.text_generation(
                prompt,
                model="google/flan-t5-xxl",
                max_new_tokens=250,
                temperature=0.3
            )
            return response.strip()
        except Exception as e:
            print(f"API generation error: {e}")
            return ""

    def _generate_extractive_answer(self, query: str, context: str) -> str:
        """Fallback extractive QA"""
        sentences = re.split(r'[.!?]+', context)
        sentences = [s.strip() for s in sentences if len(s.strip()) > 20]
        
        query_words = set(query.lower().split())
        scored = []
        
        for sent in sentences[:30]:
            sent_words = set(sent.lower().split())
            overlap = len(query_words.intersection(sent_words))
            if overlap > 0:
                scored.append((overlap, sent))
        
        scored.sort(reverse=True, key=lambda x: x[0])
        
        if scored:
            top_sentences = [s for _, s in scored[:3]]
            return ". ".join(top_sentences) + "."
        
        return "I couldn't find relevant information to answer your question."

    def _generate_answer(self, prompt: str, query: str = "", context: str = "") -> str:
        """Generate answer with fallback chain"""
        if self.llm_mode == "api":
            answer = self._generate_answer_api(prompt)
            if answer and len(answer) > 20:
                return answer
        
        if self.llm_mode == "local":
            answer = self._generate_answer_local(prompt)
            if answer and len(answer) > 20:
                return answer
        
        # Fallback to extractive
        return self._generate_extractive_answer(query, context)

    def _clean_answer(self, answer: str) -> str:
        """Clean unwanted artifacts from answer"""
        answer = re.sub(r'^(Answer|Context|Question):\s*', '', answer, flags=re.IGNORECASE)
        answer = re.sub(r'[#*_`]+', '', answer)
        answer = re.sub(r'\s+', ' ', answer).strip()
        return answer

    @requests(on="/upload")
    def upload(self, docs: DocumentArray, **kwargs):
        """Upload PDF, extract text + images, index them"""
        try:
            for doc in docs:
                file_path = doc.uri
                if file_path.startswith("file://"):
                    file_path = file_path.replace("file://", "")

                if not os.path.exists(file_path) or not file_path.endswith(".pdf"):
                    continue

                print(f"📄 Processing PDF: {file_path}")

                # Extract text
                text_chunks = self._extract_text_from_pdf(file_path)
                
                # Index text chunks
                text_count = 0
                context_text = "\n".join(text_chunks[:2])[:1000]  # Context for images
                
                for chunk in text_chunks:
                    # Split into smaller chunks
                    paragraphs = chunk.split("\n\n")
                    for para in paragraphs:
                        if para.strip() and len(para.strip()) > 50:
                            emb = self._get_embedding(para.strip())
                            self.index.add(np.array([emb]))
                            self.metadata.append({
                                "type": "text",
                                "content": para.strip()
                            })
                            text_count += 1

                print(f"✅ Indexed {text_count} text chunks")

                # Extract and analyze images
                images_data = self._extract_images_from_pdf(file_path)
                image_count = 0
                
                for img_data in images_data:
                    # Analyze image
                    description = self._analyze_image(img_data['image'], context_text)
                    
                    # Save image temporarily
                    img_path = os.path.abspath(
                        f"img_p{img_data['page']}_i{img_data['index']}.png"
                    )
                    img_data['image'].save(img_path, "PNG")
                    
                    # Embed description
                    embed_text = f"Image from page {img_data['page']}: {description}"
                    emb = self._get_embedding(embed_text)
                    self.index.add(np.array([emb]))
                    self.metadata.append({
                        "type": "image",
                        "content": f"file://{img_path}",
                        "description": description,
                        "page": img_data['page']
                    })
                    image_count += 1

                print(f"✅ Analyzed and indexed {image_count} images")

            # Save index
            faiss.write_index(self.index, self.index_file)
            with open(self.metadata_file, "wb") as f:
                pickle.dump(self.metadata, f)

            summary = f"✅ Upload complete!\n"
            summary += f"📊 Total vectors: {self.index.ntotal}\n"
            summary += f"📝 Text chunks: {text_count}\n"
            summary += f"🖼️ Images: {image_count}\n"
            summary += f"💚 100% FREE processing!"
            
            return DocumentArray([Document(text=summary)])
            
        except Exception as e:
            error_msg = f"❌ Upload failed:\n{traceback.format_exc()}"
            print(error_msg)
            return DocumentArray([Document(text=error_msg)])

    @requests(on="/query")
    def query(self, docs: DocumentArray, **kwargs):
        """Query with multimodal RAG"""
        results = DocumentArray()

        if self.index.ntotal == 0:
            return DocumentArray([
                Document(text="⚠️ No documents uploaded. Please upload a PDF first.")
            ])

        for doc in docs:
            try:
                query_text = doc.text
                
                # Get query embedding and search
                query_emb = self._get_embedding(query_text)
                D, I = self.index.search(np.array([query_emb]), k=8)

                context_parts = []
                matched_images = []
                image_descriptions = []
                
                for idx in I[0]:
                    if idx < len(self.metadata):
                        meta = self.metadata[idx]
                        if meta["type"] == "text":
                            context_parts.append(meta["content"])
                        elif meta["type"] == "image":
                            matched_images.append(Document(uri=meta["content"]))
                            image_descriptions.append(
                                f"[Image Page {meta.get('page', '?')}]: {meta['description']}"
                            )

                # Build context
                context_text = "\n\n".join(context_parts[:4])
                if image_descriptions:
                    context_text += "\n\nRelevant Images:\n" + "\n".join(image_descriptions[:2])

                # Truncate if too long
                if len(context_text) > 2000:
                    context_text = context_text[:2000] + "..."

                # Create prompt
                prompt = (
                    f"Answer the question based on the context from a research paper.\n\n"
                    f"Context:\n{context_text}\n\n"
                    f"Question: {query_text}\n\n"
                    f"Answer:"
                )

                # Generate answer
                answer = self._generate_answer(prompt, query_text, context_text)
                answer = self._clean_answer(answer)

                # Create result
                answer_doc = Document(text=answer)
                if matched_images:
                    answer_doc.chunks = DocumentArray(matched_images[:4])

                results.append(answer_doc)
                
            except Exception as e:
                error_msg = f"❌ Query failed: {str(e)}\n{traceback.format_exc()}"
                print(error_msg)
                results.append(Document(text=error_msg))
                
        return results

    @requests(on="/stats")
    def stats(self, docs: DocumentArray, **kwargs):
        """Get index statistics"""
        text_count = sum(1 for m in self.metadata if m["type"] == "text")
        image_count = sum(1 for m in self.metadata if m["type"] == "image")
        
        stats_text = (
            f"📊 Index Statistics:\n"
            f"Total vectors: {self.index.ntotal}\n"
            f"Text chunks: {text_count}\n"
            f"Images: {image_count}\n"
            f"LLM mode: {self.llm_mode}\n"
            f"💚 100% FREE models!"
        )
        
        return DocumentArray([Document(text=stats_text)])

    @requests(on="/reset")
    def reset(self, docs: DocumentArray, **kwargs):
        """Reset index and metadata"""
        try:
            self.index = faiss.IndexFlatL2(self.dim)
            self.metadata = []
            
            if os.path.exists(self.index_file):
                os.remove(self.index_file)
            if os.path.exists(self.metadata_file):
                os.remove(self.metadata_file)
            
            return DocumentArray([Document(text="✅ Index reset successfully!")])
        except Exception as e:
            return DocumentArray([
                Document(text=f"❌ Reset failed: {str(e)}")
            ])