Code Editor for app.py

import streamlit as st
import pandas as pd
from sqlalchemy import create_engine, Column, Integer, String, DateTime, Boolean, ForeignKey, Enum
from sqlalchemy.orm import Session, relationship, DeclarativeBase
from datetime import datetime
import enum
import plotly.express as px

# Database setup
class Base(DeclarativeBase):
    pass

class UserRole(enum.Enum):
    ADMIN = "admin"
    RESOLVER = "resolver"

class NotificationChannel(enum.Enum):
    EMAIL = "email"
    SMS = "sms"
    SLACK = "slack"

class IncidentStatus(enum.Enum):
    NEW = "new"
    ACKNOWLEDGED = "acknowledged"
    IN_PROGRESS = "in_progress"
    RESOLVED = "resolved"
    ESCALATED = "escalated"

class User(Base):
    __tablename__ = "users"
    id = Column(Integer, primary_key=True)
    username = Column(String, unique=True, nullable=False)
    password = Column(String, nullable=False)
    role = Column(Enum(UserRole), nullable=False)
    email = Column(String)
    phone = Column(String)
    slack_id = Column(String)
    preferences = relationship("UserPreference", back_populates="user")

class UserPreference(Base):
    __tablename__ = "user_preferences"
    id = Column(Integer, primary_key=True)
    user_id = Column(Integer, ForeignKey("users.id"))
    channel = Column(Enum(NotificationChannel))
    is_active = Column(Boolean, default=True)
    user = relationship("User", back_populates="preferences")

class OnCallSchedule(Base):
    __tablename__ = "oncall_schedules"
    id = Column(Integer, primary_key=True)
    primary_resolver_id = Column(Integer, ForeignKey("users.id"))
    secondary_resolver_id = Column(Integer, ForeignKey("users.id"))
    start_date = Column(DateTime)
    end_date = Column(DateTime)

class Incident(Base):
    __tablename__ = "incidents"
    id = Column(Integer, primary_key=True)
    title = Column(String)
    description = Column(String)
    priority = Column(Integer)  # 1 (highest) to 5 (lowest)
    status = Column(Enum(IncidentStatus))
    created_at = Column(DateTime, default=datetime.utcnow)
    assigned_to = Column(Integer, ForeignKey("users.id"))
    resolved_at = Column(DateTime)

# Initialize database
engine = create_engine("sqlite:///incident_management.db")
Base.metadata.create_all(engine)

# Initialize session state
if "logged_in" not in st.session_state:
    st.session_state.logged_in = False
    st.session_state.user_role = None
    st.session_state.user_id = None

def login_user(username, password):
    with Session(engine) as session:
        user = session.query(User).filter_by(username=username, password=password).first()
        if user:
            st.session_state.logged_in = True
            st.session_state.user_role = user.role
            st.session_state.user_id = user.id
            return True
        return False

def show_login():
    st.title("Incident Management System - Login")
    username = st.text_input("Username")
    password = st.text_input("Password", type="password")
    
    if st.button("Login"):
        if login_user(username, password):
            st.success("Login successful!")
            st.rerun()
        else:
            st.error("Invalid credentials!")

