app.py

import streamlit as st
import time
import requests
from streamlit.components.v1 import html
import os
from dotenv import load_dotenv
import numpy as np
import torchaudio
from audio_recorder_streamlit import audio_recorder
import torch
from io import BytesIO
import hashlib

# Load Whisper model (cached)
@st.cache_resource
def load_model():
    return pipeline("automatic-speech-recognition", model="openai/whisper-base")

# Audio processing function
def process_audio(audio_bytes):
    waveform, sample_rate = torchaudio.load(BytesIO(audio_bytes))
    if waveform.shape[0] > 1:  # Convert stereo to mono
        waveform = torch.mean(waveform, dim=0, keepdim=True)
    if sample_rate != 16000:  # Resample to 16kHz if needed
        resampler = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)
        waveform = resampler(waveform)
    return {"raw": waveform.numpy().squeeze(), "sampling_rate": 16000}

# Voice input component
def voice_input(key, prompt_text, default_text=""):
    col1, col2 = st.columns([4, 1])
    with col1:
        text_input = st.text_input(prompt_text, value=default_text, key=f"text_{key}")
    with col2:
        audio_bytes = audio_recorder(
            pause_threshold=0.8,
            text="🎤 Speak",
            recording_color="#e8b622",
            neutral_color="#6aa36f",
            key=f"recorder_{key}"
        )
    
    # Process audio if new recording is available
    if audio_bytes:
        current_hash = hashlib.md5(audio_bytes).hexdigest()
        if f"last_audio_hash_{key}" not in st.session_state or current_hash != st.session_state[f"last_audio_hash_{key}"]:
            st.session_state[f"last_audio_hash_{key}"] = current_hash
            try:
                audio_input = process_audio(audio_bytes)
                whisper = load_model()
                transcribed_text = whisper(audio_input)["text"]
                
                # Update the corresponding text input
                st.session_state[f"text_{key}"] = transcribed_text
                st.rerun()
                
            except Exception as e:
                st.error(f"Error in voice input: {str(e)}")
    
    return text_input

# Import transformers and cache the help agent for performance
@st.cache_resource
def get_help_agent():
    from transformers import pipeline
    # Using BlenderBot 400M Distill as the public conversational model (used elsewhere)
    return pipeline("conversational", model="facebook/blenderbot-400M-distill")

# [Rest of your existing functions remain exactly the same...]
# inject_custom_css()
# show_confetti()
# ask_llama()
# ask_help_agent()

