import gradio as gr
import pandas as pd
import sqlite3
import openai
import pinecone
import chromadb
from PyPDF2 import PdfReader
from transformers import pipeline
import os
from google.colab import auth
from googleapiclient.discovery import build

# Initialize APIs
openai.api_key = "YOUR_OPENAI_API_KEY"
pinecone.init(api_key="YOUR_PINECONE_API_KEY", environment="us-west1-gcp")
db = chromadb.Client() # Initialize ChromaDB

# Set up Gradio interface
def initialize_models(model_choice):
    if model_choice == 'OpenAI':
        return {
            'embedding': openai.Embedding.create,
            'chat': openai.ChatCompletion.create
        }
    elif model_choice == 'HuggingFace':
        embedding_model = pipeline('feature-extraction', model='bert-base-uncased')
        chat_model = pipeline('conversational', model='facebook/blenderbot-400M-distill')
        return {
            'embedding': embedding_model,
            'chat': chat_model
        }

def fetch_pdf_from_drive(file_id):
    auth.authenticate_user()
    drive_service = build('drive', 'v3')
    request = drive_service.files().get_media(fileId=file_id)
    file = io.BytesIO(request.execute())
    pdf_reader = PdfReader(file)
    text = ""
    for page in pdf_reader.pages:
        text += page.extract_text()
    return text

def query_db(query, db_type):
    conn = sqlite3.connect('data.db')
    cursor = conn.cursor()
    if db_type == 'Pinecone':
        index = pinecone.Index('your-index-name')
        results = index.query(query)
        return results
    elif db_type == 'ChromaDB':
        # Example of ChromaDB query - adapt as needed
        results = db.query(query)
        return results
    conn.close()

def generate_response(model_choice, query, chat_history, db_type):
    models = initialize_models(model_choice)
    if model_choice == 'OpenAI':
        response = models['chat'](model='gpt-3.5-turbo', messages=chat_history + [{'role': 'user', 'content': query}])
        return response['choices'][0]['message']['content'], chat_history + [{'role': 'user', 'content': query}, {'role': 'assistant', 'content': response['choices'][0]['message']['content']}]
    elif model_choice == 'HuggingFace':
        response = models['chat'](query)
        return response['generated_text'], chat_history + [{'role': 'user', 'content': query}, {'role': 'assistant', 'content': response['generated_text']}]

def process_input(model_choice, query, db_type, file_id):
    if file_id:
        pdf_text = fetch_pdf_from_drive(file_id)
        query = f"{query} {pdf_text}"
    response, updated_history = generate_response(model_choice, query, chat_history, db_type)
    return response, updated_history

def gradio_interface():
    with gr.Blocks() as demo:
        with gr.Row():
            model_choice = gr.Dropdown(['OpenAI', 'HuggingFace'], label='Model Choice', value='OpenAI')
            db_type = gr.Dropdown(['ChromaDB', 'Pinecone'], label='Database Type', value='ChromaDB')
            file_id = gr.Textbox(label='Google Drive File ID', placeholder='Enter Google Drive file ID (for PDFs)')
        
        with gr.Row():
            chat_history = gr.Chatbot()
            query = gr.Textbox(label='Query')
            submit_button = gr.Button('Submit')
        
        submit_button.click(fn=process_input, inputs=[model_choice, query, db_type, file_id], outputs=[chat_history])
    
    return demo

if __name__ == "__main__":
    gradio_interface().launch()