app.py

import re
import gradio as gr
import os
import time
import json
import sys
import requests 
import shutil   
import speech_recognition as sr
from openai import OpenAI
from dotenv import load_dotenv
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from elevenlabs.client import ElevenLabs

load_dotenv()

# --- CONFIGURATION ---
client = OpenAI(
    base_url="https://api.hyperbolic.xyz/v1", 
    api_key=os.getenv("HYPERBOLIC_API_KEY"), 
)

eleven = ElevenLabs(api_key=os.getenv("ELEVENLABS_API_KEY"))
VOICE_ID = "JBFqnCBsd6RMkjVDRZzb" 

# --- PROMPTS (INTERNET ONLY) ---

PROMPT_LEARN = """
You are Anato-Mitra, a Medical Researcher using the Web.
- Source: The Internet.
- Tool: ALWAYS use 'search_anatomy_diagrams'.
- Provide broad, clinical information.
"""

PROMPT_VIVA_ASK = """
You are Anato-Mitra, an Examiner.
- Action: 
  1. Use 'get_random_anatomy_topic' to pick a subject.
  2. Use 'search_anatomy_diagrams' to find an image.
- Output: Show the image and ask: "Identify this structure."
- Constraint: Do NOT grade previous answers. Just ask the question.
"""

PROMPT_VIVA_GRADE = """
You are Anato-Mitra, a Medical Researcher.
- Task: Verify the user's answer.
- Format:
  1. Start with "🎉 CORRECT" or "❌ INCORRECT".
  2. If INCORRECT, explicitly state: "The correct structure is [NAME]".
  3. Provide a brief clinical fact.
- Constraint: Do NOT generate a new image. Stop after the explanation.
"""

# --- HELPER: BYPASS 403 BLOCKS (Web Images) ---
def get_safe_image_path(url_or_path):
    if not url_or_path: return None
    
    # If it's a local path (starts with extracted_images or /gradio_api...)
    if not url_or_path.startswith("http"):
        # Clean it up to get the actual file system path
        clean = url_or_path.replace("/gradio_api/file=", "").replace("/file=", "")
        if os.path.exists(clean):
            return os.path.abspath(clean)
        return None

    # If it's a web URL, download it
    try:
        headers = {'User-Agent': 'Mozilla/5.0'}
        response = requests.get(url_or_path, headers=headers, stream=True, timeout=5)
        if response.status_code == 200:
            ext = ".svg" if ".svg" in url_or_path else ".png"
            temp_filename = f"temp_downloaded_img{ext}"
            with open(temp_filename, 'wb') as f:
                response.raw.decode_content = True
                shutil.copyfileobj(response.raw, f)
            return os.path.abspath(temp_filename)
    except: return None
    return None

# --- HELPER: FIX IMAGE PATHS ---
def fix_gradio_images(text):
    """
    Forces Gradio to load images via the API route.
    """
    if not text: return ""
    return text.replace("(extracted_images/", "(/gradio_api/file=extracted_images/")