def main():
    inject_custom_css()

    st.markdown('<div class="title">KASOTI</div>', unsafe_allow_html=True)
    st.markdown('<div class="subtitle">AI-Powered Guessing Game Challenge</div>', unsafe_allow_html=True)

    if 'game_state' not in st.session_state:
        st.session_state.game_state = "start"
        st.session_state.questions = []
        st.session_state.current_q = 0
        st.session_state.answers = []
        st.session_state.conversation_history = []
        st.session_state.category = None
        st.session_state.final_guess = None
        st.session_state.help_conversation = []

    # Start screen with voice input
    if st.session_state.game_state == "start":
        with st.container():
            st.markdown("""
            <div class="question-box">
                <h3 style="color: #6C63FF; margin-bottom: 1.5rem;">🎮 Welcome to KASOTI</h3>
                <p style="line-height: 1.6; color: #64748B;">
                    Think of something and I'll try to guess it in 20 questions or less!<br>
                    Choose from these categories:
                </p>
                <div style="display: grid; gap: 1rem; margin: 2rem 0;">
                    <div style="padding: 1.5rem; background: #f8f9fa; border-radius: 12px;">
                        <h4 style="margin: 0; color: #6C63FF;">🧑 Person</h4>
                        <p style="margin: 0.5rem 0 0; color: #64748B;">Celebrity, fictional character, historical figure</p>
                    </div>
                    <div style="padding: 1.5rem; background: #f8f9fa; border-radius: 12px;">
                        <h4 style="margin: 0; color: #6C63FF;">🌍 Place</h4>
                        <p style="margin: 0.5rem 0 0; color: #64748B;">City, country, landmark, geographical location</p>
                    </div>
                    <div style="padding: 1.5rem; background: #f8f9fa; border-radius: 12px;">
                        <h4 style="margin: 0; color: #6C63FF;">🎯 Object</h4>
                        <p style="margin: 0.5rem 0 0; color: #64748B;">Everyday item, tool, vehicle, or concept</p>
                    </div>
                </div>
            </div>
            """, unsafe_allow_html=True)

        with st.form("start_form"):
            # Replace text input with voice input component
            category_input = voice_input("category", "Enter category (person/place/object):").strip().lower()
            
            if st.form_submit_button("Start Game"):
                if not category_input:
                    st.error("Please enter a category!")
                elif category_input not in ["person", "place", "object"]:
                    st.error("Please enter either 'person', 'place', or 'object'!")
                else:
                    st.session_state.category = category_input
                    first_question = ask_llama([
                        {"role": "user", "content": "Ask your first strategic yes/no question."}
                    ], category_input)
                    st.session_state.questions = [first_question]
                    st.session_state.conversation_history = [
                        {"role": "assistant", "content": first_question}
                    ]
                    st.session_state.game_state = "gameplay"
                    st.experimental_rerun()

    # Gameplay screen with voice answer input
    elif st.session_state.game_state == "gameplay":
        with st.container():
            progress = (st.session_state.current_q + 1) / 20
            st.markdown(f"""
            <div class="question-count">QUESTION {st.session_state.current_q + 1} OF 20</div>
            <div class="progress-bar">
                <div class="progress-fill" style="width: {progress * 100}%"></div>
            </div>
            """, unsafe_allow_html=True)

            current_question = st.session_state.questions[st.session_state.current_q]
            
            st.markdown(f'''
            <div class="question-box">
                <div style="display: flex; align-items: center; gap: 1rem; margin-bottom: 1.5rem;">
                    <div style="background: #6C63FF; width: 40px; height: 40px; border-radius: 50%; 
                            display: flex; align-items: center; justify-content: center; color: white;">
                        <i class="fas fa-robot"></i>
                    </div>
                    <h3 style="margin: 0; color: #1E293B;">AI Question</h3>
                </div>
                <p style="font-size: 1.1rem; line-height: 1.6; color: #1E293B;">{current_question}</p>
            </div>
            ''', unsafe_allow_html=True)

        if "Final Guess:" in current_question:
            st.session_state.final_guess = current_question.split("Final Guess:")[1].strip()
            st.session_state.game_state = "confirm_guess"
            st.experimental_rerun()

        with st.form("answer_form"):
            # Replace text input with voice input component for answers
            answer_input = voice_input(f"answer_{st.session_state.current_q}", 
                                     "Your answer (yes/no/both):").strip().lower()
            
            if st.form_submit_button("Submit"):
                if answer_input not in ["yes", "no", "both"]:
                    st.error("Please answer with 'yes', 'no', or 'both'!")
                else:
                    st.session_state.answers.append(answer_input)
                    st.session_state.conversation_history.append(
                        {"role": "user", "content": answer_input}
                    )

                    next_response = ask_llama(
                        st.session_state.conversation_history,
                        st.session_state.category
                    )

                    if "Final Guess:" in next_response:
                        st.session_state.final_guess = next_response.split("Final Guess:")[1].strip()
                        st.session_state.game_state = "confirm_guess"
                    else:
                        st.session_state.questions.append(next_response)
                        st.session_state.conversation_history.append(
                            {"role": "assistant", "content": next_response}
                        )
                        st.session_state.current_q += 1

                        if st.session_state.current_q >= 20:
                            st.session_state.game_state = "result"

                    st.experimental_rerun()

        # Help assistant with voice input
        with st.expander("Need Help? Chat with AI Assistant"):
            # Replace help query input with voice input
            help_query = voice_input("help_query", "Enter your help query:")
            
            if st.button("Send", key="send_help"):
                if help_query:
                    help_response = ask_help_agent(help_query)
                    st.session_state.help_conversation.append({"query": help_query, "response": help_response})
                else:
                    st.error("Please enter a query!")
            if st.session_state.help_conversation:
                for msg in st.session_state.help_conversation:
                    st.markdown(f"**You:** {msg['query']}")
                    st.markdown(f"**Help Assistant:** {msg['response']}")

    # Guess confirmation with voice input
    elif st.session_state.game_state == "confirm_guess":
        st.markdown(f'''
        <div class="question-box">
            <div style="display: flex; align-items: center; gap: 1rem; margin-bottom: 1.5rem;">
                <div style="background: #6C63FF; width: 40px; height: 40px; border-radius: 50%; 
                        display: flex; align-items: center; justify-content: center; color: white;">
                    <i class="fas fa-lightbulb"></i>
                </div>
                <h3 style="margin: 0; color: #1E293B;">AI's Final Guess</h3>
            </div>
            <p style="font-size: 1.2rem; line-height: 1.6; color: #1E293B;">
                Is it <strong style="color: #6C63FF;">{st.session_state.final_guess}</strong>?
            </p>
        </div>
        ''', unsafe_allow_html=True)

        with st.form("confirm_form"):
            # Replace confirmation input with voice input
            confirm_input = voice_input("confirm_input", 
                                       "Type your answer (yes/no/both):").strip().lower()
            
            if st.form_submit_button("Submit"):
                if confirm_input not in ["yes", "no", "both"]:
                    st.error("Please answer with 'yes', 'no', or 'both'!")
                else:
                    if confirm_input == "yes":
                        st.session_state.game_state = "result"
                        st.experimental_rerun()
                        st.stop()
                    else:
                        st.session_state.conversation_history.append(
                            {"role": "user", "content": "no"}
                        )
                        st.session_state.game_state = "gameplay"
                        next_response = ask_llama(
                            st.session_state.conversation_history,
                            st.session_state.category
                        )
                        st.session_state.questions.append(next_response)
                        st.session_state.conversation_history.append(
                            {"role": "assistant", "content": next_response}
                        )
                        st.session_state.current_q += 1
                        st.experimental_rerun()

    # Result screen (unchanged)
    elif st.session_state.game_state == "result":
        if not st.session_state.final_guess:
            qa_history = "\n".join(
                [f"Q{i+1}: {q}\nA: {a}"
                 for i, (q, a) in enumerate(zip(st.session_state.questions, st.session_state.answers))]
            )

            final_guess = ask_llama(
                [{"role": "user", "content": qa_history}],
                st.session_state.category,
                is_final_guess=True
            )
            st.session_state.final_guess = final_guess.split("Final Guess:")[-1].strip()

        show_confetti()
        st.markdown(f'<div class="final-reveal">🎉 It\'s...</div>', unsafe_allow_html=True)
        time.sleep(1)
        st.markdown(f'<div class="final-reveal" style="font-size:3.5rem;color:#6C63FF;">{st.session_state.final_guess}</div>',
                    unsafe_allow_html=True)
        st.markdown(f"<p style='text-align:center; color:#64748B;'>Guessed in {len(st.session_state.questions)} questions</p>",
                    unsafe_allow_html=True)

        if st.button("Play Again", key="play_again"):
            st.session_state.clear()
            st.experimental_rerun()

if __name__ == "__main__":
    main()