super_server.py

# super_server.py
import os
import fitz  # PyMuPDF
import random
import io
import requests 
import math
from PIL import Image, ImageFilter, ImageStat 
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 1: DROPBOX DOWNLOADER
# ==========================================
def download_textbook_from_dropbox():
    if not os.path.exists("textbooks"): os.makedirs("textbooks")
    # If empty, download
    if not os.listdir("textbooks"):
        dropbox_url = os.getenv("DROPBOX_URL")
        if not dropbox_url: return
        if "dl=0" in dropbox_url: dropbox_url = dropbox_url.replace("dl=0", "dl=1")
        
        print("📥 Downloading Textbook from Dropbox...")
        try:
            response = requests.get(dropbox_url, stream=True)
            if response.status_code == 200:
                with open(os.path.join("textbooks", "main_textbook.pdf"), 'wb') as f:
                    for chunk in response.iter_content(chunk_size=8192): f.write(chunk)
                print("✅ Download Complete")
        except: pass

download_textbook_from_dropbox()

# ==========================================
# 📸 PART 2: SMART IMAGE EXTRACTOR (With Filter)
# ==========================================
IMAGE_DIR = "extracted_images"

def is_valid_diagram(pil_image):
    """Filters out headers, footers, and empty boxes."""
    width, height = pil_image.size
    if width < 200 or height < 200: return False 
    aspect = width / height
    if aspect > 3.5 or aspect < 0.25: return False
    
    try:
        # Entropy check (detects flat images)
        histogram = pil_image.histogram()
        entropy = 0
        total_pixels = sum(histogram)
        for count in histogram:
            if count == 0: continue
            p = count / total_pixels
            entropy -= p * math.log2(p)
        if entropy < 3.2: return False

        # Edge detection (detects complexity)
        gray_img = pil_image.convert("L")
        edges = gray_img.filter(ImageFilter.FIND_EDGES)
        stat = ImageStat.Stat(edges)
        if stat.var[0] < 800: return False
    except: return False

    return True

def extract_images_from_pdfs():
    if not os.path.exists("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]
                try:
                    base_image = doc.extract_image(xref)
                    pil_image = Image.open(io.BytesIO(base_image["image"]))
                    
                    if is_valid_diagram(pil_image):
                        image_name = f"{pdf_file.replace('.pdf', '')}_page_{page_index + 1}_img_{img_index}.{base_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(base_image["image"])
                except: continue 
    print(f"✅ VIVA Diagrams ready in '{IMAGE_DIR}/'")

extract_images_from_pdfs()

# ==========================================
# 🎲 PART 3: VIVA TOOLS (Textbook & Internet)
# ==========================================

@mcp.tool()
def get_random_quiz_image() -> str:
    """Get random image from LOCAL PDF."""
    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."
    random_image = random.choice(images)
    return f"![Quiz Diagram](extracted_images/{random_image})\n(Internal Context: {random_image})"

@mcp.tool()
def get_random_anatomy_topic() -> str:
    """Get a random high-yield topic for INTERNET search."""
    topics = [
        "Circle of Willis", "Brachial Plexus", "Femoral Triangle", "Cubital Fossa",
        "Popliteal Fossa", "Stomach Bed", "Porta Hepatis", "Hilum of Lung",
        "Interior of Heart Right Atrium", "Cavernous Sinus", "Sciatic Nerve Course",
        "Rotator Cuff Muscles", "Carpal Tunnel Contents", "Inguinal Canal"
    ]
    return random.choice(topics)

# ==========================================
# 🧠 PART 4: TEXTBOOK KNOWLEDGE (RAG)
# ==========================================
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:
    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 5: INTERNET SEARCH
# ==========================================
@mcp.tool()
def search_anatomy_diagrams(topic: str) -> str:
    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()