Update app.py

app.py

@@ -1,22 +1,19 @@
 import streamlit as st
-import uuid
 import cv2
-import numpy as np
 from PIL import Image
+import io
 import tempfile
-import pickle
 import os
+import pickle
+import requests
 from datetime import datetime
 from googleapiclient.discovery import build
 from googleapiclient.http import MediaFileUpload
 from google.auth.transport.requests import Request
 from simple_salesforce import Salesforce
+import pytz
 
-# === Google Drive Setup ===
-REG_FOLDER_ID = '1qkcR7nQTEtiMH9OFUv2bGxVn08E3dKjF'  # Registered images folder
-INTRUDER_FOLDER_ID = '1PPAUWU-wMx7fek73p-hqPqYQypYtG8Ob'  # Intruder folder
-
-# === Salesforce Setup ===
+# === SALESFORCE SETUP ===
 sf = Salesforce(
     username='suriraja822@agentforce.com',
     password='Sati@1010',
@@ -24,112 +21,137 @@ sf = Salesforce(
     domain='login'
 )
 
-def upload_to_drive(image, filename, folder_id):
+# === CONSTANTS ===
+UPLOAD_FOLDER_ID = '1qkcR7nQTEtiMH9OFUv2bGxVn08E3dKjF'
+
+def get_vacant_rooms():
+    rooms = sf.query("SELECT Id, Name FROM Hotel_Room__c WHERE Room_Status__c = 'Vacant'")
+    return {room['Name']: room['Id'] for room in rooms['records']}
+
+def upload_image_to_drive(image_np, filename):
     creds = None
     if os.path.exists("token.pickle"):
         with open("token.pickle", "rb") as token:
             creds = pickle.load(token)
-    if creds and creds.expired and creds.refresh_token:
-        creds.refresh(Request())
+    if not creds or not creds.valid:
+        if creds and creds.expired and creds.refresh_token:
+            creds.refresh(Request())
     service = build('drive', 'v3', credentials=creds)
-
-    image_pil = Image.fromarray(image)
-    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".jpg")
-    image_pil.save(temp_file.name)
-
-    file_metadata = {'name': filename, 'parents': [folder_id]}
-    media = MediaFileUpload(temp_file.name, mimetype='image/jpeg')
-    file = service.files().create(body=file_metadata, media_body=media, fields='id').execute()
-
-    file_id = file.get('id')
-    public_link = f"https://drive.google.com/uc?id={file_id}"
-    os.remove(temp_file.name)
-
-    # Make file public
+    image = Image.fromarray(cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB))
+    buf = io.BytesIO()
+    image.save(buf, format="JPEG")
+    buf.seek(0)
+
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as tmp:
+        tmp.write(buf.read())
+        tmp.flush()
+        file_metadata = {'name': filename, 'parents': [UPLOAD_FOLDER_ID]}
+        media = MediaFileUpload(tmp.name, mimetype='image/jpeg')
+        uploaded_file = service.files().create(body=file_metadata, media_body=media, fields='id').execute()
+        file_id = uploaded_file.get('id')
+        os.remove(tmp.name)
+        return f"https://drive.google.com/uc?id={file_id}"
+
+def save_guest_to_sf(data):
     try:
-        service.permissions().create(
-            fileId=file_id,
-            body={'role': 'reader', 'type': 'anyone'}
-        ).execute()
+        sf.Hotel_Guest__c.create(data)
+        st.success("✅ Guest registered successfully.")
     except Exception as e:
-        st.warning(f"Failed to set permission: {e}")
+        st.error(f"❌ Salesforce Guest Error: {e}")
 
-    return public_link
-
-# === Streamlit UI ===
-st.set_page_config(page_title="Hotel Registration", layout="wide")
-st.title("📸 Hotel Guest & Staff Registration")
+def save_staff_to_sf(data):
+    try:
+        sf.Hotel_Staff__c.create(data)
+        st.success("✅ Staff registered successfully.")
+    except Exception as e:
+        st.error(f"❌ Salesforce Staff Error: {e}")
 
-tab1, tab2 = st.tabs(["🛎️ Guest Registration", "👮 Staff Registration"])
+st.title("🏨 Hotel Guest and Staff Registration")
+tab1, tab2 = st.tabs(["👤 Guest", "👨‍💼 Staff"])
 
-# === Guest Tab ===
 with tab1:
