import streamlit as st
import boto3
import json
import chromadb
import pandas as pd
import time
import re
from datetime import datetime

# Sample Bollywood movies data (simplified for demo)
SAMPLE_MOVIES = [
    {"title": "Sholay", "year": 1975, "genre": "Action", "director": "Ramesh Sippy", 
     "plot": "Two criminals are hired by a retired police officer to capture a bandit terrorizing a village."},
    {"title": "Dilwale Dulhania Le Jayenge", "year": 1995, "genre": "Romance", "director": "Aditya Chopra",
     "plot": "A young man and woman fall in love during a trip to Europe, but face family opposition."},
    {"title": "Lagaan", "year": 2001, "genre": "Drama", "director": "Ashutosh Gowariker",
     "plot": "Villagers accept a challenge from British officers to play cricket to avoid paying tax."},
    {"title": "3 Idiots", "year": 2009, "genre": "Comedy", "director": "Rajkumar Hirani",
     "plot": "Two friends search for their missing college friend and recall their engineering days."},
    {"title": "Dangal", "year": 2016, "genre": "Sports", "director": "Nitesh Tiwari",
     "plot": "A former wrestler trains his daughters to become world-class wrestlers."},
    {"title": "Anand", "year": 1971, "genre": "Drama", "director": "Hrishikesh Mukherjee",
     "plot": "A terminally ill man spreads joy and teaches the meaning of life to a doctor."},
    {"title": "Golmaal", "year": 1979, "genre": "Comedy", "director": "Hrishikesh Mukherjee",
     "plot": "A man creates chaos by lying about his identity to get a job."},
    {"title": "Chupke Chupke", "year": 1975, "genre": "Comedy", "director": "Hrishikesh Mukherjee",
     "plot": "A newlywed plays pranks on his wife's family by pretending to be someone else."},
    {"title": "Don", "year": 1978, "genre": "Action", "director": "Chandra Barot",
     "plot": "A police officer impersonates a crime boss to infiltrate his gang."},
    {"title": "Andaz Apna Apna", "year": 1994, "genre": "Comedy", "director": "Rajkumar Santoshi",
     "plot": "Two friends compete to marry a wealthy heiress but get caught up in a kidnapping plot."},
    {"title": "Mughal-E-Azam", "year": 1960, "genre": "Romance", "director": "K. Asif",
     "plot": "A Mughal prince falls in love with a court dancer, defying his father the emperor."},
    {"title": "Deewaar", "year": 1975, "genre": "Action", "director": "Yash Chopra",
     "plot": "Two brothers choose different paths in life - one becomes a police officer, the other a criminal."},
    {"title": "Queen", "year": 2013, "genre": "Comedy", "director": "Vikas Bahl",
     "plot": "A woman goes on her honeymoon alone after her wedding is called off."},
    {"title": "Zindagi Na Milegi Dobara", "year": 2011, "genre": "Adventure", "director": "Zoya Akhtar",
     "plot": "Three friends go on a bachelor trip to Spain and face their fears."},
    {"title": "Taare Zameen Par", "year": 2007, "genre": "Drama", "director": "Aamir Khan",
     "plot": "An art teacher helps a dyslexic child overcome his learning difficulties."},
    {"title": "Rang De Basanti", "year": 2006, "genre": "Drama", "director": "Rakeysh Omprakash Mehra",
     "plot": "College students making a documentary about freedom fighters become revolutionaries themselves."},
    {"title": "Gol Maal", "year": 1979, "genre": "Comedy", "director": "Hrishikesh Mukherjee",
     "plot": "A young man lies about having a mustache to keep his job with a strict boss."},
    {"title": "Namak Haraam", "year": 1973, "genre": "Drama", "director": "Hrishikesh Mukherjee",
     "plot": "A friendship is tested when one friend betrays the other for money and power."},
    {"title": "Kuch Kuch Hota Hai", "year": 1998, "genre": "Romance", "director": "Karan Johar",
     "plot": "A man's daughter tries to reunite him with his college sweetheart."},
    {"title": "My Name is Khan", "year": 2010, "genre": "Drama", "director": "Karan Johar",
     "plot": "A man with Asperger's syndrome embarks on a journey to meet the President of the United States."}
]

