import os
import streamlit as st
import google.generativeai as genai
from dotenv import load_dotenv

# Load environment variables if running locally
load_dotenv()

# Streamlit web app layout
st.title("Gemini Model Suggestions Generator")
api_key = st.sidebar.text_input("Enter your GEMINI API Key", type="password")

if api_key:
    # Configure Gemini with the user-provided API key
    genai.configure(api_key=api_key)

    # Gemini model configuration
    MODEL_NAME = 'gemini-1.5-pro'
    SAFETY_SETTINGS = {
        'HATE': 'BLOCK_NONE',
        'HARASSMENT': 'BLOCK_NONE',
        'SEXUAL': 'BLOCK_NONE',
        'DANGEROUS': 'BLOCK_NONE'
    }

    # Function to generate suggestions
    def generate_suggestions(free_text: str, patient_info: str, instructions: list) -> list:
        responses = []
        configs = [
            {'instruction': instructions[0], 'temperature': 0.3, 'max_tokens': 1000},
            {'instruction': instructions[1], 'temperature': 0.7, 'max_tokens': 1000},
            {'instruction': instructions[2], 'temperature': 1.0, 'max_tokens': 1000}
        ]

        # Create model instance
        model = genai.GenerativeModel(MODEL_NAME)

        for config in configs:
            try:
                # Construct the full prompt
                full_prompt = f"""
                [Free Text]: {free_text}
                [Patient Information]: {patient_info}
                [Instruction]: {config['instruction']}
                Please provide your response below:
                """

                # Generate content
                response = model.generate_content(
                    contents=full_prompt,
                    generation_config=genai.types.GenerationConfig(
                        temperature=config['temperature'],
                        max_output_tokens=config['max_tokens']
                    ),
                    safety_settings=SAFETY_SETTINGS
                )

                responses.append({
                    'success': True,
                    'instruction': config['instruction'],
                    'temperature': config['temperature'],
                    'response': response.text
                })

            except Exception as e:
                responses.append({
                    'success': False,
                    'error': str(e),
                    'instruction': config['instruction'],
                    'temperature': config['temperature']
                })

        return responses

    # Example usage in Streamlit
    if st.button("Generate Suggestions"):
        free_text = st.text_area("Enter Free Text Information")
        patient_info = st.text_area("Enter Patient Information")
        instructions = [
            "a) Provide possible diagnoses in bullet points, b) Suggest lifestyle modifications, c) Recommend diagnostic tests",
            "a) Provide possible course of action for physio in bullet points, b) Suggest lifestyle modifications that can be achieved by the patient",
            "a) General preventive advice for the patient through physio routines, b) Suggest lifestyle modifications that can be achieved by the patient"
        ]

        results = generate_suggestions(free_text, patient_info, instructions)

        st.write("Generated Suggestions:")
        for idx, result in enumerate(results, 1):
            st.markdown(f"### Suggestion {idx}:")
            st.write(f"Temperature: {result['temperature']}")
            st.write(f"Instruction: {result['instruction']}")

            if result['success']:
                st.write(f"Response:\n{result['response']}")
            else:
                st.error(f"Error: {result['error']}")

else:
    st.sidebar.warning("Please enter a valid API key.")