import streamlit as st
import pandas as pd
from processor import DataProcessor
from llm_handler import DDoSInference

class DDoSResistanceHelper:
    def __init__(self):
        # Set up Streamlit page configuration
        st.set_page_config(
            page_title="DDoS Resistance Helper | LLM Network Analyzer",
            page_icon="🛡️",
            layout="wide"
        )

        # Initialize session state
        self.setup_session_state()

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

        # Define themes
        self.themes = {
            'Light': {
                'base': 'light',
                'primaryColor': '#3B82F6',
                'backgroundColor': '#FFFFFF',
                'secondaryBackgroundColor': '#F3F4F6',
                'textColor': '#000000'
            },
            'Dark': {
                'base': 'dark',
                'primaryColor': '#10B981',
                'backgroundColor': '#1F2937',
                'secondaryBackgroundColor': '#374151',
                'textColor': '#FFFFFF'
            },
            'Neon': {
                'base': 'dark',
                'primaryColor': '#22D3EE',
                'backgroundColor': '#000000',
                'secondaryBackgroundColor': '#1F2937',
                'textColor': '#22D3EE'
            }
        }

    def setup_session_state(self):
        """Initialize Streamlit session state variables."""
        st.write("Initializing session state...")
        session_keys = ['chat_history', 'current_file', 'analysis_complete', 'selected_theme', 'preview_data']
        for key in session_keys:
            if key not in st.session_state:
                st.session_state[key] = [] if key == 'chat_history' else None
        st.write("Session state initialized.")

    def apply_theme(self):
        """Apply selected theme to Streamlit app."""
        theme_name = st.session_state.get('selected_theme', 'Light')
        theme = self.themes[theme_name]
        st.markdown(f"""
        <style>
        .stApp {{
            background-color: {theme['backgroundColor']};
            color: {theme['textColor']};
        }}
        </style>
        """, unsafe_allow_html=True)

    def render_theme_selector(self):
        """Render theme selection icons on the top-right corner."""
        st.sidebar.empty()
        st.write(f'<div style="position: absolute; top: 15px; right: 15px;">', unsafe_allow_html=True)
        if st.button("🌞 Light", key="light_theme"):
            st.session_state.selected_theme = 'Light'
            self.apply_theme()
        if st.button("🌙 Dark", key="dark_theme"):
            st.session_state.selected_theme = 'Dark'
            self.apply_theme()
        if st.button("💡 Neon", key="neon_theme"):
            st.session_state.selected_theme = 'Neon'
            self.apply_theme()

    def render_file_upload(self):
        """Render file upload section with an upload icon in the main header."""
        st.title("DDoS Resistance Helper | LLM Network Analyzer")
        uploaded_file = st.file_uploader("Upload Data File 📂", type=['csv'])
        if uploaded_file:
            st.session_state.current_file = uploaded_file
            st.session_state.preview_data = pd.read_csv(uploaded_file).head(5)
            if st.button("🚀 Process Data"):
                self.process_uploaded_file(uploaded_file)

    def process_uploaded_file(self, file):
        """Process the uploaded file and perform analysis."""
        try:
            st.write("Processing file...")
            with st.spinner("Analyzing data..."):
                # Preprocess the uploaded file
                df = self.processor.preprocess_data(file)

                # Perform LLM analysis
                results = self.llm_handler.analyze_data(df)

                st.session_state.analysis_complete = True
                st.success("Analysis completed!")
        except Exception as e:
            st.error(f"Error: {e}")

    def render_chat_interface(self):
        """Render chat interface in the sidebar."""
        st.sidebar.header("💬 Chat Interface")
        for msg in st.session_state.chat_history:
            st.sidebar.write(f"[{msg['role'].capitalize()}]: {msg['content']}")
        new_chat = st.sidebar.text_input("Ask a question...")
        if new_chat:
            st.session_state.chat_history.append({"role": "user", "content": new_chat})
            response = self.llm_handler.get_chat_response(new_chat)
            st.session_state.chat_history.append({"role": "assistant", "content": response})

    def render_analysis_interface(self):
        """Render analysis results interface."""
        st.subheader("📋 Data Preview")
        if st.session_state.preview_data is not None:
            st.dataframe(st.session_state.preview_data)
        if st.session_state.analysis_complete:
            st.subheader("🔍 Analysis Results")
            st.info("Analysis complete. Visualizations and results here.")

    def main(self):
        """Main application runner."""
        self.apply_theme()
        self.render_theme_selector()
        self.render_file_upload()
        self.render_analysis_interface()
        self.render_chat_interface()

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