# Simple function to connect to AWS Bedrock
def connect_to_bedrock():
    try:
        client = boto3.client('bedrock-runtime', region_name='us-east-1')
        return client
    except:
        st.error("⚠️ AWS Bedrock not configured. Using mock responses for demo.")
        return None

# Get embeddings from Bedrock
def get_embeddings(bedrock_client, text):
    if not bedrock_client:
        # Return dummy embedding for demo
        import random
        return [random.random() for _ in range(1536)]
    
    try:
        body = json.dumps({"inputText": text})
        response = bedrock_client.invoke_model(
            modelId="amazon.titan-embed-text-v1",
            body=body
        )
        result = json.loads(response['body'].read())
        return result['embedding']
    except:
        # Return dummy embedding if API fails
        import random
        return [random.random() for _ in range(1536)]

# Create movie documents and store in ChromaDB
def setup_movie_database(bedrock_client):
    st.write("🎬 Setting up Bollywood movies database...")
    
    # Create ChromaDB client
    chroma_client = chromadb.Client()
    
    # Create or recreate collection
    try:
        chroma_client.delete_collection("bollywood_movies")
    except:
        pass
    
    collection = chroma_client.create_collection("bollywood_movies")
    
    # Prepare data for ChromaDB
    ids = []
    documents = []
    metadatas = []
    embeddings = []
    
    progress_bar = st.progress(0)
    
    for i, movie in enumerate(SAMPLE_MOVIES):
        # Create document text
        doc_text = f"Title: {movie['title']}\nYear: {movie['year']}\nGenre: {movie['genre']}\nDirector: {movie['director']}\nPlot: {movie['plot']}"
        
        # Get embedding
        embedding = get_embeddings(bedrock_client, doc_text)
        
        # Prepare data
        ids.append(str(i))
        documents.append(doc_text)
        metadatas.append({
            'title': movie['title'],
            'year': movie['year'],
            'genre': movie['genre'].lower(),
            'director': movie['director'].lower(),
            'decade': f"{(movie['year'] // 10) * 10}s"
        })
        embeddings.append(embedding)
        
        progress_bar.progress((i + 1) / len(SAMPLE_MOVIES))
    
    # Add to ChromaDB
    collection.add(
        ids=ids,
        documents=documents,
        metadatas=metadatas,
        embeddings=embeddings
    )
    
    st.success(f"✅ Added {len(SAMPLE_MOVIES)} movies to database!")
    return collection

# Simple query filter detection
def detect_filters(query):
    query_lower = query.lower()
    filters = {}
    
    # Genre detection
    genres = ['action', 'comedy', 'drama', 'romance', 'sports', 'adventure']
    for genre in genres:
        if genre in query_lower:
            filters['genre'] = genre
            break
    
    # Decade detection
    decades = ['1960s', '1970s', '1980s', '1990s', '2000s', '2010s']
    for decade in decades:
        if decade in query_lower:
            filters['decade'] = decade
            break
    
    # Year detection
    years = re.findall(r'\b(19\d{2}|20\d{2})\b', query)
    if years:
        year = int(years[0])
        filters['decade'] = f"{(year // 10) * 10}s"
    
    # Director detection (simple)
    directors = ['hrishikesh mukherjee', 'rajkumar hirani', 'aamir khan', 'yash chopra']
    for director in directors:
        if director in query_lower:
            filters['director'] = director
            break
    
    return filters

# Retrieve without metadata filter
def retrieve_without_filter(collection, bedrock_client, query, top_k=5):
    start_time = time.time()
    
    # Get query embedding
    query_embedding = get_embeddings(bedrock_client, query)
    
    # Search without filters
    results = collection.query(
        query_embeddings=[query_embedding],
        n_results=top_k
    )
    
    end_time = time.time()
    
    # Format results
    movies = []
    for i in range(len(results['documents'][0])):
        movies.append({
            'document': results['documents'][0][i],
            'metadata': results['metadatas'][0][i],
            'distance': results['distances'][0][i]
        })
    
    return movies, end_time - start_time

