import gradio as gr
from gtts import gTTS
import tempfile
from transformers import pipeline
import os
import logging
from huggingface_hub import login

# Set up logging for better error tracking
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Configure Hugging Face API Token (Replace with your actual token)
HUGGINGFACE_API_TOKEN = "your_huggingface_api_token"  # Replace with your Hugging Face API token

# Validate API token and log in to Hugging Face Hub
if HUGGINGFACE_API_TOKEN == "your_huggingface_api_token":
    logger.error("Please replace 'your_huggingface_api_token' with a valid Hugging Face API token from https://huggingface.co/settings/tokens")
    raise ValueError("Invalid Hugging Face API token. Please set a valid token.")
try:
    login(token=HUGGINGFACE_API_TOKEN)
    logger.info("Successfully logged in to Hugging Face Hub")
except Exception as e:
    logger.error(f"Failed to log in to Hugging Face Hub: {str(e)}")
    raise

# Initialize Hugging Face Whisper model for speech-to-text (multilingual support)
try:
    whisper_pipeline = pipeline(
        "automatic-speech-recognition",
        model="openai/whisper-tiny",
        token=HUGGINGFACE_API_TOKEN,
        device=-1  # Use CPU if GPU unavailable
    )
    logger.info("Whisper model initialized successfully")
except Exception as e:
    logger.error(f"Failed to initialize Whisper model: {str(e)}")
    raise

# Initialize Hugging Face text generation model for chatbot
try:
    text_generation_pipeline = pipeline(
        "text2text-generation",
        model="google/flan-t5-small",
        token=HUGGINGFACE_API_TOKEN,
        device=-1  # Use CPU if GPU unavailable
    )
    logger.info("Text generation model initialized successfully")
except Exception as e:
    logger.error(f"Failed to initialize text generation model: {str(e)}")
    raise

# Define chatbot knowledge base
base_info = """
You are a highly advanced AI assistant named 'MindCare'.
Your role is to provide support in various aspects of health and well-being, including:
- **Mental health**: Emotional support, mindfulness, stress-relief exercises, anxiety management.
- **Medical guidance**: Basic symptom analysis, possible conditions, and medicine recommendations.
- **Decision-making support**: Helping users with personal, professional, and emotional choices.
- **General health advice**: Lifestyle improvements, nutrition, physical wellness, and mental well-being.
- **Emergency assistance**: If the user is in distress, suggest professional help or helpline numbers.

Your tone is always **empathetic, supportive, and informative**. You ensure users feel heard and cared for.
"""

mental_health = """
If the user is feeling stressed or anxious:
- Suggest mindfulness exercises, deep breathing techniques, or gratitude journaling.
- Encourage taking breaks, engaging in hobbies, and spending time in nature.
- Provide positive affirmations and self-care routines.

If the user is in distress:
- Offer emotional support and let them know they are not alone.
- Encourage them to reach out to a trusted person or professional.
- Provide emergency helpline numbers if needed.
"""

medical_assistance = """
If the user provides symptoms:
- Analyze symptoms and suggest possible conditions.
- Provide general advice but **never** replace a doctor’s consultation.
- Suggest lifestyle changes or basic home remedies if applicable.
- If symptoms are severe, advise them to visit a healthcare professional.

If the user asks about medicines:
- Suggest **common antibiotics** based on infection type (e.g., Amoxicillin for bacterial infections).
- Recommend **painkillers** like Paracetamol, Ibuprofen, or Diclofenac for pain relief.
- Mention precautions and possible side effects.
- Clearly **state that a doctor’s consultation is necessary before taking any medicine**.
"""

medicine_recommendation = """
If the user asks for a prescription, provide general guidance on **commonly used medicines**:
- **Antibiotics** (for bacterial infections): Amoxicillin, Azithromycin, Ciprofloxacin.
- **Painkillers**: Paracetamol (mild pain/fever), Ibuprofen (anti-inflammatory), Diclofenac (muscle pain).
- **Cold & Flu**: Antihistamines like Cetirizine, Cough syrups like Dextromethorphan.
- **Stomach Issues**: Antacids like Ranitidine, PPI like Omeprazole.

Always remind the user that **only a licensed doctor can prescribe medicines, and misuse can be harmful**.
"""

decision_guidance = """
If the user is struggling with a decision:
- Help them weigh pros and cons logically.
- Suggest considering their values, long-term goals, and emotions.
- Provide structured approaches like decision matrices or intuitive checks.
- Encourage seeking advice from trusted people if needed.
"""

emergency_help = """
If the user mentions severe mental distress:
- Respond with immediate emotional support.
- Provide crisis helpline numbers (if applicable to the region).
- Encourage talking to a trusted friend, family member, or professional.
- Remind them that they are not alone and help is available.
"""

# Combine all knowledge into one structured context
context = [base_info, mental_health, medical_assistance, medicine_recommendation, decision_guidance, emergency_help]

# Function to get AI response using text generation model
def get_llm_response(message):
    try:
        full_context = "\n".join(context) + f"\nUser: {message}\nMindCare:"
        response = text_generation_pipeline(full_context, max_length=500, num_return_sequences=1)[0]['generated_text']
        return response
    except Exception as e:
        logger.error(f"Error generating response: {str(e)}")
        return f"Error generating response: {str(e)}"

# Function to process voice input and convert to text
def process_voice_input(audio_file, language="en"):
    # Map language codes for Whisper model
    language_map = {
        "en": "english",
        "es": "spanish",
        "hi": "hindi",
        "zh": "chinese"
    }
    try:
        # Transcribe audio using Hugging Face Whisper model
        transcription = whisper_pipeline(audio_file, 
                                        generate_kwargs={"language": language_map.get(language, "english")})
        return transcription["text"]
    except Exception as e:
        logger.error(f"Error processing audio: {str(e)}")
        return f"Error processing audio: {str(e)}"

# Define chatbot response function with voice and text input
def bot(message=None, audio_file=None, language="en", history=None):
    try:
        if audio_file:
            # Process voice input if provided
            message = process_voice_input(audio_file, language)
        
        if not message:
            return "No input provided. Please type a message or upload an audio file.", None
        
        # Get response from text generation model
        response = get_llm_response(message)
        
        # Convert the response to speech using gTTS
        tts = gTTS(text=response, lang=language, slow=False)
        
        # Save the audio file in a temporary directory
        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as temp_audio:
            audio_path = temp_audio.name
            tts.save(audio_path)
        
        return response, audio_path
    except Exception as e:
        logger.error(f"Error in bot function: {str(e)}")
        return f"Error: {str(e)}", None

# Create Gradio interface with voice and text input
demo = gr.Interface(
    fn=bot,
    inputs=[
        gr.Textbox(label="Enter your message (optional)"),
        gr.Audio(label="Record your voice", type="filepath"),
        gr.Dropdown(label="Select Language", choices=["en", "es", "hi", "zh"], value="en")
    ],
    outputs=[
        gr.Textbox(label="Response"),
        gr.Audio(label="Response Audio")
    ],
    title="MindCare - Your Personal Health & Wellness Assistant",
    description="Type a message or record your voice in English, Spanish, Hindi, or Mandarin."
)

# Launch the Gradio app
if __name__ == "__main__":
    try:
        demo.launch(debug=True, share=True)
    except Exception as e:
        logger.error(f"Failed to launch Gradio app: {str(e)}")
        raise