import streamlit as st
import requests
from geopy.geocoders import Nominatim
import whisper
import tempfile
from pydub import AudioSegment
from io import BytesIO
from streamlit_js_eval import streamlit_js_eval

# Set your Hugging Face API URL and API key
API_URL = "https://api-inference.huggingface.co/models/dmis-lab/biobert-base-cased-v1.1"
headers = {"Authorization": f"secret"}

# Initialize Whisper model
whisper_model = whisper.load_model("base")

# Function to query the Hugging Face model
def query(payload):
    response = requests.post(API_URL, headers=headers, json=payload)
    if response.status_code == 200:
        return response.json()
    else:
        st.error(f"Error: Unable to fetch response from model (status code: {response.status_code})")
        st.error(response.text)
        return None

# Function to find nearby clinics/pharmacies using geopy
def find_nearby_clinics(address):
    geolocator = Nominatim(user_agent="healthcare_companion")
    location = geolocator.geocode(address)
    if location:
        return (location.latitude, location.longitude)
    else:
        st.error("Error: Address not found")
        return None

# Function to transcribe audio to text using Whisper
def transcribe_audio(audio_bytes):
    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_audio_file:
        audio = AudioSegment.from_file(BytesIO(audio_bytes), format="wav")
        audio.export(temp_audio_file.name, format="wav")
        result = whisper_model.transcribe(temp_audio_file.name)
        return result["text"]

# Main function to create the Streamlit app
def main():
    st.title("Healthcare Companion")
    st.write("This app provides healthcare guidance, prescription information, and locates nearby clinics or pharmacies.")

    # JavaScript code to capture audio
    js_code = """
    async function recordAudio() {
        const stream = await navigator.mediaDevices.getUserMedia({ audio: true });
        const mediaRecorder = new MediaRecorder(stream);
        let audioChunks = [];
        
        mediaRecorder.ondataavailable = event => {
            audioChunks.push(event.data);
        };
        
        mediaRecorder.onstop = async () => {
            const audioBlob = new Blob(audioChunks, { type: 'audio/wav' });
            const audioBuffer = await audioBlob.arrayBuffer();
            const audioBase64 = arrayBufferToBase64(audioBuffer);
            document.getElementById('audio_data').value = audioBase64;
            document.getElementById('audio_form').submit();
        };
        
        mediaRecorder.start();
        setTimeout(() => mediaRecorder.stop(), 5000);  // Record for 5 seconds
        
        function arrayBufferToBase64(buffer) {
            let binary = '';
            const bytes = new Uint8Array(buffer);
            const len = bytes.byteLength;
            for (let i = 0; i < len; i++) {
                binary += String.fromCharCode(bytes[i]);
            }
            return window.btoa(binary);
        }
    }
    
    recordAudio();
    """
    
    # Placeholder for audio data
    st_js_code = streamlit_js_eval(js_code, key="record_audio")
    
    # Form to receive audio data from JavaScript
    with st.form("audio_form", clear_on_submit=True):
        audio_data = st.text_input("audio_data", type="hidden")
        submit_button = st.form_submit_button("Submit")
        
        if submit_button and audio_data:
            audio_bytes = BytesIO(base64.b64decode(audio_data))
            symptoms = transcribe_audio(audio_bytes)
            st.write(f"Transcribed symptoms: {symptoms}")
    
    if 'symptoms' in locals() and symptoms:
        context = """
        This is a healthcare question and answer platform. The following text contains typical symptoms, treatments, and medical conditions commonly asked about in healthcare settings.
        For example, symptoms of COVID-19 include fever, dry cough, and tiredness. Treatment options for hypertension include lifestyle changes and medications. The platform is designed to assist with general medical inquiries.
        """
        payload = {"inputs": {"question": symptoms, "context": context}}
        st.write(f"Debug: Payload sent to model: {payload}")  # Debugging: Check payload
        result = query(payload)
        st.write(f"Debug: Response from model: {result}")  # Debugging: Check response
        if result:
            st.write("**Medical Advice:**")
            # Check the response structure and extract the answer appropriately
            answer = result.get('answer') if 'answer' in result else "Sorry, I don't have information on that."
            st.write(answer)

    # User input for address to find nearby clinics/pharmacies
    address = st.text_input("Enter your address to find nearby clinics/pharmacies:")
    if address:
        location = find_nearby_clinics(address)
        if location:
            st.write(f"**Nearby Clinics/Pharmacies (Coordinates):** {location}")

    # Additional features like prescription info can be added similarly

if __name__ == "__main__":
    main()