# Retrieve with metadata filter
def retrieve_with_filter(collection, bedrock_client, query, filters, top_k=5):
    start_time = time.time()
    
    # Get query embedding
    query_embedding = get_embeddings(bedrock_client, query)
    
    # Create where clause for filtering
    where_clause = {}
    for key, value in filters.items():
        where_clause[key] = value
    
    # Search with filters
    try:
        results = collection.query(
            query_embeddings=[query_embedding],
            n_results=top_k,
            where=where_clause
        )
    except:
        # If filtering fails, fall back to no filter
        results = collection.query(
            query_embeddings=[query_embedding],
            n_results=top_k
        )
    
    end_time = time.time()
    
    # Format results
    movies = []
    for i in range(len(results['documents'][0])):
        movies.append({
            'document': results['documents'][0][i],
            'metadata': results['metadatas'][0][i],
            'distance': results['distances'][0][i]
        })
    
    return movies, end_time - start_time

# Generate answer using Bedrock
def generate_answer(bedrock_client, query, movies):
    if not bedrock_client:
        return "🎬 Based on the retrieved movies, here are some recommendations that match your query!"
    
    # Create context from movies
    context = "\n\n".join([movie['document'] for movie in movies])
    
    prompt = f"""
    Based on the following Bollywood movies information, please answer the user's question.
    
    Question: {query}
    
    Movies Information:
    {context}
    
    Please provide a helpful and informative answer about the movies.
    """
    
    try:
        body = json.dumps({
            "anthropic_version": "bedrock-2023-05-31",
            "max_tokens": 400,
            "messages": [{"role": "user", "content": prompt}]
        })
        
        response = bedrock_client.invoke_model(
            modelId="anthropic.claude-3-haiku-20240307-v1:0",
            body=body
        )
        
        result = json.loads(response['body'].read())
        return result['content'][0]['text']
    except:
        return "🎬 Based on the retrieved movies, here are some great recommendations that match your query!"