# --- AGENT LOGIC ---
async def run_agent(user_message, history, custom_prompt=None, force_no_tools=False):
    
    # Default to Learning prompt if no custom prompt provided
    system_prompt = custom_prompt if custom_prompt else PROMPT_LEARN
    
    server_params = StdioServerParameters(
        command=sys.executable, 
        args=["super_server.py"], 
    )

    async with stdio_client(server_params) as (read, write):
        async with ClientSession(read, write) as session:
            await session.initialize()
            tools_list = await session.list_tools()
            
            openai_tools = [{
                "type": "function",
                "function": {
                    "name": t.name,
                    "description": t.description,
                    "parameters": t.inputSchema
                }
            } for t in tools_list.tools]

            messages = [{"role": "system", "content": system_prompt}]
            for msg in history:
                messages.append({"role": "user" if msg['role'] == "user" else "assistant", "content": msg['content']})
            messages.append({"role": "user", "content": user_message})

            print(f"🧠 Thinking...")
            
            # LOGIC FIX: If force_no_tools is True, we tell OpenAI "none"
            current_tool_choice = "none" if force_no_tools else "auto"

            response = client.chat.completions.create(
                model="meta-llama/Meta-Llama-3.1-70B-Instruct",
                messages=messages,
                tools=openai_tools,
                tool_choice=current_tool_choice 
            )

            final_response = response.choices[0].message.content or ""
            
            # Only process tools if we allowed them
            if not force_no_tools and response.choices[0].message.tool_calls:
                tool_call = response.choices[0].message.tool_calls[0]
                fn_name = tool_call.function.name
                fn_args = tool_call.function.arguments
                
                args_dict = json.loads(fn_args)
                result = await session.call_tool(fn_name, arguments=args_dict)
                tool_output = result.content[0].text
                
                messages.append({
                    "role": "user", 
                    "content": f"SYSTEM DATA from '{fn_name}': {tool_output}\n\nProceed."
                })
                
                final_response_obj = client.chat.completions.create(
                    model="meta-llama/Meta-Llama-3.1-70B-Instruct",
                    messages=messages
                )
                final_response_w_tool = final_response_obj.choices[0].message.content
                
                if "![" in tool_output and "![" not in final_response_w_tool:
                    final_response_w_tool += f"\n\n{tool_output}"
                
                final_response = final_response_w_tool

            # Apply path fix
            final_response = fix_gradio_images(final_response)

            extracted_image_url = None
            img_match = re.search(r"!\[.*?\]\((.*?)\)", final_response)
            if img_match:
                extracted_image_url = img_match.group(1)

            return final_response, extracted_image_url

# --- AUDIO & STT ---
def generate_audio(chat_history):
    if not chat_history: return None
    last_msg = chat_history[-1]['content']
    clean = re.sub(r'!\[.*?\]\(.*?\)', '', last_msg)
    clean = re.sub(r'https?://\S+', '', clean).replace("🎉", "").replace("❌", "").replace("/gradio_api/file=", "").strip()
    try:
        gen = eleven.text_to_speech.convert(text=clean, voice_id=VOICE_ID, model_id="eleven_multilingual_v2")
        path = f"voice_{int(time.time())}.mp3"
        with open(path, "wb") as f:
            for chunk in gen: f.write(chunk)
        return path
    except: return None

def transcribe_audio(audio_path):
    if not audio_path: return ""
    rec = sr.Recognizer()
    try:
        with sr.AudioFile(audio_path) as s: return rec.recognize_google(rec.record(s))
    except: return ""