def show_resolver_dashboard():
    st.title("Resolver Dashboard")
    
    tabs = st.tabs(["My Incidents", "Preferences"])
    
    with tabs[0]:
        st.header("My Incidents")
        with Session(engine) as session:
            incidents = session.query(Incident).filter_by(assigned_to=st.session_state.user_id).all()
            if incidents:
                incident_data = []
                for incident in incidents:
                    incident_data.append({
                        "ID": incident.id,
                        "Title": incident.title,
                        "Priority": incident.priority,
                        "Status": incident.status.value,
                        "Created": incident.created_at
                    })
                df = pd.DataFrame(incident_data)
                st.dataframe(df)
                
                # Incident actions
                selected_incident = st.selectbox("Select incident to update", df["ID"])
                new_status = st.selectbox("Update status", 
                                        ["acknowledged", "in_progress", "resolved"])
                if st.button("Update Status"):
                    incident = session.query(Incident).filter_by(id=selected_incident).first()
                    incident.status = IncidentStatus[new_status.upper()]
                    if new_status == "resolved":
                        incident.resolved_at = datetime.utcnow()
                    session.commit()
                    st.success("Status updated successfully!")
                    st.rerun()
    
    with tabs[1]:
        st.header("Notification Preferences")
        with Session(engine) as session:
            user = session.query(User).filter_by(id=st.session_state.user_id).first()
            
            email = st.text_input("Email", value=user.email or "")
            phone = st.text_input("Phone", value=user.phone or "")
            slack_id = st.text_input("Slack ID", value=user.slack_id or "")
            
            channels = st.multiselect(
                "Preferred Notification Channels",
                ["email", "sms", "slack"],
                default=[pref.channel.value for pref in user.preferences]
            )
            
            if st.button("Update Preferences"):
                user.email = email
                user.phone = phone
                user.slack_id = slack_id
                
                # Update preferences
                session.query(UserPreference).filter_by(user_id=user.id).delete()
                for channel in channels:
                    pref = UserPreference(
                        user_id=user.id,
                        channel=NotificationChannel[channel.upper()],
                        is_active=True
                    )
                    session.add(pref)
                session.commit()
                st.success("Preferences updated successfully!")

def show_admin_dashboard():
    st.title("Admin Dashboard")
    
    tabs = st.tabs(["Onboard Resolvers", "On-Call Schedule", "Incident Reports"])
    
    with tabs[0]:
        st.header("Onboard New Resolver")
        new_username = st.text_input("Username")
        new_password = st.text_input("Password", type="password")
        new_email = st.text_input("Email")
        new_phone = st.text_input("Phone")
        
        if st.button("Create Resolver"):
            with Session(engine) as session:
                new_user = User(
                    username=new_username,
                    password=new_password,
                    role=UserRole.RESOLVER,
                    email=new_email,
                    phone=new_phone
                )
                session.add(new_user)
                session.commit()
                st.success("Resolver created successfully!")
    
    with tabs[1]:
        st.header("Manage On-Call Schedule")
        with Session(engine) as session:
            resolvers = session.query(User).filter_by(role=UserRole.RESOLVER).all()
            resolver_options = {user.username: user.id for user in resolvers}
            
            primary = st.selectbox("Primary Resolver", list(resolver_options.keys()))
            secondary = st.selectbox("Secondary Resolver", list(resolver_options.keys()))
            start_date = st.date_input("Start Date")
            end_date = st.date_input("End Date")
            
            if st.button("Create Schedule"):
                schedule = OnCallSchedule(
                    primary_resolver_id=resolver_options[primary],
                    secondary_resolver_id=resolver_options[secondary],
                    start_date=datetime.combine(start_date, datetime.min.time()),
                    end_date=datetime.combine(end_date, datetime.max.time())
                )
                session.add(schedule)
                session.commit()
                st.success("Schedule created successfully!")
    
    with tabs[2]:
        st.header("Incident Reports")
        with Session(engine) as session:
            incidents = session.query(Incident).all()
            if incidents:
                incident_data = []
                for incident in incidents:
                    incident_data.append({
                        "ID": incident.id,
                        "Title": incident.title,
                        "Priority": incident.priority,
                        "Status": incident.status.value,
                        "Created": incident.created_at,
                        "Resolved": incident.resolved_at
                    })
                df = pd.DataFrame(incident_data)
                
                # Status distribution
                fig_status = px.pie(df, names="Status", title="Incidents by Status")
                st.plotly_chart(fig_status)
                
                # Priority distribution
                fig_priority = px.bar(df, x="Priority", title="Incidents by Priority")
                st.plotly_chart(fig_priority)
                
                st.subheader("Raw Data")
                st.dataframe(df)

# Main app logic
if not st.session_state.logged_in:
    show_login()
else:
    if st.session_state.user_role == UserRole.RESOLVER:
        show_resolver_dashboard()
    else:
        show_admin_dashboard()
    
    if st.sidebar.button("Logout"):
        st.session_state.logged_in = False
        st.session_state.user_role = None
        st.session_state.user_id = None
        st.rerun()