utils/admin_interface.py

import streamlit as st  # Add this at the top
from datetime import datetime
import json
import os
from pathlib import Path

class AdminInterface:
    def __init__(self, case_manager, vector_store, document_processor):
        self.case_manager = case_manager
        self.vector_store = vector_store
        self.document_processor = document_processor
        
        # Initialize session state for admin
        if 'admin_authenticated' not in st.session_state:
            st.session_state.admin_authenticated = False
        if 'users' not in st.session_state:
            st.session_state.users = {}
        if 'prompts' not in st.session_state:
            st.session_state.prompts = self._load_default_prompts()
            
    def _load_default_prompts(self):
        return {
            "system_prompt": """You are a legal AI assistant analyzing legal documents. 
Follow these guidelines:
1. Provide accurate legal analysis
2. Cite specific sections from documents
3. Maintain formal legal language
4. Consider jurisdictional context""",
            "user_prompt": """Please analyze the following legal question:
{query}

Based on these documents:
{context}

Provide a structured analysis with:
1. Key findings
2. Relevant citations
3. Legal implications"""
        }
            
    def render(self):
        if not st.session_state.admin_authenticated:
            self._render_login()
        else:
            self._render_admin_panel()
            
    def _render_login(self):
        st.title("🔒 Admin Login")
        
        with st.form("admin_login"):
            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.admin_authenticated = True
                    st.success("Login successful!")
                    st.rerun()
                else:
                    st.error("Invalid credentials")
                    
    def _render_admin_panel(self):
        st.title("⚙️ Admin Panel")
        
        tab1, tab2, tab3, tab4 = st.tabs([
            "Prompt Management",
            "Ontology Editor",
            "User Management",
            "System Settings"
        ])
        
        with tab1:
            self._render_prompt_management()
            
        with tab2:
            self._render_ontology_editor()
            
        with tab3:
            self._render_user_management()
            
        with tab4:
            self._render_system_settings()
            
    def _render_prompt_management(self):
        st.subheader("🔤 Prompt Templates")
        
        # System prompt
        st.markdown("### System Prompt")
        new_system_prompt = st.text_area(
            "System Instructions",
            value=st.session_state.prompts["system_prompt"],
            height=200
        )
        
        # User prompt
        st.markdown("### User Prompt Template")
        new_user_prompt = st.text_area(
            "Query Template",
            value=st.session_state.prompts["user_prompt"],
            height=200
        )
        
        if st.button("Save Prompts"):
            st.session_state.prompts = {
                "system_prompt": new_system_prompt,
                "user_prompt": new_user_prompt
            }
            # Save to file
            self._save_prompts()
            st.success("Prompts updated successfully!")
            
    def _render_ontology_editor(self):
        st.subheader("📚 Legal Ontology Editor")
        
        # Load current ontology
        ontology = self._load_ontology()
        
        # Display ontology editor
        st.json(ontology)
        
        # Upload new ontology
        uploaded_file = st.file_uploader("Upload new ontology (JSON)", type="json")
        if uploaded_file:
            try:
                new_ontology = json.load(uploaded_file)
                self._save_ontology(new_ontology)
                st.success("Ontology updated successfully!")
                st.rerun()
            except Exception as e:
                st.error(f"Error updating ontology: {str(e)}")
                
        # Add new concept
        st.markdown("### Add New Concept")
        with st.form("add_concept"):
            concept_label = st.text_input("Concept Label")
            concept_type = st.selectbox("Type", ["LegalConcept", "Jurisdiction", "Process"])
            concept_desc = st.text_area("Description")
            related = st.text_input("Related Concepts (comma-separated)")
            
            if st.form_submit_button("Add Concept"):
                try:
                    new_concept = {
                        "@id": f"concept:{concept_label.replace(' ', '')}",
                        "@type": f"vocab:{concept_type}",
                        "rdfs:label": concept_label,
                        "rdfs:comment": concept_desc,
                        "vocab:relatedConcepts": [x.strip() for x in related.split(",")]
                    }
                    ontology["@graph"].append(new_concept)
                    self._save_ontology(ontology)
                    st.success("Concept added successfully!")
                    st.rerun()
                except Exception as e:
                    st.error(f"Error adding concept: {str(e)}")
                    
    def _render_user_management(self):
        st.subheader("👥 User Management")
        
        # List existing users
        st.markdown("### Current Users")
        for username, user_data in st.session_state.users.items():
            col1, col2, col3 = st.columns([2, 2, 1])
            with col1:
                st.write(f"**{username}**")
            with col2:
                st.write(f"Access: {', '.join(user_data['case_access'])}")
            with col3:
                if st.button("🗑️", key=f"del_user_{username}"):
                    del st.session_state.users[username]
                    st.success(f"User {username} deleted")
                    st.rerun()
        
        # Add new user
        st.markdown("### Add User")
        with st.form("add_user"):
            new_username = st.text_input("Username")
            new_password = st.text_input("Password", type="password")
            cases = self.case_manager.get_all_cases()
            case_access = st.multiselect(
                "Case Access",
                options=[case['id'] for case in cases],
                format_func=lambda x: next(c['title'] for c in cases if c['id'] == x)
            )
            
            if st.form_submit_button("Add User"):
                if new_username in st.session_state.users:
                    st.error("Username already exists")
                else:
                    st.session_state.users[new_username] = {
                        "password": new_password,
                        "case_access": case_access
                    }
                    st.success(f"User {new_username} added successfully!")
                    st.rerun()
                    
    def _render_system_settings(self):
        st.subheader("⚙️ System Settings")
        
        # Vector store settings
        st.markdown("### Vector Store Settings")
        with st.form("vector_store_settings"):
            chunk_size = st.number_input(
                "Document Chunk Size",
                min_value=100,
                max_value=2000,
                value=500,
                help="Number of characters per document chunk"
            )
            
            top_k = st.number_input(
                "Search Results (Top K)",
                min_value=1,
                max_value=10,
                value=3,
                help="Number of similar chunks to return in search"
            )
            
            similarity_threshold = st.slider(
                "Similarity Threshold",
                min_value=0.0,
                max_value=1.0,
                value=0.7,
                help="Minimum similarity score for chunk matching"
            )
            
            if st.form_submit_button("Save Settings"):
                self._save_system_settings({
                    "chunk_size": chunk_size,
                    "top_k": top_k,
                    "similarity_threshold": similarity_threshold
                })
                st.success("Settings saved successfully!")

        # Model settings
        st.markdown("### Model Settings")
        with st.form("model_settings"):
            model = st.selectbox(
                "Claude Model",
                ["claude-3-sonnet-20240229", "claude-3-opus-20240229", "claude-3-haiku-20240229"],
                help="Select Claude model version"
            )
            
            temperature = st.slider(
                "Temperature",
                min_value=0.0,
                max_value=1.0,
                value=0.2,
                help="Controls response creativity"
            )
            
            max_tokens = st.number_input(
                "Max Response Tokens",
                min_value=100,
                max_value=4096,
                value=1500,
                help="Maximum length of responses"
            )
            
            if st.form_submit_button("Save Model Settings"):
                self._save_model_settings({
                    "model": model,
                    "temperature": temperature,
                    "max_tokens": max_tokens
                })
                st.success("Model settings saved!")

        # Backup and restore
        st.markdown("### 💾 Backup & Restore")
        col1, col2 = st.columns(2)
        
        with col1:
            if st.button("📤 Export All Data"):
                try:
                    data = self._export_all_data()
                    st.download_button(
                        "Download Backup",
                        data=json.dumps(data, indent=2),
                        file_name=f"suomoto_backup_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json",
                        mime="application/json"
                    )
                except Exception as e:
                    st.error(f"Error exporting data: {str(e)}")

        with col2:
            uploaded_file = st.file_uploader("📥 Restore from Backup", type="json")
            if uploaded_file:
                try:
                    backup_data = json.load(uploaded_file)
                    if st.button("Restore Data"):
                        self._restore_from_backup(backup_data)
                        st.success("Data restored successfully!")
                        st.rerun()
                except Exception as e:
                    st.error(f"Error restoring data: {str(e)}")

    def _save_prompts(self):
        """Save prompts to JSON file."""
        try:
            with open(os.path.join('/data', 'prompts.json'), 'w') as f:
                json.dump(st.session_state.prompts, f, indent=2)
        except Exception as e:
            st.error(f"Error saving prompts: {str(e)}")

    def _load_ontology(self):
        """Load current ontology from file."""
        try:
            with open(os.path.join('/data', 'legal_ontology.json'), 'r') as f:
                return json.load(f)
        except Exception:
            return {"@graph": []}

    def _save_ontology(self, ontology: dict):
        """Save updated ontology to file."""
        try:
            with open(os.path.join('/data', 'legal_ontology.json'), 'w') as f:
                json.dump(ontology, f, indent=2)
        except Exception as e:
            st.error(f"Error saving ontology: {str(e)}")

    def _save_system_settings(self, settings: dict):
        """Save system settings to file."""
        try:
            with open(os.path.join('/data', 'system_settings.json'), 'w') as f:
                json.dump(settings, f, indent=2)
        except Exception as e:
            st.error(f"Error saving settings: {str(e)}")

    def _save_model_settings(self, settings: dict):
        """Save model settings to file."""
        try:
            with open(os.path.join('/data', 'model_settings.json'), 'w') as f:
                json.dump(settings, f, indent=2)
        except Exception as e:
            st.error(f"Error saving model settings: {str(e)}")

    def _export_all_data(self):
        """Export all system data for backup."""
        return {
            "prompts": st.session_state.prompts,
            "ontology": self._load_ontology(),
            "users": st.session_state.users,
            "system_settings": self._load_system_settings(),
            "model_settings": self._load_model_settings(),
            "export_date": datetime.now().isoformat()
        }

    def _restore_from_backup(self, backup_data: dict):
        """Restore system from backup data."""
        try:
            # Validate backup data
            required_keys = ["prompts", "ontology", "users", "system_settings", "model_settings"]
            if not all(key in backup_data for key in required_keys):
                raise ValueError("Invalid backup data format")

            # Restore each component
            st.session_state.prompts = backup_data["prompts"]
            self._save_prompts()
            
            self._save_ontology(backup_data["ontology"])
            st.session_state.users = backup_data["users"]
            
            self._save_system_settings(backup_data["system_settings"])
            self._save_model_settings(backup_data["model_settings"])
            
        except Exception as e:
            raise Exception(f"Error restoring from backup: {str(e)}")

    def _load_system_settings(self):
        """Load system settings from file."""
        try:
            with open(os.path.join('/data', 'system_settings.json'), 'r') as f:
                return json.load(f)
        except Exception:
            return {
                "chunk_size": 500,
                "top_k": 3,
                "similarity_threshold": 0.7
            }

    def _load_model_settings(self):
        """Load model settings from file."""
        try:
            with open(os.path.join('/data', 'model_settings.json'), 'r') as f:
                return json.load(f)
        except Exception:
            return {
                "model": "claude-3-sonnet-20240229",
                "temperature": 0.2,
                "max_tokens": 1500
            }