import gradio as gr
from datasets import load_dataset
from transformers import pipeline

# Load the WikiArt dataset in streaming mode
dataset = load_dataset("huggan/wikiart", streaming=True)

# Function to safely get a field value or return a default message
def get_field(record, field, default="Unknown"):
    return record[field] if field in record and record[field] is not None else default

# Function to display artwork details
def display_artwork(index):
    try:
        for i, record in enumerate(dataset["train"]):  # Stream through the dataset
            if i == index:
                return (
                    get_field(record, "image"), 
                    f"Title: {get_field(record, 'title')}\n"
                    f"Artist: {get_field(record, 'artist')}\n"
                    f"Style: {get_field(record, 'style')}\n"
                    f"Genre: {get_field(record, 'genre')}"
                )
        return None, "Error: Index out of range or invalid."
    except Exception as e:
        return None, f"Error: {str(e)}"

# Function to filter artworks based on metadata
def filter_artworks(artist=None, genre=None, style=None):
    results = []
    try:
        for record in dataset["train"]:
            if (artist is None or get_field(record, "artist") == artist) and \
               (genre is None or get_field(record, "genre") == genre) and \
               (style is None or get_field(record, "style") == style):
                results.append(record)
    except Exception as e:
        return []
    return results

# Function to display filtered artworks
def display_filtered_artworks(artist, genre, style):
    filtered_results = filter_artworks(artist, genre, style)
    if len(filtered_results) == 0:
        return None, "No artworks found with the specified filters."
    return [(get_field(r, "image"), 
             f"Title: {get_field(r, 'title')}\n"
             f"Artist: {get_field(r, 'artist')}\n"
             f"Style: {get_field(r, 'style')}\n"
             f"Genre: {get_field(r, 'genre')}")
            for r in filtered_results]

# Chatbot functionality for museum guide persona using a publicly available Hugging Face model
chatbot = pipeline("text-generation", model="gpt2")  # Replace with a valid Hugging Face model

def museum_guide_query(prompt):
    try:
        response = chatbot(prompt, max_length=100, num_return_sequences=1)
        return response[0]["generated_text"]
    except Exception as e:
        return f"Error: {str(e)}"

# Gradio interfaces
artwork_interface = gr.Interface(
    fn=display_artwork,
    inputs=gr.Number(label="Artwork Index"),
    outputs=[gr.Image(label="Artwork"), gr.Text(label="Details")],
    title="Exhibit AI - Virtual Art Gallery"
)

filter_interface = gr.Interface(
    fn=display_filtered_artworks,
    inputs=[gr.Text(label="Artist"), gr.Text(label="Genre"), gr.Text(label="Style")],
    outputs=gr.Gallery(label="Filtered Artworks"),  # Removed the 'caption' argument
    title="Filter Artworks"
)

chatbot_interface = gr.Interface(
    fn=museum_guide_query,
    inputs=gr.Textbox(label="Ask the Museum Guide"),
    outputs=gr.Text(label="Guide Response"),
    title="Museum Guide Chatbot"
)

# Launch Gradio Blocks to combine all interfaces
def launch_app():
    with gr.Blocks() as demo:
        gr.Markdown("# Exhibit AI - Virtual Art Gallery")
        gr.TabbedInterface(
            [artwork_interface, filter_interface, chatbot_interface],
            ["View Artwork", "Filter Artworks", "Ask the Museum Guide"]
        )
    demo.launch()

if __name__ == "__main__":
    launch_app()