import streamlit as st
import pandas as pd
from pathlib import Path
from processor import DataProcessor
from llm_handler import LLMHandler

class DDoSResistanceHelper:
    def __init__(self):
        # Configure Streamlit app
        st.set_page_config(
            page_title="DDoS Resistance Helper URF LLM Network Analyzer",
            page_icon=":shield:",
            layout="wide",
            initial_sidebar_state="expanded"
        )

        # Initialize session state
        self.initialize_session_state()

        # Initialize processor and LLM handler
        self.processor = DataProcessor()
        self.llm_handler = LLMHandler()

    def initialize_session_state(self):
        """Set up Streamlit session state variables."""
        session_keys = [
            'current_file', 'preprocessed_data', 'analysis_results', 'chat_history'
        ]
        for key in session_keys:
            if key not in st.session_state:
                st.session_state[key] = None if key != 'chat_history' else []

    def render_top_bar(self):
        """Render the top bar with theme and upload options."""
        col1, col2 = st.columns([8, 2])

        with col1:
            st.title("🛡️ DDoS Resistance Helper URF LLM Network Analyzer")

        with col2:
            st.markdown("### Theme Selector")
            if st.button("Light"):
                st.markdown("<style>.stApp { background-color: #ffffff; }</style>", unsafe_allow_html=True)
            elif st.button("Dark"):
                st.markdown("<style>.stApp { background-color: #1f1f1f; color: white; }</style>", unsafe_allow_html=True)

    def render_file_upload(self):
        """Render the file upload component."""
        uploaded_file = st.file_uploader("Upload Network Traffic Data (CSV)", type=["csv"],
                                         label_visibility="collapsed")
        if uploaded_file:
            try:
                df = pd.read_csv(uploaded_file)
                st.session_state.current_file = df
                st.success("File uploaded successfully!")
            except Exception as e:
                st.error(f"Error reading file: {e}")

    def render_analysis(self):
        """Render the analysis results."""
        if st.session_state.current_file is None:
            st.info("Please upload a CSV file to start analysis.")
            return

        # Preprocess the data
        st.subheader("Preprocessing Data")
        with st.spinner("Preprocessing data..."):
            try:
                preprocessed_data = self.processor.preprocess_data(st.session_state.current_file)
                st.session_state.preprocessed_data = preprocessed_data
                st.success("Data preprocessed successfully!")
            except Exception as e:
                st.error(f"Error during preprocessing: {e}")

        # Perform LLM analysis
        st.subheader("Performing LLM Analysis")
        with st.spinner("Analyzing data with LLM..."):
            try:
                results = self.llm_handler.analyze_data(st.session_state.preprocessed_data)
                st.session_state.analysis_results = results
                st.success("Analysis completed successfully!")
            except Exception as e:
                st.error(f"Error during LLM analysis: {e}")

        # Show results
        if st.session_state.analysis_results is not None:
            st.subheader("Analysis Results")
            st.dataframe(st.session_state.analysis_results)
            csv_path = Path("~/.dataset/PROBABILITY_OF_EACH_ROW_DDOS_AND_BENGNIN.csv").expanduser()
            st.download_button("Download Results as CSV", csv_path.read_bytes(), "analysis_results.csv")

    def render_chat_interface(self):
        """Render a chat interface for interacting with the LLM."""
        st.sidebar.header("💬 Chat Interface")
        # Display chat history
        for message in st.session_state.chat_history:
            with st.chat_message(message['role']):
                st.write(message['content'])

        # Chat input
        if prompt := st.sidebar.text_input("Ask about the analysis or mitigation steps..."):
            # Add user message to chat history
            st.session_state.chat_history.append({
                'role': 'user',
                'content': prompt
            })

            # Get LLM response
            response = self.llm_handler.get_chat_response(prompt)

            # Add LLM response to chat history
            st.session_state.chat_history.append({
                'role': 'assistant',
                'content': response
            })

    def run(self):
        """Run the Streamlit app."""
        self.render_top_bar()
        self.render_file_upload()
        self.render_analysis()
        self.render_chat_interface()

if __name__ == "__main__":
    app = DDoSResistanceHelper()
    app.run()