import streamlit as st
from transformers import pipeline
import time

# Page configuration
st.set_page_config(page_title="AI Q&A Assistant", layout="wide")

# Custom styling
st.markdown("""
    <style>
    .header {
        background-color: #1565C0;
        padding: 60px;
        text-align: center;
        border-radius: 15px;
        color: white;
        font-family: 'Arial', sans-serif;
    }
    .header h1 {
        font-size: 50px;
        font-weight: bold;
    }
    .qa-box {
        background-color: #f5f5f5;
        padding: 20px;
        border-radius: 10px;
        margin: 10px 0;
    }
    </style>
    <div class="header">
        <h1>🤔 AI Q&A Assistant</h1>
        <p>Get answers to your questions from any text</p>
    </div>
    """, unsafe_allow_html=True)

# Load model
@st.cache_resource
def load_qa_model():
    return pipeline("question-answering", model="distilbert-base-cased-distilled-squad")

qa_model = load_qa_model()

# Main content
col1, col2 = st.columns([6, 4])

with col1:
    st.subheader("Context")
    context = st.text_area(
        "Enter the text passage:",
        height=300,
        help="This is the text that the AI will use to answer questions"
    )

with col2:
    st.subheader("Question")
    question = st.text_input(
        "Ask a question about the text:",
        help="Make sure your question can be answered using the provided text"
    )

# History tracking
if 'qa_history' not in st.session_state:
    st.session_state.qa_history = []

if qa_model and st.button("Get Answer"):
    if context and question:
        with st.spinner("Finding answer..."):
            start_time = time.time()
            
            result = qa_model(
                question=question,
                context=context
            )
            
            end_time = time.time()
            
            # Add to history
            st.session_state.qa_history.append({
                'question': question,
                'answer': result['answer'],
                'confidence': result['score']
            })
            
            # Display result
            st.markdown(f"""
                <div class="qa-box">
                    <h3>Answer:</h3>
                    <p>{result['answer']}</p>
                    <small>Confidence: {result['score']:.2%}</small><br>
                    <small>Response time: {(end_time - start_time):.2f} seconds</small>
                </div>
                """, unsafe_allow_html=True)
    else:
        st.warning("Please provide both context and a question")

# Display history
if st.session_state.qa_history:
    st.subheader("Question History")
    for i, qa in enumerate(reversed(st.session_state.qa_history[-5:])):
        st.markdown(f"""
            <div class="qa-box">
                <strong>Q: {qa['question']}</strong><br>
                A: {qa['answer']}<br>
                <small>Confidence: {qa['confidence']:.2%}</small>
            </div>
            """, unsafe_allow_html=True)

# Clear history button
if st.session_state.qa_history:
    if st.button("Clear History"):
        st.session_state.qa_history = []
        st.experimental_rerun()