import streamlit as st
import os
from groq import Groq
from langdetect import detect, DetectorFactory
import re

# Ensure consistent language detection
DetectorFactory.seed = 0

# Load API key from environment
GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
if not GROQ_API_KEY:
    st.error("❌ Missing GROQ API Key. Set 'GROQ_API_KEY' in environment variables.")

# Initialize Groq client
groq_client = Groq(api_key=GROQ_API_KEY)

# Define chatbot role as a specialized homeopathic doctor
CHATBOT_ROLE = """You are a highly skilled homeopathic doctor specializing in natural remedies. 
You analyze symptoms provided by users and suggest effective homeopathic treatments based on research-backed sources. 
Your response must include:
- Best possible remedy
- Proper dosage and administration
- Scientific or historical sources for reference
Ensure accuracy in medical details and maintain professionalism."""

# Improve language detection, especially for Roman Urdu
ROMAN_URDU_PATTERN = re.compile(r'\b(mujhe|kamar|dard|hai|aap|kyun|kaisi|bohot|zyada|masla|kar|raha|rahi|hun|ho|hai)\b', re.IGNORECASE)

def detect_language(text):
    """Detect the language of the input text, with special handling for Roman Urdu."""
    try:
        if ROMAN_URDU_PATTERN.search(text):
            return "roman_ur"  # Detect Roman Urdu correctly
        return detect(text)
    except:
        return "en"  # Default to English if detection fails

def get_homeopathic_advice(symptoms):
    """Search the web for remedies and generate a professional homeopathic response."""
    try:
        response = groq_client.chat.completions.create(
            model="llama3-70b-8192",
            messages=[
                {"role": "system", "content": CHATBOT_ROLE},
                {"role": "user", "content": f"Find the best homeopathic remedy for: {symptoms}. Provide remedy details, usage, dose, and sources."}
            ],
            temperature=0.3
        )
        return response.choices[0].message.content
    except Exception as e:
        return f"❌ Error fetching homeopathic advice: {str(e)}"

def translate_text(text, target_lang):
    """Translate text using Groq's LLaMA model."""
    try:
        if target_lang == "roman_ur":
            translation_prompt = f"Translate this medical advice into **Roman Urdu**, keeping it in **Latin script**: {text}"
        else:
            translation_prompt = f"Translate this medical advice into {target_lang}: {text}"

        response = groq_client.chat.completions.create(
            model="llama3-70b-8192",
            messages=[
                {"role": "system", "content": "You are a multilingual AI that accurately translates medical advice while preserving technical accuracy."},
                {"role": "user", "content": translation_prompt}
            ],
            temperature=0.2
        )
        return response.choices[0].message.content
    except Exception as e:
        return f"⚠️ Translation error: {str(e)}"

# 🎨 Streamlit UI
st.set_page_config(page_title="Homeo Doctor AI", page_icon="🌿")
st.title("🌍 AI-Powered Homeopathic Doctor")

# Chat interface
if "messages" not in st.session_state:
    st.session_state.messages = []

# Display chat history
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# User input box
if prompt := st.chat_input("Describe your symptoms..."):
    st.session_state.messages.append({"role": "user", "content": prompt})

    with st.spinner("🔍 Analyzing symptoms..."):
        # Detect user language
        detected_lang = detect_language(prompt)

        # Get homeopathic remedy
        homeo_advice = get_homeopathic_advice(prompt)

        # Translate advice if needed
        if detected_lang != "en":
            translated_advice = translate_text(homeo_advice, detected_lang)
        else:
            translated_advice = homeo_advice

        # Format final response
        final_response = f"""
        **💡 Expert Homeopathic Advice (English):**  
        {homeo_advice}  
        
        **🌍 Translated Advice ({'Roman Urdu' if detected_lang == 'roman_ur' else detected_lang.upper()}):**  
        {translated_advice}
        """

    # Display response
    with st.chat_message("assistant"):
        st.markdown(final_response)

    # Save chat history
    st.session_state.messages.append({"role": "assistant", "content": final_response})

# Disclaimer
st.caption("⚠️ This AI does not replace professional medical consultation.")