# --- UI ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    
    gr.Markdown("# 🩻 Anato-Mitra: AI Medical Companion")
    
    # === VIEW 1: MODE SELECTION (Starts Here) ===
    with gr.Column(visible=True) as view_mode:
        gr.Markdown("### Select Activity")
        with gr.Row():
            btn_mode_learn = gr.Button("🎓 Learning Mode", variant="secondary", size="lg")
            btn_mode_viva = gr.Button("🎲 VIVA Mode", variant="primary", size="lg")

    # === VIEW 2: LEARNING MODE ===
    with gr.Column(visible=False) as view_learn:
        lbl_learn_header = gr.Markdown("### 🎓 Learning Mode (Internet)")
        with gr.Row():
            with gr.Column(scale=1):
                chat_learn = gr.Chatbot(type="messages", height=500)
                msg_learn = gr.Textbox(placeholder="Ask about any anatomy topic (e.g., 'Show me the Heart')...")
                with gr.Row():
                    btn_speak_learn = gr.Button("🔊 Speak")
                    btn_home_learn = gr.Button("🏠 Main Menu")
                aud_learn = gr.Audio(visible=False, autoplay=True)
            with gr.Column(scale=1):
                img_learn = gr.Image(label="Diagram", interactive=False, height=500)

    # === VIEW 3: VIVA MODE ===
    with gr.Column(visible=False) as view_viva:
        lbl_viva_header = gr.Markdown("### 🎲 VIVA Mode (Internet)")
        chat_viva = gr.Chatbot(type="messages", height=500)
        with gr.Row():
            msg_viva = gr.Textbox(placeholder="Answer here...", scale=4)
            mic_viva = gr.Audio(sources=["microphone"], type="filepath", scale=1)
        with gr.Row():
            btn_speak_viva = gr.Button("🔊 Speak")
            btn_next_viva = gr.Button("🎲 Start / Next Question ➡️", variant="primary")
            btn_home_viva = gr.Button("🏠 Main Menu")
        aud_viva = gr.Audio(visible=False, autoplay=True)

    # --- NAVIGATION LOGIC ---
    def go_home():
        return gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)

    def select_learn(): 
        return gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)

    def select_viva(): 
        return gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)

    # --- EVENT WIRING ---
    demo.load(go_home, outputs=[view_mode, view_learn, view_viva])
    
    btn_mode_learn.click(select_learn, outputs=[view_mode, view_learn, view_viva])
    btn_mode_viva.click(select_viva, outputs=[view_mode, view_learn, view_viva])
    
    btn_home_learn.click(go_home, outputs=[view_mode, view_learn, view_viva])
    btn_home_viva.click(go_home, outputs=[view_mode, view_learn, view_viva])

    # --- CHAT INTERACTIONS: LEARNING ---
    async def run_learn(msg, hist):
        # Uses Default PROMPT_LEARN (Internet)
        txt, url = await run_agent(msg, hist, custom_prompt=None)
        hist.append({"role": "user", "content": msg})
        hist.append({"role": "assistant", "content": txt})
        safe_path = get_safe_image_path(url)
        return "", hist, safe_path if safe_path else gr.update()

    msg_learn.submit(run_learn, [msg_learn, chat_learn], [msg_learn, chat_learn, img_learn])
    btn_speak_learn.click(generate_audio, chat_learn, aud_learn)

    # --- CHAT INTERACTIONS: VIVA (SPLIT LOGIC) ---
    
    # 1. ASK (Start / Next) -> ALLOW TOOLS (force_no_tools=False)
    async def run_viva_next_question(hist):
        # Uses PROMPT_VIVA_ASK (Internet)
        prompt = PROMPT_VIVA_ASK
        
        # force_no_tools=False because we NEED the image here
        txt, url = await run_agent("Generate next question", hist, custom_prompt=prompt, force_no_tools=False)
        
        final = txt
        if url and "![" not in txt:
             final += f"\n\n![Quiz Image](/file={url})" if not url.startswith("/file=") else f"\n\n![Quiz Image]({url})"
        
        hist.append({"role": "assistant", "content": final})
        return hist

    # 2. GRADE (Submit Answer) -> BLOCK TOOLS (force_no_tools=True)
    async def run_viva_grade_answer(msg, hist):
        # Uses PROMPT_VIVA_GRADE (Internet)
        prompt = PROMPT_VIVA_GRADE
        
        # force_no_tools=True prevents fetching new images
        txt, _ = await run_agent(msg, hist, custom_prompt=prompt, force_no_tools=True)
        
        hist.append({"role": "user", "content": msg})
        hist.append({"role": "assistant", "content": txt})
        return "", hist

    # Viva Wiring
    msg_viva.submit(run_viva_grade_answer, [msg_viva, chat_viva], [msg_viva, chat_viva])
    btn_next_viva.click(run_viva_next_question, [chat_viva], [chat_viva])
    
    btn_speak_viva.click(generate_audio, chat_viva, aud_viva)
    mic_viva.stop_recording(transcribe_audio, mic_viva, msg_viva).then(run_viva_grade_answer, [msg_viva, chat_viva], [msg_viva, chat_viva])

if __name__ == "__main__":
    cwd = os.path.abspath(".")
    # Ensure specific folders exist or allow generic temp access
    allowed = [cwd, "/tmp"]
    if os.path.exists("extracted_images"): allowed.append(os.path.abspath("extracted_images"))
    
    demo.launch(share=True, allowed_paths=allowed)