import streamlit as st
import pandas as pd
import numpy as np
import plotly.express as px
import plotly.graph_objects as go
from ydata_profiling import ProfileReport
from streamlit_pandas_profiling import st_profile_report
import os
import requests
import json
from datetime import datetime
import re
import tempfile
from scipy import stats
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import StandardScaler, LabelEncoder, OneHotEncoder
from sklearn.decomposition import PCA
import streamlit.components.v1 as components
from io import StringIO
from dotenv import load_dotenv
from flask import Flask, request, jsonify
from openai import OpenAI
import threading
from sentence_transformers import SentenceTransformer

# Load environment variables
load_dotenv()

# Initialize Flask app
flask_app = Flask(__name__)
FLASK_PORT = 5000  # Internal port for Flask, not exposed externally

# Initialize OpenAI client
api_key = os.getenv("OPENAI_API_KEY")
if not api_key:
    st.error("OPENAI_API_KEY not set. Please configure it in the Hugging Face Space secrets.")
    st.stop()
client = OpenAI(api_key=api_key)

# Flask RAG Endpoint
@flask_app.route('/rag_chat', methods=['POST'])
def rag_chat():
    data = request.get_json()
    user_input = data.get('user_input', '')
    app_mode = data.get('app_mode', 'Data Upload')
    dataset_text = data.get('dataset_text', '')

    # RAG Logic: Use dataset_text as retrieval context
    system_prompt = (
        "You are an AI assistant in Data-Vision Pro, a data analysis app with RAG capabilities. "
        "The app has three pages:\n"
        "- **Data Upload**: Upload CSV/XLSX files, view stats, or generate reports.\n"
        "- **Data Cleaning**: Clean data (e.g., handle missing values, encode variables).\n"
        "- **EDA**: Visualize data (e.g., scatter plots, histograms).\n"
        f"The user is on the '{app_mode}' page.\n"
    )

    if dataset_text:
        system_prompt += (
            "Using the following dataset context, augment your response:\n"
            f"{dataset_text}\n"
            "Answer based on this data where relevant, otherwise provide general assistance."
        )
    else:
        system_prompt += "No dataset is loaded. Assist based on app functionality."

    try:
        response = client.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_input}
            ],
            max_tokens=100,  # Increased for RAG context
            temperature=0.7
        )
        return jsonify({"response": response.choices[0].message.content})
    except Exception as e:
        return jsonify({"error": str(e)}), 500

# Run Flask in a background thread
def run_flask():
    flask_app.run(host='0.0.0.0', port=FLASK_PORT, debug=False, use_reloader=False)

# Start Flask thread
flask_thread = threading.Thread(target=run_flask, daemon=True)
flask_thread.start()

# Helper Functions
def enhance_section_title(title):
    st.markdown(f"<h2 style='border-bottom: 2px solid #ccc; padding-bottom: 5px;'>{title}</h2>", unsafe_allow_html=True)

def update_cleaned_data(df):
    st.session_state.cleaned_data = df
    if 'data_versions' not in st.session_state:
        st.session_state.data_versions = [st.session_state.raw_data.copy()]
    st.session_state.data_versions.append(df.copy())
    st.success("✅ Action completed successfully!")
    st.rerun()

def convert_csv_to_json_and_text(df):
    """Convert DataFrame to JSON and then to plain text."""
    json_data = df.to_json(orient="records")
    data_dict = json.loads(json_data)
    text_summary = f"Dataset Summary: {df.shape[0]} rows, {df.shape[1]} columns\n"
    text_summary += f"Missing Values: {df.isna().sum().sum()}\n"
    text_summary += "Columns:\n"
    for col in df.columns:
        text_summary += f"- {col} ({df[col].dtype}): "
        if pd.api.types.is_numeric_dtype(df[col]):
            text_summary += f"Mean={df[col].mean():.2f}, Min={df[col].min()}, Max={df[col].max()}"
        else:
            text_summary += f"Unique={df[col].nunique()}, Top={df[col].mode()[0] if not df[col].mode().empty else 'N/A'}"
        text_summary += f", Missing={df[col].isna().sum()}\n"
    return text_summary

def get_chatbot_response(user_input, app_mode, dataset_text=""):
    """Send request to internal Flask RAG endpoint."""
    payload = {
        "user_input": user_input,
        "app_mode": app_mode,
        "dataset_text": dataset_text
    }
    try:
        response = requests.post(f"http://localhost:{FLASK_PORT}/rag_chat", json=payload, timeout=5)
        response.raise_for_status()
        return response.json().get("response", "Error: No response from server")
    except requests.exceptions.RequestException as e:
        return f"Error: Could not connect to RAG server. {str(e)}"

