import os
import gradio as gr
from llama_index.llms.huggingface import HuggingFaceLLM
from parse_tabular import symptom_index
import json

# --- System prompt ---
SYSTEM_PROMPT = """
You are a medical assistant helping a user narrow down to the most likely ICD-10 code.
At each turn, EITHER ask one focused clarifying question (e.g. "Is your cough dry or productive?")
or, if you have enough info, output a final JSON with fields:
{"diagnoses":[…], "confidences":[…]}.
"""

def process_speech(new_transcript, history):
    # Skip if no new transcript
    if not new_transcript:
        return history
    
    try:
        # Build conversation context
        context = "\n".join([f"{role}: {msg}" for role, msg in history])
        
        # Query symptom index for relevant codes
        response = symptom_index.as_query_engine().query(new_transcript)
        
        # Format response as structured JSON
        formatted_response = {
            "diagnoses": [str(response).split(":")[0]],  # Extract ICD code
            "confidences": [0.8],  # Add confidence scoring
            "follow_up": "Is the cough productive or dry?"  # Add interactive questioning
        }
        
        # Append exchange to history
        history.append((new_transcript, json.dumps(formatted_response, indent=2)))
        
    except Exception as e:
        # Handle errors gracefully for MCP clients
        error_response = {
            "error": str(e),
            "status": "error"
        }
        history.append((new_transcript, json.dumps(error_response, indent=2)))
    
    return history

# Build Gradio interface
demo = gr.Blocks()
with demo:
    gr.Markdown("# Symptom to ICD-10 Code Lookup (Audio Input)")
    chatbot = gr.Chatbot(label="Conversation")
    audio = gr.Audio(source="microphone", type="text", streaming=True)
    
    # Add MCP-specific metadata
    demo.config = {
        "mcp": {
            "title": "Medical Symptom to ICD-10 Code Assistant",
            "description": "Convert spoken medical symptoms to ICD-10 codes",
            "version": "1.0.0",
            "capabilities": {
                "speech_input": True,
                "streaming": True
            }
        }
    }
    
    audio.stream(
        process_speech,
        inputs=[audio, chatbot], 
        outputs=chatbot,
        show_progress="hidden"
    )

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0", 
        server_port=7860, 
        mcp_server=True,
        mcp_polling_interval=1000  # 1 second polling interval
    )