# app.py
import gradio as gr
import os
import time
import json
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"), 
)

# Initialize ElevenLabs (We will only use this when the button is clicked)
eleven = ElevenLabs(api_key=os.getenv("ELEVENLABS_API_KEY"))
VOICE_ID = "JBFqnCBsd6RMkjVDRZzb" 

SYSTEM_PROMPT = """
You are Anato-Mitra, a strict Indian Medical College Professor.
- If the student asks a question, ALWAYS use the 'search_anatomy_diagrams' tool.
- Be concise and strict.
- The diagram will be shown automatically, so you just need to explain the clinical significance.
"""

# --- CORE LOGIC ---
async def run_agent(user_message, history):
    # Connect to MCP Server
    server_params = StdioServerParameters(
        command="uv", 
        args=["run", "anatomy_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}]
            messages.append({"role": "user", "content": user_message})

            # 1. First Call (Thinking)
            print("🧠 Thinking...")
            response = client.chat.completions.create(
                model="meta-llama/Meta-Llama-3.1-70B-Instruct",
                messages=messages,
                tools=openai_tools,
                tool_choice="auto"
            )

            final_response = response.choices[0].message.content or ""
            tool_image_markdown = "" # We will store the image here

            # 2. Tool Handling
            if 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
                
                print(f"🔧 Tool Call: {fn_name}")
                
                args_dict = json.loads(fn_args)
                result = await session.call_tool(fn_name, arguments=args_dict)
                tool_output = result.content[0].text
                
                # CAPTURE THE IMAGE MARKDOWN
                # We save this to append it manually later
                tool_image_markdown = f"\n\n---\n**Reference Diagram:**\n{tool_output}"
                
                # Inject result for LLM context
                messages.append({
                    "role": "user", 
                    "content": f"SYSTEM DATA: {tool_output}\n\nAnswer the student."
                })
                
                # 3. Final Synthesis
                final_response_obj = client.chat.completions.create(
                    model="meta-llama/Meta-Llama-3.1-70B-Instruct",
                    messages=messages
                )
                final_response = final_response_obj.choices[0].message.content

            # FORCE APPEND IMAGE (This fixes your missing diagram issue!)
            if tool_image_markdown:
                final_response += tool_image_markdown

            return final_response

# --- MANUAL AUDIO FUNCTION (Saves Credits!) ---
def generate_audio(chat_history):
    """Only generates audio for the LAST message when button is clicked"""
    if not chat_history:
        return None
        
    # Get the last message from the bot
    last_bot_message = chat_history[-1]['content']
    
    # Remove the markdown image links so the voice doesn't read "Image dot png"
    text_to_speak = last_bot_message.split("Reference Diagram:")[0] 
    
    print("🗣️ Generating Manual Audio...")
    try:
        audio_generator = eleven.text_to_speech.convert(
            text=text_to_speak,
            voice_id=VOICE_ID,
            model_id="eleven_multilingual_v2"
        )
        
        audio_path = f"voice_{int(time.time())}.mp3"
        with open(audio_path, "wb") as f:
            for chunk in audio_generator:
                f.write(chunk)
        return audio_path
    except Exception as e:
        print(f"Audio Error: {e}")
        return None

# --- UI ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🩻 Anato-Mitra: VIVA Companion")
    
    with gr.Row():
        # Left Side: Chat
        with gr.Column(scale=2):
            chatbot = gr.Chatbot(type="messages", height=500)
            msg = gr.Textbox(placeholder="Ask: 'Show me the Circle of Willis'")
            
        # Right Side: Controls
        with gr.Column(scale=1):
            gr.Markdown("### 🔊 VIVA Examiner Voice")
            audio_out = gr.Audio(label="Audio Output", type="filepath")
            # NEW BUTTON: Saves your credits!
            btn_speak = gr.Button("🔊 Speak Response", variant="primary")

    # Chat Logic
    async def respond(message, chat_history):
        bot_message = await run_agent(message, chat_history)
        chat_history.append({"role": "user", "content": message})
        chat_history.append({"role": "assistant", "content": bot_message})
        return "", chat_history

    msg.submit(respond, [msg, chatbot], [msg, chatbot])
    
    # Audio Logic (Only runs on click)
    btn_speak.click(generate_audio, inputs=chatbot, outputs=audio_out)

if __name__ == "__main__":
    demo.launch(share=True, auth=("waheed", "rehan123"))