# Streamlit App
# Sidebar Navigation
with st.sidebar:
    st.title("🔮 Data-Vision Pro")
    st.markdown("Your AI-powered data analysis suite with RAG.")
    st.markdown("---")
    app_mode = st.selectbox(
        "Navigation",
        ["Data Upload", "Data Cleaning", "EDA"],
        format_func=lambda x: f"📌 {x}"
    )
    if app_mode == "Data Upload":
        st.info("⬆️ Upload your CSV or XLSX dataset to begin.")
    elif app_mode == "Data Cleaning":
        st.info("🧹 Clean and preprocess your data using various tools.")
    elif app_mode == "EDA":
        st.info("🔍 Explore your data visually and statistically.")

    st.markdown("---")
    st.markdown("**Note**: Requires dependencies in `requirements.txt`.")
    if 'cleaned_data' in st.session_state:
        csv = st.session_state.cleaned_data.to_csv(index=False)
        st.download_button(
            label="Download Cleaned Data as CSV",
            data=csv,
            file_name='cleaned_data.csv',
            mime='text/csv',
        )
    st.markdown("Created by Calvin Allen-Crawford")
    st.markdown("v1.0 | © 2025")

# Main App Pages
if app_mode == "Data Upload":
    st.title("📤 Data Upload & Profiling")
    st.header("Upload Your Dataset")
    st.write("Supported formats: CSV, XLSX")

    if 'raw_data' not in st.session_state:
        st.info("It looks like no dataset has been uploaded yet. Would you like to upload a CSV or XLSX file?")
    
    uploaded_file = st.file_uploader("Choose a file", type=["csv", "xlsx"], key="file_uploader")
    if uploaded_file:
        st.session_state.pop('raw_data', None)
        st.session_state.pop('cleaned_data', None)
        st.session_state.pop('data_versions', None)
        try:
            if uploaded_file.name.endswith('.csv'):
                df = pd.read_csv(uploaded_file)
            else:
                df = pd.read_excel(uploaded_file)
            if df.empty:
                st.error("Uploaded file is empty.")
                st.stop()
            st.session_state.raw_data = df
            st.session_state.dataset_text = convert_csv_to_json_and_text(df)
            if 'data_versions' not in st.session_state:
                st.session_state.data_versions = [df.copy()]
            col1, col2, col3 = st.columns(3)
            with col1: st.metric("Rows", df.shape[0])
            with col2: st.metric("Columns", df.shape[1])
            with col3: st.metric("Missing Values", df.isna().sum().sum())
            if st.checkbox("Show Data Preview"):
                st.dataframe(df.head(10), use_container_width=True)
            if st.button("Generate Full Profile Report"):
                with st.spinner("Generating report..."):
                    pr = ProfileReport(df, explorative=True)
                    st_profile_report(pr)
            st.success("✅ Data loaded successfully!")
        except Exception as e:
            st.error(f"An error occurred: {str(e)}")

elif app_mode == "Data Cleaning":
    st.title("🧹 Smart Data Cleaning")
    st.header("Preprocess and Transform Your Data")
    if 'raw_data' not in st.session_state:
        st.warning("Please upload data first in the Data Upload section.")
        st.stop()
    if 'cleaned_data' not in st.session_state:
        st.session_state.cleaned_data = st.session_state.raw_data.copy()
    df = st.session_state.cleaned_data.copy()

    enhance_section_title("📊 Data Health Dashboard")
    with st.expander("Explore Data Health Metrics", expanded=True):
        col1, col2, col3 = st.columns(3)
        with col1: st.metric("Columns", len(df.columns))
        with col2: st.metric("Rows", len(df))
        with col3: st.metric("Missing Values", df.isna().sum().sum())
        if st.button("Generate Detailed Health Report"):
            with st.spinner("Generating report..."):
                profile = ProfileReport(df, minimal=True)
                st_profile_report(profile)
        if 'data_versions' in st.session_state and len(st.session_state.data_versions) > 1:
            if st.button("Undo Last Action"):
                st.session_state.data_versions.pop()
                st.session_state.cleaned_data = st.session_state.data_versions[-1].copy()
                st.session_state.dataset_text = convert_csv_to_json_and_text(st.session_state.cleaned_data)
                st.rerun()

elif app_mode == "EDA":
    st.title("🔍 Interactive Data Explorer")
    if 'cleaned_data' not in st.session_state:
        st.warning("Please upload and clean data first.")
        st.stop()
    df = st.session_state.cleaned_data.copy()

    enhance_section_title("Dataset Overview")
    with st.container():
        col1, col2, col3, col4 = st.columns(4)
        col1.metric("Total Rows", df.shape[0])
        col2.metric("Total Columns", df.shape[1])
        missing_percentage = df.isna().sum().sum() / df.size * 100
        col3.metric("Missing Values", f"{df.isna().sum().sum()} ({missing_percentage:.1f}%)")
        col4.metric("Duplicates", df.duplicated().sum())

# Chatbot Section
st.markdown("---")
st.subheader("💬 AI Chatbot Assistant (RAG Enabled)")
st.info("Ask me about the app or your data! Try: 'What can I do here?' or 'What’s in the dataset?'")
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

for message in st.session_state.chat_history:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

user_input = st.chat_input("Ask me anything about the app or your data...")
if user_input:
    st.session_state.chat_history.append({"role": "user", "content": user_input})
    with st.chat_message("user"):
        st.markdown(user_input)

    with st.spinner("Thinking with RAG..."):
        dataset_text = st.session_state.get("dataset_text", "")
        response = get_chatbot_response(user_input, app_mode, dataset_text)
        st.session_state.chat_history.append({"role": "assistant", "content": response})
    with st.chat_message("assistant"):
        st.markdown(response)