# super_server.py
import os
import fitz  # PyMuPDF
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: IMAGE EXTRACTOR
# ==========================================
IMAGE_DIR = "extracted_images"

def extract_images_from_pdfs():
    """
    Runs on startup. Scans PDFs and saves images.
    Skips saving if the image already exists.
    """
    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 images...")
    
    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"]
                
                # Naming convention: bookname_page_X_img_Y.png
                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)
                
                # CHECK: Only save if it doesn't exist
                if not os.path.exists(image_path):
                    with open(image_path, "wb") as f:
                        f.write(image_bytes)
    
    print(f"✅ PDF Images ready in '{IMAGE_DIR}/'")

extract_images_from_pdfs()

# ==========================================
# 🧠 PART 1: TEXTBOOK KNOWLEDGE (With Page Citations)
# ==========================================
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 not documents: return

        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. 
    Returns answer + Diagrams + Page Numbers.
    """
    global QUERY_ENGINE
    if not QUERY_ENGINE:
        return "Textbook unavailable."
    
    try:
        # 1. Get Text Answer
        response = QUERY_ENGINE.query(question)
        text_answer = str(response)
        
        # 2. Find Matching Images & Page Numbers
        found_data = [] # We will store (path, page, book) here
        
        for node in response.source_nodes:
            page_num = node.metadata.get("page_label")
            file_name = node.metadata.get("file_name")
            
            if page_num and file_name:
                clean_name = file_name.replace('.pdf', '')
                search_prefix = f"{clean_name}_page_{page_num}_"
                
                # Check if we extracted an image from this specific page
                for img_file in os.listdir(IMAGE_DIR):
                    if img_file.startswith(search_prefix):
                        found_data.append({
                            "path": os.path.join(IMAGE_DIR, img_file),
                            "page": page_num,
                            "book": clean_name
                        })
        
        # 3. Format the Output
        if found_data:
            text_answer += "\n\n---\n### 📖 Textbook Reference Diagrams"
            
            # Limit to top 2 images to avoid spamming
            for item in found_data[:2]:
                # Add Image
                text_answer += f"\n![Textbook Diagram]({item['path']})"
                # Add Caption with Page Number
                text_answer += f"\n\n*Source: **{item['book']}**, Page **{item['page']}***\n"
                
        return text_answer

    except Exception as e:
        return f"Error: {str(e)}"

# ==========================================
# 👁️ PART 2: EXTERNAL SEARCH
# ==========================================
@mcp.tool()
def search_anatomy_diagrams(topic: str) -> str:
    """Searches Google Images (Backup)."""
    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()