import gradio as gr

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import torchaudio
from transformers import Wav2Vec2ForCTC, Wav2Vec2Tokenizer
from TTS.api import TTS






# Device configuration
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")








# Load language model
language_model_name = "microsoft/BioGPT-large"  # Example model
tokenizer = AutoTokenizer.from_pretrained(language_model_name)
language_model = AutoModelForCausalLM.from_pretrained(language_model_name).to(device)

# Load STT model
stt_model_name = "facebook/wav2vec2-base-960h"
stt_tokenizer = Wav2Vec2Tokenizer.from_pretrained(stt_model_name)
stt_model = Wav2Vec2ForCTC.from_pretrained(stt_model_name).to(device)

# Initialize TTS (optional if using browser-based TTS)
# tts = TTS(model_name="tts_models/en/ljspeech/tacotron2-DDC", progress_bar=False, gpu=torch.cuda.is_available())

def transcribe_audio(audio_path):
    waveform, sample_rate = torchaudio.load(audio_path)

    
    # Resample to 16000 if necessary
    if sample_rate != 16000:
        resampler = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)
        waveform = resampler(waveform)















    
    input_values = stt_tokenizer(waveform.squeeze().numpy(), return_tensors="pt", sampling_rate=16000).input_values.to(device)
    
    with torch.no_grad():
        logits = stt_model(input_values).logits
    
    predicted_ids = torch.argmax(logits, dim=-1)
    transcription = stt_tokenizer.decode(predicted_ids[0])
    return transcription.lower()

def get_bot_response(user_input):
    # Prepare the input
    input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors="pt").to(device)
    
    # Generate a response
    chat_history_ids = language_model.generate(
        input_ids,
        max_length=500,
        pad_token_id=tokenizer.eos_token_id,
        no_repeat_ngram_size=3,
        do_sample=True,
        top_k=50,
        top_p=0.95,
        temperature=0.7
    )
    
    # Decode the response
    response = tokenizer.decode(chat_history_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
    return response

def chatbot_interface(audio):
    if audio is None:
        return "Please provide an audio input."
    
    try:
        # Step 1: Transcribe audio to text
        user_input = transcribe_audio(audio)
    except Exception as e:
        return f"Error in transcription: {str(e)}"
    
    if not user_input:
        return "Sorry, I could not understand the audio."
    
    # Step 2: Get bot response
    try:
        bot_response = get_bot_response(user_input)
    except Exception as e:
        return f"Error in generating response: {str(e)}"
    
    # Step 3: Return response with browser-based TTS
    response_html = f"""
    <p><strong>You:</strong> {user_input}</p>
    <p><strong>Bot:</strong> {bot_response}</p>
    <button onclick="speakBotResponse()">🗣️ Hear Response</button>
    <script>
    function speakBotResponse() {{
        var msg = new SpeechSynthesisUtterance("{bot_response}");
        window.speechSynthesis.speak(msg);
    }}
    </script>
    """
    return response_html

# Create Gradio interface
demo = gr.Interface(
    fn=chatbot_interface,
    inputs=gr.Audio(source="microphone", type="filepath"),
    outputs=gr.HTML(),
    title="Healthcare AI Chatbot",
    description="Ask health-related questions. Please consult a professional for medical advice.",
    examples=None,
    allow_flagging="never",












)

# Launch the interface
demo.launch()


Traceback (most recent call last):
  File "/home/user/app/app.py", line 126, in <module>
    inputs=gr.Audio(source="microphone", type="filepath"),
  File "/usr/local/lib/python3.10/site-packages/gradio/component_meta.py", line 163, in wrapper
    return fn(self, **kwargs)
TypeError: Audio.__init__() got an unexpected keyword argument 'source'