import streamlit as st
from dotenv import load_dotenv
import os
from PIL import Image
import datetime


from brain_of_the_doctor import encode_image, analyze_image_with_query
from voice_of_the_patient import record_audio, transcribe_with_groq
from voice_of_the_doctor import text_to_speech_with_elevenlabs


load_dotenv()


system_prompt = """You have to act as a professional doctor, i know you are not but this is for learning purpose. 
            What's in this image?. Do you find anything wrong with it medically? 
            If you make a differential, suggest some remedies for them. Donot add any numbers or special characters in 
            your response. Your response should be in one long paragraph. Also always answer as if you are answering to a real person.
            Donot say 'In the image I see' but say 'With what I see, I think you have ....'
            Dont respond as an AI model in markdown, your answer should mimic that of an actual doctor not an AI bot, 
            Keep your answer concise (max 2 sentences). No preamble, start your answer right away please. Dont tell u shoud go to doctor. if user ask for medicine then suggest medicine also
            Dont say this : I cannot help you with that. I strongly advise against taking any medication without a prescription from a healthcare professional.
            Suggest meditions to user.
            """


def display_chat_message(speaker, text, timestamp):
    if speaker == "Patient":
        # User message (aligned to the right)
        st.markdown(
            f"""
            <div style="display: flex; justify-content: flex-end; margin-bottom: 10px;">
                <div style="background: #DCF8C6; padding: 10px; border-radius: 10px; max-width: 70%;">
                    <p style="margin: 0;">{text}</p>
                    <small style="color: gray; text-align: right;">{timestamp}</small>
                </div>
            </div>
            """,
            unsafe_allow_html=True,
        )
    else:
        # Doctor message (aligned to the left)
        st.markdown(
            f"""
            <div style="display: flex; justify-content: flex-start; margin-bottom: 10px;">
                <div style="background: #ECECEC; padding: 10px; border-radius: 10px; max-width: 70%;">
                    <p style="margin: 0;">{text}</p>
                    <small style="color: gray;">{timestamp}</small>
                </div>
            </div>
            """,
            unsafe_allow_html=True,
        )

def main():
    st.title("AI Doctor By NueSpaarx")

    
    if 'conversation' not in st.session_state:
        st.session_state.conversation = []

    # Display the chat interface
    st.header("Chat with AI Doctor")

    # Chat container
    chat_container = st.container()

    # Display the entire conversation history in chat format
    with chat_container:
        for entry in st.session_state.conversation:
            display_chat_message(entry["speaker"], entry["text"], entry["timestamp"])

    # Image input
    st.header("Upload an Image")
    image_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

    if image_file is not None:
        st.image(Image.open(image_file), caption="Uploaded Image", use_column_width=True)

    # Audio input
    st.header("Record Your Voice")
    audio_filepath = "patient_voice_test_for_patient.mp3"
    if st.button("Record Audio"):
        record_audio(file_path=audio_filepath)
        st.audio(audio_filepath, format="audio/mp3")

    if st.button("Analyze"):
        if image_file is not None:
            # Save the uploaded image to a temporary file
            image_filepath = "temp_image.jpg"
            with open(image_filepath, "wb") as f:
                f.write(image_file.getbuffer())

            # Process the inputs
            speech_to_text_output = transcribe_with_groq(
                GROQ_API_KEY=os.environ.get("GROQ_API_KEY"),
                audio_filepath=audio_filepath,
                stt_model="whisper-large-v3"
            )

            # Add user query to conversation history with timestamp
            st.session_state.conversation.append({
                "speaker": "Patient",
                "text": speech_to_text_output,
                "timestamp": datetime.datetime.now().strftime("%H:%M")
            })

            # Analyze the image and user query
            doctor_response = analyze_image_with_query(
                query=system_prompt + speech_to_text_output,
                encoded_image=encode_image(image_filepath),
                model="llama-3.2-11b-vision-preview"
            )

            # Add doctor's response to conversation history with timestamp
            st.session_state.conversation.append({
                "speaker": "Doctor",
                "text": doctor_response,
                "timestamp": datetime.datetime.now().strftime("%H:%M")
            })

            # Convert doctor's response to speech
            voice_of_doctor = text_to_speech_with_elevenlabs(
                input_text=doctor_response,
                output_filepath="final.mp3"
            )

            # Display the doctor's voice response
            st.header("Doctor's Voice")
            st.audio("final.mp3", format="audio/mp3")

            # Rerun the app to update the chat interface
            st.rerun()
        else:
            st.error("Please upload an image to analyze.")

if __name__ == "__main__":
    main()