import os
import gradio as gr
from neo4j import GraphDatabase
import pandas as pd
import plotly.express as px

class MovieRecommender:
    def __init__(self):
        uri = os.getenv("NEO4J_URI")
        user = os.getenv("NEO4J_USER")
        password = os.getenv("NEO4J_PASSWORD")
        self.driver = GraphDatabase.driver(uri, auth=(user, password))

    def get_genres(self):
        with self.driver.session() as session:
            query = "MATCH (g:Genere) RETURN DISTINCT g.genres ORDER BY g.genres"
            result = session.run(query)
            return [record['g.genres'] for record in result]

    def find_movies_by_genre(self, genre):
        with self.driver.session() as session:
            query = """ 
            MATCH (t:Title)-[:TITLE_TO_GENRE]->(g:Genre {genres: $genre})
            MATCH (t)-[:TITLE_TO_YEAR]->(y:Year)
            MATCH (t)-[:TITLE_TO_VOTES]->(v:Votes)
            RETURN 
                t.title as movie, 
                y.releaseYear as year, 
                v.avgVote as avgVote
            ORDER BY avgVote DESC
            LIMIT 50
            """
            result = session.run(query, genre=genre)
            df = pd.DataFrame([dict(record) for record in result])
            
            # Add rating category column
            df['vote_rating'] = df['avgVote']
            df['rating_category'] = pd.cut(
                df['vote_rating'], 
                bins=[0, 5, 7, float('inf')], 
                labels=['Low Rating', 'Medium Rating', 'High Rating']
            )
            return df

def create_interface():
    recommender = MovieRecommender()
    
    def recommend_movies(genre):
        if not genre:
            return pd.DataFrame(), None
        
        df = recommender.find_movies_by_genre(genre)
        
        if df.empty:
            return df, None
        
        fig = px.bar(df, x='movie', y='vote_rating', 
                     color='rating_category', 
                     title=f'Movies in {genre} Genre',
                     labels={'vote_rating': 'Average Vote', 'movie': 'Movie Title'},
                     hover_data=['year'])
        
        return df, fig

    # Create Gradio interface
    with gr.Blocks() as demo:
        gr.Markdown("# Movie Recommendation System")
        
        with gr.Row():
            genre_dropdown = gr.Dropdown(
                choices=recommender.get_genres(), 
                label="Select Genre"
            )
            recommend_btn = gr.Button("Get Recommendations")
        
        output_table = gr.DataFrame(label="Movie Recommendations")
        output_plot = gr.Plot(label="Movie Ratings Visualization")
        
        recommend_btn.click(
            fn=recommend_movies, 
            inputs=genre_dropdown, 
            outputs=[output_table, output_plot]
        )
    
    return demo

# Launch the app
demo = create_interface()
demo.launch(debug=True)