import streamlit as st
from transformers import OpenAIGPT
import datetime
import json
from langchain_community.llms import OpenAI

# Initialize OpenAI client
client = OpenAIGPT(api_key=None)

# Define function for interacting with GPT-3.5
def chat_gpt(prompt):
    response = client.chat_completions(
        model_name_or_path="openai/gpt-3.5-turbo",
        messages=[{"role": "user", "content": prompt}]
    )
    return response.choices[0].message.strip()

# Step 2: Data Preprocessing
def parse_user_query(query):
    # Assuming default dates are in ISO 8601 format (YYYY-MM-DD)
    default_start_date = datetime.date.today() - datetime.timedelta(days=365)
    default_end_date = datetime.date.today()
    
    # Implement parsing logic here
    return {
        "entity": "Company",
        "parameter": "Metric",
        "start_date": default_start_date.isoformat(),
        "end_date": default_end_date.isoformat()
    }

# Step 3: Interacting with the GPT-3.5 API
def query_gpt(prompt):
    return chat_gpt(prompt)

# Step 4: Generating JSON Output
def generate_json_output(entity, parameter, start_date, end_date):
    # Parse start_date and end_date strings to datetime objects
    start_date = datetime.datetime.strptime(start_date, '%Y-%m-%d')
    end_date = datetime.datetime.strptime(end_date, '%Y-%m-%d')
    
    json_output = {
        "entity": entity,
        "parameter": parameter,
        "start_date": start_date.date().isoformat(),  # Extract date part and convert to ISO format
        "end_date": end_date.date().isoformat()       # Extract date part and convert to ISO format
    }
    return json_output

# Step 5: Testing and Validation
def test_application():
    # Test user queries
    user_queries = [
        "What is the revenue of Apple from 2023-01-01 to 2023-12-31?",
        "Show me the profit margin of Amazon for the last quarter.",
        "Compare the GMV of Alibaba and Walmart from 2022-01-01 to 2022-12-31."
    ]
    
    for query in user_queries:
        st.write("User Query:", query)
        parsed_query = parse_user_query(query)
        start_date, end_date = parsed_query["start_date"], parsed_query["end_date"]
        gpt_response = query_gpt(query)
        st.write("GPT-3.5 Response:", gpt_response)
        json_output = generate_json_output(parsed_query["entity"], parsed_query["parameter"], start_date, end_date)
        st.write("JSON Output:", json_output)
        st.write("-" * 50)

# Initialize Streamlit app
st.set_page_config(page_title="Company Performance Metrics Chatbot")
st.header("Company Performance Metrics Chatbot")

# Run the application
test_application()