from openai import OpenAI
import streamlit as st
import os
from datetime import datetime
import pandas as pd
import json
import xml.etree.ElementTree as ET
from io import StringIO, BytesIO
from time import sleep
from tabulate import tabulate

# Load API key securely
API_KEY = os.getenv("NV_API_KEY", "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")
if not API_KEY:
    st.error("API key is missing! Please set NV_API_KEY as an environment variable.")
    st.stop()

client = OpenAI(
    base_url="https://integrate.api.nvidia.com/v1",
    api_key=API_KEY
)

st.set_page_config(page_title="Nemotron 4 340B", layout="wide")
st.title("Nemotron 4 340B")

# Sidebar content
with st.sidebar:
    st.markdown("This is an advanced chatbot focused on coding and file analysis. Supported by Nazmul Hasan Nihal.")
    if st.button("Clear Session"):
        st.session_state.clear()
    st.write(f"Copyright 2023-{datetime.now().year} Present Nazmul Hasan Nihal")

# Initialize session state
if "openai_model" not in st.session_state:
    st.session_state['openai_model'] = "nvidia/nemotron-4-340b-instruct"

if "messages" not in st.session_state:
    st.session_state.messages = [{"role": "system", "content": "You are a helpful coding assistant."}]

# Function to display typing animation
def display_typing_animation(text):
    for char in text:
        st.markdown(char, unsafe_allow_html=True)
        sleep(0.02)

# Function to process uploaded files
def process_uploaded_file(file):
    file_type = file.type
    if file_type == "text/csv":
        df = pd.read_csv(file)
        return f"**CSV File Analysis:**\n{tabulate(df.head(), headers='keys', tablefmt='grid')}"
    elif file_type in ["application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", "application/vnd.ms-excel"]:
        df = pd.read_excel(file)
        return f"**Excel File Analysis:**\n{tabulate(df.head(), headers='keys', tablefmt='grid')}"
    elif file_type == "application/json":
        data = json.load(file)
        return f"**JSON File Analysis:**\n```json\n{json.dumps(data, indent=4)}\n```"
    elif file_type in ["text/xml", "application/xml"]:
        tree = ET.parse(file)
        root = tree.getroot()
        return f"**XML File Root Tag:** {root.tag}\n**Attributes:** {root.attrib}"
    elif file_type == "application/pdf":
        return "PDF uploads are not processed for content. Try text-based analysis."
    else:
        return "Unsupported file format. Please upload CSV, Excel, JSON, XML, or PDF files."

# Display previous messages
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# Handle user input or file uploads
if user_input := st.chat_input("Type a message or upload a file below."):
    st.session_state.messages.append({"role": "user", "content": user_input})
    with st.chat_message("user"):
        st.markdown(user_input)

    # Assistant response
    with st.chat_message("assistant"):
        with st.spinner("The assistant is processing..."):
            try:
                # Generate response
                stream = client.chat.completions.create(
                    model=st.session_state["openai_model"],
                    messages=st.session_state.messages,
                    temperature=0.5,
                    top_p=0.7,
                    max_tokens=1024,
                    stream=True,
                )
                response_chunks = []
                for chunk in stream:
                    if chunk.choices[0].delta.content:
                        response_chunks.append(chunk.choices[0].delta.content)
                        display_typing_animation(chunk.choices[0].delta.content)
                response = "".join(response_chunks)
                st.markdown(response)

                # Save assistant message
                st.session_state.messages.append({"role": "assistant", "content": response})

            except Exception as e:
                st.error(f"An error occurred: {e}")

# Add file upload to chat input
st.markdown("### File Upload:")
uploaded_file = st.file_uploader("Upload your file (CSV, XLSX, JSON, XML, PDF)", type=["csv", "xlsx", "json", "xml", "pdf"])
if uploaded_file:
    with st.chat_message("user"):
        st.markdown(f"**Uploaded File:** {uploaded_file.name}")
    with st.chat_message("assistant"):
        with st.spinner("Processing the uploaded file..."):
            file_analysis = process_uploaded_file(uploaded_file)
            st.markdown(file_analysis)