import os
import pickle
import streamlit as st
import speech_recognition as sr
import pyttsx3

try:
    import torch
    from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
except ModuleNotFoundError as e:
    st.error(f"❌ Missing dependency: {e}. Please install required packages.")
    st.stop()

# ==============================
# 1) Model Configuration
# ==============================
MODEL_DIR = "flan-t5-small"

torch_dtype = torch.float32 if torch.cuda.is_available() else torch.float32

def load_model():
    """Load the FLAN-T5 model from local storage."""
    st.write("🚀 Loading FLAN-T5 model from local storage...")
    try:
        tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR, local_files_only=True)
        model = AutoModelForSeq2SeqLM.from_pretrained(
            MODEL_DIR,
            torch_dtype=torch_dtype,
        )
        st.write("✅ Model loaded successfully from local storage!")
        return tokenizer, model
    except Exception as e:
        st.error(f"❌ Model failed to load: {e}")
        st.stop()

# ==============================
# 2) Initialize Streamlit UI
# ==============================
st.title("🩺 Healthcare Chatbot (FLAN-T5)")

# Load model
try:
    tokenizer, model = load_model()
except Exception as e:
    st.error(f"❌ Model load error: {e}")
    st.stop()

# ==============================
# 3) Load Medical Q&A Data
# ==============================
try:
    st.write("📂 Loading medical Q&A data...")
    with open("train_data.pkl", "rb") as file:
        medical_qna = pickle.load(file)
    st.write("✅ Q&A data loaded!")
except FileNotFoundError:
    st.error("❌ 'train_data.pkl' not found. Please ensure it exists.")
    st.stop()
except Exception as e:
    st.error(f"❌ Failed to load Q&A data: {e}")
    st.stop()

# ==============================
# 4) Chatbot Response Logic
# ==============================
def chatbot_response(user_input: str) -> str:
    """
    1. Check if user_input matches any question in medical_qna.
    2. If so, return that pre-written answer.
    3. Otherwise, feed FLAN-T5 a prompt to generate a medical response.
    """

    # Check Q&A first
    for qa in medical_qna:
        if user_input.lower() in qa["question"].lower():
            return qa["answer"]

    # System-style prompt for medical answers
    prompt = (
        "You are a helpful medical assistant. The user asked:\n"
        f"Question: {user_input}\n\n"
        "Answer in a concise, accurate way. If you're unsure, advise seeing a doctor."
    )

    # Encode
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, padding=True)
    # Generate
    outputs = model.generate(
        **inputs,
        max_length=256,
        num_beams=2,  # for more deterministic, coherent responses
        no_repeat_ngram_size=2
    )
    # Decode
    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)

    return answer

# ==============================
# 5) Text-to-Speech
# ==============================
tts_engine = pyttsx3.init()

def speak(text: str):
    """Speak text aloud."""
    tts_engine.say(text)
    tts_engine.runAndWait()

# ==============================
# 6) Streamlit UI Logic
# ==============================
if st.button("What can you help me with?"):
    st.write(
        "I can provide general information about medical symptoms, treatments, "
        "and offer guidance. If you have serious concerns, please contact a doctor."
    )

user_input = st.text_input("Ask me a medical question:")
if st.button("Get Answer"):
    if user_input.strip():
        response = chatbot_response(user_input)
        st.write(f"**Bot:** {response}")
        speak(response)
    else:
        st.warning("Please enter a question.")

if st.button("🎙 Use Voice Input"):
    recognizer = sr.Recognizer()
    with sr.Microphone() as source:
        st.write("🎤 Listening...")
        audio = recognizer.listen(source)
    try:
        voice_input = recognizer.recognize_google(audio)
        st.write(f"🗣 You said: {voice_input}")
        response = chatbot_response(voice_input)
        st.write(f"**Bot:** {response}")
        speak(response)
    except sr.UnknownValueError:
        st.write("❌ Could not understand audio.")
    except sr.RequestError:
        st.write("❌ Speech recognition service error.")