import openai
from pinecone import Pinecone, ServerlessSpec
import pandas as pd
import gradio as gr
from typing import List, Tuple

# Function to get embeddings from OpenAI's model
def get_embedding(text: str, openai_api_key: str, model: str = "text-embedding-ada-002") -> List[float]:
    openai.api_key = openai_api_key
    try:
        response = openai.Embedding.create(
            model=model,
            input=text
        )
        return response['data'][0]['embedding']
    except Exception as e:
        print(f"Error getting embedding: {e}")
        return []

# Function to process the uploaded CSV and store embeddings in Pinecone
def process_csv(file, openai_api_key: str, pinecone_api_key: str, pinecone_env: str) -> str:
    try:
        df = pd.read_csv(file.name)
        
        # Initialize Pinecone
        pc = Pinecone(api_key=pinecone_api_key)
        index_name = "product-recommendations"
        
        # Check if index exists
        if index_name not in pc.list_indexes().names():
            try:
                pc.create_index(
                    name=index_name,
                    dimension=1536,
                    spec=ServerlessSpec(cloud="aws", region=pinecone_env)
                )
            except Exception as e:
                print(f"Error creating Pinecone index: {e}")
                return "Failed to create Pinecone index."

        index = pc.Index(index_name)
        
        embeddings = []
        for i, row in df.iterrows():
            embedding = get_embedding(row['description'], openai_api_key)
            if embedding:
                embeddings.append((str(row['product_id']), embedding, {'product_name': row['product_name'], 'image_url': row['image_url']}))
        
        if embeddings:
            try:
                index.upsert(embeddings)
            except Exception as e:
                print(f"Error upserting embeddings to Pinecone: {e}")
                return "Failed to upsert embeddings."
        
        return "Product catalog processed and embeddings stored in Pinecone."
    except Exception as e:
        print(f"Error processing CSV file: {e}")
        return "Failed to process CSV file."

# Recommendation logic
def recommend_products(query: str, openai_api_key: str, pinecone_api_key: str, pinecone_env: str, top_k: int = 5) -> List[Tuple[str, str]]:
    query_embedding = get_embedding(query, openai_api_key)
    
    if not query_embedding:
        return []

    try:
        # Initialize Pinecone
        pc = Pinecone(api_key=pinecone_api_key)
        index = pc.Index("product-recommendations")
        
        results = index.query(vector=query_embedding, top_k=top_k, include_metadata=True)
        recommended_products = [(match['metadata']['image_url'], f"{match['metadata']['product_name']} (Score: {match['score']})") for match in results['matches']]
        return recommended_products
    except Exception as e:
        print(f"Error querying Pinecone: {e}")
        return []

# Function to generate contextual message
def generate_contextual_message(query: str, recommendations: List[Tuple[str, str]], openai_api_key: str) -> str:
    openai.api_key = openai_api_key
    product_names = [rec[1] for rec in recommendations]
    prompt = f"User query: {query}\nRecommended products: {', '.join(product_names)}\nGenerate a personalized message for the user based on these recommendations."
    
    try:
        response = openai.ChatCompletion.create(
            model="gpt-4",  # or use "gpt-3.5-turbo" if preferred
            messages=[{"role": "system", "content": "You are a helpful assistant."},
                      {"role": "user", "content": prompt}]
        )
        return response['choices'][0]['message']['content']
    except Exception as e:
        print(f"Error generating contextual message: {e}")
        return "Failed to generate contextual message."

# Gradio interface
def handle_file_upload(file, openai_api_key, pinecone_api_key, pinecone_env):
    return process_csv(file, openai_api_key, pinecone_api_key, pinecone_env)

def display_recommendations(user_input, openai_api_key, pinecone_api_key, pinecone_env):
    recommendations = recommend_products(user_input, openai_api_key, pinecone_api_key, pinecone_env)
    contextual_message = generate_contextual_message(user_input, recommendations, openai_api_key)
    return recommendations, contextual_message

# Function to update outputs
def update_outputs(query_input, openai_api_key, pinecone_api_key, pinecone_env, chat_history):
    recommendations, contextual_message = display_recommendations(query_input, openai_api_key, pinecone_api_key, pinecone_env)
    
    # Update chat history
    new_chat_history = chat_history + [("user", query_input),("assistant", contextual_message)]
    
    return recommendations, new_chat_history

# Create Gradio Interface
def build_interface():
    with gr.Blocks() as interface:
        gr.Markdown("## Product Recommender System")
        
        with gr.Tab("API Keys"):
            openai_api_key_input = gr.Textbox(label="OpenAI API Key", type="password")
            pinecone_api_key_input = gr.Textbox(label="Pinecone API Key", type="password")
            pinecone_env_input = gr.Textbox(label="Pinecone Environment", placeholder="e.g., us-west1-gcp")
        
        with gr.Tab("Upload Catalog"):
            upload_button = gr.File(label="Upload CSV", type="filepath")
            output = gr.Textbox()
            upload_button.upload(handle_file_upload, inputs=[upload_button, openai_api_key_input, pinecone_api_key_input, pinecone_env_input], outputs=output)
        
        with gr.Tab("Get Recommendations"):
            with gr.Row():
                with gr.Column(scale=1):
                    chatbot = gr.Chatbot(label="Chat")
                    query_input = gr.Textbox(label="Enter your product preference...", show_label=False, placeholder="Type your query here...")
                    recommend_button = gr.Button("Get Recommendations")
                    # Define state for chat history
                    chat_history = gr.State([])
                    
                    # Define outputs first
                with gr.Column(scale=1):
                    recommendations_output = gr.Gallery(label="Recommendations")
                    
                recommend_button.click(
                    update_outputs,
                    inputs=[query_input, openai_api_key_input, pinecone_api_key_input, pinecone_env_input, chat_history],
                    outputs=[recommendations_output, chatbot]
                )
                
    return interface

# Run the interface
if __name__ == "__main__":
    interface = build_interface()
    interface.launch()