import streamlit as st
from utils.vector_store import VectorStore
from utils.document_processor import DocumentProcessor
from utils.case_manager import CaseManager
from utils.legal_notebook_interface import LegalNotebookInterface
from utils.case_manager_interface import CaseManagerInterface
from utils.admin_interface import AdminInterface
from datetime import datetime
import os
from pathlib import Path
import nltk
import spacy

# Ensure this is the first Streamlit command
st.set_page_config(
    page_title="Synaptyx.AI - Legal Agent : Empowering Legal Intelligence",
    page_icon="⚖️",
    layout="wide",
    initial_sidebar_state="expanded"
)

def init_huggingface_directories():
    """Initialize directory structure for Hugging Face Spaces."""
    if os.environ.get('SPACE_ID'):
        base_dir = Path("/data")
    else:
        base_dir = Path(os.getcwd()) / "data"
    
    directories = {
        'data': base_dir,
        'cases': base_dir / "cases",
        'vectors': base_dir / "vectors",
        'nltk_data': base_dir / "nltk_data",
        'temp': base_dir / "temp",
        'logs': base_dir / "logs",
        'prompts': base_dir / "prompts",
        'ontology': base_dir / "ontology"
    }
    
    # Create directories
    for dir_name, dir_path in directories.items():
        dir_path.mkdir(parents=True, exist_ok=True)
        st.sidebar.write(f"✓ {dir_name} directory ready")
        
    return directories

@st.cache_resource(show_spinner=False)
def init_components():
    """Initialize all components silently."""
    try:
        # Set up directories
        directories = init_huggingface_directories()

        # Initialize NLTK
        nltk.data.path.append(str(directories['nltk_data']))
        resources = ['punkt', 'averaged_perceptron_tagger', 'maxent_ne_chunker', 'words', 'stopwords']
        for resource in resources:
            try:
                nltk.download(resource, download_dir=str(directories['nltk_data']), quiet=True)
            except Exception:
                pass

        # Initialize spaCy
        try:
            nlp = spacy.load("en_core_web_sm")
        except OSError:
            os.system("python -m spacy download en_core_web_sm > /dev/null 2>&1")
            nlp = spacy.load("en_core_web_sm")

        # Initialize components with proper paths
        case_manager = CaseManager(base_path=str(directories['cases']))
        vector_store = VectorStore(storage_path=str(directories['vectors']))
        doc_processor = DocumentProcessor(base_path=str(directories['data']))
        
        return case_manager, vector_store, doc_processor

    except Exception as e:
        st.error(f"Failed to initialize components: {str(e)}")
        raise

def main():
    try:
        # Custom styles
        st.markdown("""
            <style>
            .main-header {
                font-size: 40px;
                font-weight: bold;
                color: #2B547E;
                margin-bottom: 0px;
            }
            .main-tagline {
                font-size: 18px;
                color: #6A7B8C;
                margin-top: 0px;
                font-style: italic;
            }
            .stButton button {
                width: 100%;
            }
            .status-message {
                padding: 1rem;
                border-radius: 0.5rem;
                margin: 1rem 0;
            }
            .status-success {
                background-color: #d1e7dd;
                color: #0f5132;
            }
            .case-card {
                padding: 1rem;
                border-radius: 0.5rem;
                border: 1px solid #e0e0e0;
                margin-bottom: 1rem;
            }
            </style>
            """, unsafe_allow_html=True)

        # Initialize components
        case_manager, vector_store, doc_processor = init_components()

        # Session state for user authentication
        if 'user_authenticated' not in st.session_state:
            st.session_state.user_authenticated = False
            st.session_state.current_user = None
        
        # Header
        st.markdown("<h1 class='main-header'>Synaptyx.AI - Legal AI agent based solution</h1>", unsafe_allow_html=True)
        st.markdown("<p class='main-tagline'>Empowering Legal Intelligence: Automate, Analyze, Act.</p>", unsafe_allow_html=True)

        # Sidebar navigation with admin access
        if st.session_state.user_authenticated:
            is_admin = st.session_state.current_user == 'admin'
            
            navigation_options = ["📁 Case Manager", "📄 Document Analysis", "🤖 Legal Assistant"]
            if is_admin:
                navigation_options.append("⚙️ Admin")
            
            tab = st.sidebar.radio("Navigation", navigation_options)
            
            # Add logout button
            if st.sidebar.button("Logout"):
                st.session_state.user_authenticated = False
                st.session_state.current_user = None
                st.rerun()

            if tab == "📁 Case Manager":
                case_manager_interface = CaseManagerInterface(case_manager, vector_store, doc_processor)
                case_manager_interface.render()

            elif tab == "📄 Document Analysis":
                notebook = LegalNotebookInterface(case_manager, vector_store, doc_processor)
                notebook.render()

            elif tab == "🤖 Legal Assistant":
                notebook = LegalNotebookInterface(case_manager, vector_store, doc_processor)
                notebook.render()

            elif tab == "⚙️ Admin" and is_admin:
                admin_interface = AdminInterface(case_manager, vector_store, doc_processor)
                admin_interface.render()

        else:
            # Login form
            st.sidebar.title("Login")
            with st.sidebar.form("login_form"):
                username = st.text_input("Username")
                password = st.text_input("Password", type="password")
                submit = st.form_submit_button("Login")

                if submit:
                    if username == "admin" and password == "demo":
                        st.session_state.user_authenticated = True
                        st.session_state.current_user = "admin"
                        st.rerun()
                    # Here you can add more user authentication logic
                    else:
                        st.error("Invalid credentials")

            # Show welcome message for non-authenticated users
            st.markdown("""
            ## Welcome to Synaptyx.AI Legal Agent
            Please login to access the legal analysis tools.
            
            ### Available Features:
            - 📁 Case Management
            - 📄 Document Analysis
            - 🤖 Legal Assistant
            - ⚙️ Admin Panel (for administrators)
            """)

    except Exception as e:
        st.error("An unexpected error occurred")
        st.exception(e)

if __name__ == "__main__":
    main()