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

class DDoSResistanceHelper:
    def __init__(self):
        """Initialize the DDoS Resistance Helper app."""
        self.setup_streamlit_page()
        self.setup_session_state()
        self.processor = DataProcessor()
        self.llm_handler = DDoSInference()
        self.themes = self.define_themes()

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

    def setup_session_state(self):
        """Initialize necessary session state variables."""
        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 if key != 'selected_theme' else 'Light'

    def define_themes(self):
        """Define available themes for the app."""
        return {
            '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 apply_theme(self):
        """Apply the selected theme to the Streamlit app."""
        theme = self.themes.get(st.session_state.selected_theme, self.themes['Light'])
        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 buttons."""
        col1, col2, col3 = st.columns(3)
        
        with col1:
            if st.button("🌞 Light"):
                st.session_state.selected_theme = 'Light'
                self.apply_theme()
        
        with col2:
            if st.button("🌙 Dark"):
                st.session_state.selected_theme = 'Dark'
                self.apply_theme()
        
        with col3:
            if st.button("💡 Neon"):
                st.session_state.selected_theme = 'Neon'
                self.apply_theme()

    def render_file_upload(self):
        """Render file upload section."""
        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:
            with st.spinner("Processing file..."):
                df = self.processor.preprocess_data(file)
                results = self.llm_handler.process_dataset()  # Assuming this method returns results
                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."""
        if st.session_state.preview_data is not None:
            st.subheader("📋 Data Preview")
            st.dataframe(st.session_state.preview_data)
        
        if st.session_state.analysis_complete:
            st.subheader("🔍 Analysis Results")
            # Display analysis results (e.g., visualizations)

    def main(self):
        """Run the main application."""
        self.apply_theme()
        
        # Layout sections
        cols = st.columns([2, 1])
        
        with cols[0]:
            self.render_file_upload()
        
        with cols[1]:
            self.render_analysis_interface()
        
        self.render_chat_interface()

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