import streamlit as st
import # your database manager here!
from cryptography.fernet import Fernet
from PIL import Image, PngImagePlugin
import base64, hashlib, uuid

def initialize_session_state():
   pass

class BadgeConfig:
        def __init__(self):
          initialize_session_state()
          # Ensure badge_id is stored as a class attribute
          self.badge_id = None

        # Generate a SHA-256 hash from image data
        def generate_image_hash(self, image_data):
            return hashlib.sha256(image_data).hexdigest()

        # Create the encryption signature using Fernet
        def create_signature(self, unique_id):
            id_image_bytes = unique_id.encode("utf-8")
            # Ensure 32-byte length
            padded_id_image_bytes = id_image_bytes.ljust(32)[:32] 
            encoded_key = base64.urlsafe_b64encode(padded_id_image_bytes)
            return Fernet(encoded_key)

        # Check if image was created on Canva using PngImagePlugin
        def is_canva_image(self, image):
            if isinstance(image, PngImagePlugin.PngImageFile):
                # Extract metadata from the image
                metadata = image.info  # Info contains the metadata

                # Checking if 'Canva' appears in the 'xmp:CreatorTool' field
                xmp_metadata = metadata.get('XML:com.adobe.xmp', '')
                if "Canva" in xmp_metadata:
                    return True
                return False

        # Display the uploaded badge and validate its dimensions and source
        def process_image(self, user_badge):
            try:
                image = Image.open(user_badge)
                WIDTH, HEIGHT = image.size
                if WIDTH != 1080 or HEIGHT != 1920:
                    st.warning("This is not a valid dnakey-badge!")
                    st.stop()

                # Check if the image is created on Canva
                if not self.is_canva_image(image):
                    st.warning("The uploaded image is not a Canva PNG image!")
                    st.stop()

                st.image(user_badge, caption="Uploaded Image", use_column_width=True)

                # Reset the file pointer and read the image data for hashing
                user_badge.seek(0)
                return user_badge.read()
            except Exception as e:
                st.error(f"Error processing the image: {str(e)}")
                st.stop()

        # Handle badge activation and update session
        def activate_badge(self, badge_usage, config_manager):
            if not st.session_state['toast_shown']:
                st.toast("**:blue[Your Id Badge is activated now!]**", icon="🧁")
                st.session_state['toast_shown'] = True

            if not st.session_state['usage_updated'] and badge_usage > 0:
                config_manager.update_usage_badge_count()
                st.session_state['usage_updated'] = True

        # Main function to create a session and handle badge logic
        def create_session(self, user_badge):
            # Process image and generate its unique ID
            image_data = self.process_image(user_badge)
            unique_id = self.generate_image_hash(image_data)

            # Create an encryption signature
            signature = self.create_signature(unique_id)

            # Create a UUID (version 5) based on the existing unique_id
            self.badge_id = str(uuid.uuid5(uuid.NAMESPACE_DNS, unique_id))

            # Initialize config manager
            config_manager = ConfigManager(self.badge_id)
            badge_usage = config_manager.get_badge_usage()

            # Handle badge activation and session updates
            self.activate_badge(badge_usage, config_manager)

            return signature, self.badge_id  # Return the badge_id as well

# call the script
with st.sidebar:
     st.title("Log-In Here:")
     with st.popover("Upload Your Agent Badge!", use_container_width=True):
     user_badge = st.file_uploader("Your Agent Badge!", type=["png"], key="agent_badge")

     if user_badge:
        # User badge is uploaded
        signature, badge_id = BadgeConfig().create_session(user_badge)