-    st.subheader("Guest Registration")
-
-    guest_name = st.text_input("Guest Name")
-    guest_aadhar = st.text_input("Aadhar Number")
-    guest_phone = st.text_input("Phone Number")
-    guest_room = st.text_input("Room Number")
-
-    id_photo = st.file_uploader("Upload ID Card Photo", type=["jpg", "jpeg", "png"])
-    guest_image = st.camera_input("Capture Guest Photo")
-
-    if st.button("✅ Register Guest"):
-        if guest_image and id_photo and all([guest_name, guest_aadhar, guest_phone, guest_room]):
-            guest_np = np.array(Image.open(guest_image).convert('RGB'))
-
-            filename = f"G_{datetime.now().strftime('%Y%m%d_%H%M%S')}.jpg"
-            try:
-                image_link = upload_to_drive(guest_np, filename, REG_FOLDER_ID)
-                st.image(guest_np, caption="Guest Registered")
-
-                result = sf.Guest__c.create({
-                    "Name": guest_name,
-                    "Aadhar_Number__c": guest_aadhar,
-                    "Phone_Number__c": guest_phone,
-                    "Room_Number__c": guest_room,
-                    "Check_In__c": datetime.utcnow().isoformat(),
-                    "Image_Link__c": image_link
-                })
-
-                st.success("🎉 Guest Registered in Salesforce!")
-                st.json(result)
-            except Exception as e:
-                st.error(f"Salesforce Error: {e}")
+    st.header("Register a Guest")
+    fname = st.text_input("First Name")
+    lname = st.text_input("Last Name")
+    email = st.text_input("Email")
+    phone = st.text_input("Phone")
+    id_number = st.text_input("ID Number")
+    special_requests = st.text_area("Special Requests")
+    emergency_contact = st.text_input("Emergency Contact")
+    vacant_rooms = get_vacant_rooms()
+    room_name = st.selectbox("Assign Room", list(vacant_rooms.keys()))
+    capture = st.camera_input("Capture Guest Photo")
+
+    if st.button("📝 Register Guest"):
+        if capture:
+            timestamp = datetime.now(pytz.timezone('Asia/Kolkata')).strftime("%Y%m%d%H%M%S")
+            image = Image.open(capture)
+            image_np = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+            file_name = f"G_{timestamp}.jpg"
+            photo_url = upload_image_to_drive(image_np, file_name)
+
+            guest_data = {
+                'First_Name__c': fname,
+                'Last_Name__c': lname,
+                'Email__c': email,
+                'Phone__c': phone,
+                'ID_Number__c': id_number,
+                'Special_Requests__c': special_requests,
+                'Emergency_Contact__c': emergency_contact,
+                'Room_Number__c': vacant_rooms[room_name],
+                'Check_In_Date__c': datetime.utcnow().isoformat(),
+                'Photo_Uploaded__c': photo_url,
+                'HF_Registered__c': True,
+                'Guest_Status__c': 'Checked In',
+                'Name': f"{fname} {lname}"
+            }
+            save_guest_to_sf(guest_data)
         else:
-            st.warning("📋 Please fill all fields and upload both photos.")
+            st.warning("📸 Please capture a photo before registering.")
 
-# === Staff Tab ===
 with tab2:
-    st.subheader("Staff Registration")
-
-    staff_name = st.text_input("Staff Name")
-    staff_id = st.text_input("Staff ID")
-    staff_dept = st.selectbox("Department", ["Security", "Reception", "Housekeeping", "Maintenance", "Kitchen"])
-
-    id_card = st.file_uploader("Upload ID Card (Staff)", type=["jpg", "jpeg", "png"], key="staff")
-    staff_image = st.camera_input("Capture Staff Photo")
-
-    if st.button("✅ Register Staff"):
-        if staff_image and id_card and all([staff_name, staff_id, staff_dept]):
-            staff_np = np.array(Image.open(staff_image).convert('RGB'))
-
-            filename = f"S_{datetime.now().strftime('%Y%m%d_%H%M%S')}.jpg"
-            try:
-                image_link = upload_to_drive(staff_np, filename, REG_FOLDER_ID)
-                st.image(staff_np, caption="Staff Registered")
-
-                result = sf.Staff__c.create({
-                    "Name": staff_name,
-                    "Staff_ID__c": staff_id,
-                    "Department__c": staff_dept,
-                    "Registered_On__c": datetime.utcnow().isoformat(),
-                    "Image_Link__c": image_link
-                })
-
-                st.success("✅ Staff Registered in Salesforce!")
-                st.json(result)
-            except Exception as e:
-                st.error(f"Salesforce Error: {e}")
+    st.header("Register a Staff Member")
+    fname = st.text_input("Staff First Name")
+    lname = st.text_input("Staff Last Name")
+    email = st.text_input("Staff Email")
+    phone = st.text_input("Staff Phone")
+    emp_number = st.text_input("Employee Number")
+    department = st.text_input("Department")
+    position = st.text_input("Position")
+    shift = st.selectbox("Shift", ["Morning", "Evening", "Night"])
+    status = st.selectbox("Status", ["Active", "Inactive"])
+    access = st.selectbox("Access Level", ["Low", "Medium", "High"])
+    badge = st.text_input("Badge Number")
+    supervisor = st.text_input("Supervisor")
+    hire_date = st.date_input("Hire Date")
+    capture = st.camera_input("Capture Staff Photo")
+
+    if st.button("📝 Register Staff"):
+        if capture:
+            timestamp = datetime.now(pytz.timezone('Asia/Kolkata')).strftime("%Y%m%d%H%M%S")
+            image = Image.open(capture)
+            image_np = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
+            file_name = f"S_{timestamp}.jpg"
+            photo_url = upload_image_to_drive(image_np, file_name)
+
+            staff_data = {
+                'First_Name__c': fname,
+                'Last_Name__c': lname,
+                'Email__c': email,
+                'Phone__c': phone,
+                'Employee_Number__c': emp_number,
+                'Department__c': department,
+                'Position__c': position,
+                'Shift__c': shift,
+                'Employment_Status__c': status,
+                'Access_Level__c': access,
+                'Badge_Number__c': badge,
+                'Supervisor__c': supervisor,
+                'Hire_Date__c': hire_date.isoformat(),
+                'Photo_Uploaded__c': photo_url,
+                'HF_Registered__c': True,
+                'Name': f"{fname} {lname}"
+            }
+            save_staff_to_sf(staff_data)
         else:
-            st.warning("📋 Please fill all fields and upload both photos.")
+            st.warning("📸 Please capture a photo before registering.")