# Main app
def main():
    st.title("🎬 Bollywood Movies RAG with Metadata Filtering")
    st.write("Ask questions about Bollywood movies and see how metadata filtering speeds up retrieval!")
    
    # Initialize session state
    if 'collection' not in st.session_state:
        st.session_state.collection = None
    if 'setup_done' not in st.session_state:
        st.session_state.setup_done = False
    
    # Setup section
    if not st.session_state.setup_done:
        st.subheader("🛠️ Setup Movie Database")
        
        if st.button("🚀 Load Bollywood Movies Data"):
            try:
                bedrock_client = connect_to_bedrock()
                collection = setup_movie_database(bedrock_client)
                st.session_state.collection = collection
                st.session_state.bedrock_client = bedrock_client
                st.session_state.setup_done = True
                st.balloons()
            except Exception as e:
                st.error(f"❌ Setup failed: {str(e)}")
    
    else:
        st.success("✅ Movie database is ready!")
        
        # Sample queries
        st.subheader("🔍 Try These Sample Queries")
        sample_queries = [
            "What are some good action movies?",
            "Tell me a few comedy movies from the 1970s",
            "What is the movie Sholay about?",
            "Tell me a few movies directed by Hrishikesh Mukherjee",
            "What are some romantic movies from the 1990s?"
        ]
        
        query_option = st.radio("Choose a query:", ["Custom Query"] + sample_queries)
        
        if query_option == "Custom Query":
            query = st.text_input("Enter your question about Bollywood movies:")
        else:
            query = query_option
            st.write(f"Selected: **{query}**")
        
        if query:
            if st.button("🔍 Search Movies"):
                try:
                    bedrock_client = st.session_state.bedrock_client
                    collection = st.session_state.collection
                    
                    # Detect filters
                    filters = detect_filters(query)
                    
                    st.write("---")
                    
                    # Method 1: Without metadata filter
                    st.subheader("📊 Method 1: Without Metadata Filter")
                    movies_no_filter, time_no_filter = retrieve_without_filter(collection, bedrock_client, query)
                    
                    st.write(f"⏱️ **Time taken: {time_no_filter:.4f} seconds**")
                    st.write("**Retrieved Movies:**")
                    for i, movie in enumerate(movies_no_filter, 1):
                        with st.expander(f"{i}. {movie['metadata']['title']} ({movie['metadata']['year']})"):
                            st.write(f"**Genre:** {movie['metadata']['genre'].title()}")
                            st.write(f"**Director:** {movie['metadata']['director'].title()}")
                            st.write(f"**Distance:** {movie['distance']:.4f}")
                    
                    # Method 2: With metadata filter
                    st.subheader("🎯 Method 2: With Metadata Filter")
                    
                    if filters:
                        st.write(f"**Detected Filters:** {filters}")
                        movies_with_filter, time_with_filter = retrieve_with_filter(collection, bedrock_client, query, filters)
                        
                        st.write(f"⏱️ **Time taken: {time_with_filter:.4f} seconds**")
                        st.write("**Filtered Retrieved Movies:**")
                        for i, movie in enumerate(movies_with_filter, 1):
                            with st.expander(f"{i}. {movie['metadata']['title']} ({movie['metadata']['year']})"):
                                st.write(f"**Genre:** {movie['metadata']['genre'].title()}")
                                st.write(f"**Director:** {movie['metadata']['director'].title()}")
                                st.write(f"**Distance:** {movie['distance']:.4f}")
                        
                        # Performance comparison
                        st.subheader("⚡ Performance Comparison")
                        col1, col2, col3 = st.columns(3)
                        with col1:
                            st.metric("Without Filter", f"{time_no_filter:.4f}s")
                        with col2:
                            st.metric("With Filter", f"{time_with_filter:.4f}s")
                        with col3:
                            speedup = ((time_no_filter - time_with_filter) / time_no_filter) * 100 if time_no_filter > 0 else 0
                            st.metric("Speedup", f"{speedup:.1f}%")
                        
                        # Generate final answer
                        st.subheader("🤖 AI Generated Answer")
                        answer = generate_answer(bedrock_client, query, movies_with_filter)
                        st.success(answer)
                    
                    else:
                        st.write("**No specific filters detected** - using general retrieval")
                        st.write(f"⏱️ **Time taken: {time_no_filter:.4f} seconds**")
                        
                        # Generate answer with no filter results
                        st.subheader("🤖 AI Generated Answer")
                        answer = generate_answer(bedrock_client, query, movies_no_filter)
                        st.success(answer)
                    
                except Exception as e:
                    st.error(f"❌ Search failed: {str(e)}")
        
        # Show movie database
        if st.checkbox("📋 Show All Movies in Database"):
            st.subheader("Movie Database")
            df = pd.DataFrame(SAMPLE_MOVIES)
            st.dataframe(df)
        
        # Reset button
        if st.button("🔄 Reset Database"):
            st.session_state.collection = None
            st.session_state.setup_done = False
            st.rerun()

# Installation and deployment guide
def show_guides():
    col1, col2 = st.columns(2)
    
    with col1:
        with st.expander("📖 Installation Guide"):
            st.markdown("""
            **Step 1: Install Libraries**
            ```bash
            pip install streamlit boto3 chromadb pandas
            ```
            
            **Step 2: Setup AWS**
            ```bash
            aws configure
            ```
            
            **Step 3: Run Locally**
            ```bash
            streamlit run bollywood_rag.py
            ```
            """)
    
    with col2:
        with st.expander("🚀 Deploy to Hugging Face"):
            st.markdown("""
            **Step 1: Create files**
            - `app.py` (this code)
            - `requirements.txt`
            - `README.md`
            
            **Step 2: requirements.txt**
            ```
            streamlit
            boto3
            chromadb
            pandas
            ```
            
            **Step 3: Deploy**
            1. Push to GitHub
            2. Connect to Hugging Face Spaces
            3. Select Streamlit SDK
            4. Add AWS secrets in settings
            """)

# Run the app
if __name__ == "__main__":
    show_guides()
    main()