# super_server.py
import os
import fitz  # PyMuPDF
import random
from PIL import Image # Needed for size filtering
import io
from mcp.server.fastmcp import FastMCP
from serpapi import GoogleSearch
from dotenv import load_dotenv

# --- IMPORTS FOR TEXTBOOK (RAG) ---
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.openai import OpenAI as LlamaOpenAI

load_dotenv()

mcp = FastMCP("Anato-Mitra Super Server")

# ==========================================
# 📸 PART 0: SMART IMAGE EXTRACTOR (Filters Trash)
# ==========================================
IMAGE_DIR = "extracted_images"

def extract_images_from_pdfs():
    """
    Scans PDFs and saves ONLY large, relevant diagrams.
    Filters out small icons, bullets, and emojis.
    """
    if not os.path.exists("textbooks"):
        os.makedirs("textbooks")
        return

    if not os.path.exists(IMAGE_DIR):
        os.makedirs(IMAGE_DIR)

    print("📸 Scanning PDFs for VIVA Diagrams...")
    
    for pdf_file in os.listdir("textbooks"):
        if not pdf_file.endswith(".pdf"):
            continue
            
        pdf_path = os.path.join("textbooks", pdf_file)
        doc = fitz.open(pdf_path)
        
        for page_index, page in enumerate(doc):
            image_list = page.get_images(full=True)
            
            for img_index, img in enumerate(image_list):
                xref = img[0]
                base_image = doc.extract_image(xref)
                image_bytes = base_image["image"]
                image_ext = base_image["ext"]
                
                # --- TRASH FILTERING ---
                # Use PIL to check dimensions
                try:
                    pil_image = Image.open(io.BytesIO(image_bytes))
                    width, height = pil_image.size
                    
                    # Logic: Anatomy diagrams are usually big. Icons/Bullets are small.
                    # Filter: Keep only images larger than 200x200 pixels
                    if width < 200 or height < 200:
                        continue
                        
                    # Filter: Remove extremely wide/flat images (headers/footers)
                    aspect_ratio = width / height
                    if aspect_ratio > 4 or aspect_ratio < 0.25:
                        continue
                        
                except:
                    continue # If image is broken, skip
                
                # Save valid diagram
                image_name = f"{pdf_file.replace('.pdf', '')}_page_{page_index + 1}_img_{img_index}.{image_ext}"
                image_path = os.path.join(IMAGE_DIR, image_name)
                
                if not os.path.exists(image_path):
                    with open(image_path, "wb") as f:
                        f.write(image_bytes)
    
    print(f"✅ VIVA Diagrams ready in '{IMAGE_DIR}/'")

extract_images_from_pdfs()

# ==========================================
# 🎲 PART 1: VIVA QUIZ TOOLS
# ==========================================

@mcp.tool()
def get_random_quiz_image() -> str:
    """
    Selects a random anatomy diagram from the extracted textbook images for a VIVA quiz.
    Returns the path to the image and its context (Book/Page).
    """
    if not os.path.exists(IMAGE_DIR):
        return "No images found."
        
    images = [f for f in os.listdir(IMAGE_DIR) if f.endswith(('.png', '.jpg', '.jpeg'))]
    
    if not images:
        return "No diagrams available for quiz."
        
    random_image = random.choice(images)
    image_path = os.path.join(IMAGE_DIR, random_image)
    
    return f"QUIZ IMAGE SELECTED:\n![Quiz Diagram]({image_path})\n\n(Internal Context for Agent: {random_image})"

# ==========================================
# 🧠 PART 2: TEXTBOOK KNOWLEDGE
# ==========================================
QUERY_ENGINE = None

def initialize_textbook():
    global QUERY_ENGINE
    try:
        print("📚 Configuring Textbook Engine...")
        Settings.embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")
        Settings.llm = LlamaOpenAI(
            model="meta-llama/Meta-Llama-3.1-70B-Instruct",
            api_key=os.getenv("HYPERBOLIC_API_KEY"), 
            api_base="https://api.hyperbolic.xyz/v1"
        )
        documents = SimpleDirectoryReader("textbooks").load_data()
        if documents:
            index = VectorStoreIndex.from_documents(documents)
            QUERY_ENGINE = index.as_query_engine(similarity_top_k=3)
            print("✅ Textbook Indexed!")
    except Exception as e:
        print(f"❌ Textbook Error: {e}")

initialize_textbook()

@mcp.tool()
def consult_medical_textbook(question: str) -> str:
    """Consults uploaded medical textbooks."""
    global QUERY_ENGINE
    if not QUERY_ENGINE: return "Textbook unavailable."
    try:
        return str(QUERY_ENGINE.query(question))
    except Exception as e: return str(e)

# ==========================================
# 👁️ PART 3: EXTERNAL SEARCH
# ==========================================
@mcp.tool()
def search_anatomy_diagrams(topic: str) -> str:
    """Searches Google Images."""
    api_key = os.getenv("SERPAPI_KEY")
    if api_key:
        try:
            search = GoogleSearch({"q": f"{topic} anatomy schematic", "engine": "google_images", "api_key": api_key})
            res = search.get_dict().get("images_results", [])
            if res: return f"![Diagram]({res[0]['original']})"
        except: pass
    return "No external diagram found."

if __name__ == "__main__":
    mcp.run()