import os
import json
import pandas as pd
import gradio as gr
import openai
import time
from cpr import brd_g
from typing import List, Dict
from langchain.document_loaders import AirtableLoader

username = os.getenv('USERNAME')
password = os.getenv('PASSWORD')

# Set up API keys
AIRTABLE_API_KEY = os.getenv("AIRTABLE_API_KEY")
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
BASE_ID = os.getenv("base_id")
TABLE_ID = os.getenv("table_id")
VIEW = os.getenv("view")

# Set up OpenAI client
openai.api_key = OPENAI_API_KEY
client = openai.Client(api_key=OPENAI_API_KEY)

# Set up assistant
ASSISTANT_ID = os.getenv('assistant_id')
assistant = client.beta.assistants.retrieve(ASSISTANT_ID)
thread = client.beta.threads.create()

# Function to load data from Airtable and return as a DataFrame
def load_airtable_data() -> pd.DataFrame:
    loader = AirtableLoader(AIRTABLE_API_KEY, TABLE_ID, BASE_ID, view=VIEW)
    documents = loader.load()
    data = []
    for doc in documents:
        try:
            record = json.loads(doc.page_content)
            data.append(record)
        except json.JSONDecodeError:
            print(f"Warning: Could not parse JSON for document: {doc.page_content[:100]}...")
            data.append({"raw_content": doc.page_content})
    return pd.DataFrame(data)

# Load Airtable data into DataFrame
try:
    airtable_data_df = load_airtable_data()
    print(f"Successfully loaded data from Airtable.")
except Exception as e:
    print(f"Error loading Airtable data: {str(e)}")
    airtable_data_df = pd.DataFrame()

# Function to chat with the assistant
def chat_with_assistant(message: str, dataframe: pd.DataFrame) -> str:
    dataframe_csv = dataframe.to_csv(index=False)
    
    full_message = f"""
    Here is the student information:
    {dataframe_csv}
    
    Here is my question: {message}
    
    Only output your answer without code.
    """
    
    client.beta.threads.messages.create(
        thread_id=thread.id,
        role="user",
        content=full_message
    )
    
    run = client.beta.threads.runs.create(
        thread_id=thread.id,
        assistant_id=ASSISTANT_ID
    )
    
    while True:
        run_status = client.beta.threads.runs.retrieve(thread_id=thread.id, run_id=run.id)
        if run_status.status == 'completed':
            messages = client.beta.threads.messages.list(thread_id=thread.id)
            assistant_response = messages.data[0].content[0].text.value
            break
        time.sleep(1)
    
    return assistant_response

# Define the Gradio interface
def gradio_interface(question: str) -> str:
    return chat_with_assistant(question, airtable_data_df)

with gr.Blocks() as demo:
    iface = gr.Interface(
        fn=gradio_interface,
        inputs="text",
        outputs=gr.Markdown(),
        title="📅 Summative Assessment Tracker",
        description="I am here to help you schedule summative assessments for your students",
    )
    gr.Markdown(brd_g)

# Launch the Gradio app
demo.launch(auth=(username,password))