import streamlit as st
from transformers import GPT2LMHeadModel, GPT2Tokenizer
from streamlit_webrtc import webrtc_streamer, VideoProcessorBase, WebRtcMode

# Load the pretrained DialoGPT model
tokenizer = GPT2Tokenizer.from_pretrained("microsoft/DialoGPT-medium")
model = GPT2LMHeadModel.from_pretrained("microsoft/DialoGPT-medium")

# Streamlit UI Setup
st.title("AI Multimodal Chat & File Processing App")

# Session state for chat history
if "history" not in st.session_state:
    st.session_state.history = []

# Function to process the chat
def chat_with_model(user_input):
    new_user_input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors="pt")
    st.session_state.history.append(new_user_input_ids)

    bot_input_ids = new_user_input_ids
    for history in st.session_state.history:
        bot_input_ids = history if len(history) < 2048 else history[-1024:]

    chat_history_ids = model.generate(bot_input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id)
    bot_output = tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
    return bot_output

# Chat Input Box
st.subheader("Chat with AI")
user_input = st.text_input("You: ", "")
if user_input:
    response = chat_with_model(user_input)
    st.session_state.history.append(tokenizer.encode(user_input + tokenizer.eos_token, return_tensors="pt"))
    st.write(f"Bot: {response}")

# Show chat history (safer decoding)
if st.session_state.history:
    st.subheader("📝 Chat History")
    for i in range(len(st.session_state.history) - 1, -1, -1):
        try:
            user_msg = tokenizer.decode(st.session_state.history[i][0].tolist(), skip_special_tokens=True)
            st.write(f"You: {user_msg}")
        except Exception as e:
            st.warning(f"Could not decode history message: {e}")

# --- File Upload and Processing ---
st.subheader("📁 Upload a File for AI to Read")

uploaded_file = st.file_uploader("Choose a text file", type=["txt", "csv", "md", "log"])
if uploaded_file:
    content = uploaded_file.read().decode("utf-8")
    st.text_area("File Content", content, height=200)

    # Allow interaction with file content
    file_question = st.text_input("Ask something about the file:")
    if file_question:
        combined_input = file_question + "\n" + content[:1000]  # Prevent token overload
        response = chat_with_model(combined_input)
        st.write(f"Bot: {response}")

# --- Video/Audio Stream ---
st.subheader("🎥 Video/Audio Stream")

class VideoProcessor(VideoProcessorBase):
    def recv(self, frame):
        return frame

webrtc_streamer(key="example", mode=WebRtcMode.SENDRECV, video_processor_factory=VideoProcessor)