#@title Dependency Installation
# Install necessary libraries
%pip install hvac google-cloud-kms loguru pandas matplotlib seaborn

import os
import time
import subprocess

# 1. Download standalone Vault binary if not present
if not os.path.exists('./vault'):
    print("Downloading Vault binary...")
    !wget https://releases.hashicorp.com/vault/1.15.4/vault_1.15.4_linux_amd64.zip
    !unzip -o vault_1.15.4_linux_amd64.zip
    !chmod +x vault

# 2. Start Vault in dev mode in the background
print("Starting Vault dev server...")
!pkill vault || true
time.sleep(2)
get_ipython().system_raw('./vault server -dev -dev-root-token-id="myroottoken" > vault.log 2>&1 &')

# Wait for Vault to be ready
print("Waiting for Vault to be ready...")
for i in range(10):
    try:
        env = os.environ.copy()
        env['VAULT_ADDR'] = 'http://127.0.0.1:8200'
        result = subprocess.run(['./vault', 'status'], env=env, capture_output=True)
        if result.returncode in [0, 2]:
            print("Vault is ready!")
            break
    except:
        pass
    time.sleep(1)

# 3. Configure environment for hvac and CLI
os.environ['VAULT_ADDR'] = 'http://127.0.0.1:8200'
os.environ['VAULT_TOKEN'] = 'myroottoken'

# 4. Enable engines
print("Enabling KV and Transit engines...")
!./vault secrets enable -path=kv kv-v2 || true
!./vault secrets enable transit || true
!./vault write -f transit/keys/my-key || true

print("Vault is configured and running at http://127.0.0.1:8200")

# Enable the Transit secrets engine for local encryption/decryption
print("Enabling Transit secrets engine...")
!./vault secrets enable transit || echo "Transit engine might already be enabled"

# Create a named encryption key
print("Creating local encryption key 'my-key'...")
!./vault write -f transit/keys/my-key

#@title Library Imports
# Standard library imports
import os
import time
import logging
import random
from collections import deque
from typing import Dict, Any, Optional

# Third-party library imports
import hvac # HashiCorp Vault client
from google.cloud import kms_v1 # Google Cloud KMS client
from loguru import logger # Enhanced logging
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

import sys
from loguru import logger
from typing import Dict, Any, Optional

def create_app_state(vault_addr: Optional[str] = None, kms_project_id: Optional[str] = None, log_level: str = "INFO") -> Dict[str, Any]:
    """
    Initializes the application's state dictionary.

    Parameters
    ----------
    vault_addr : Optional[str]
        The address for the HashiCorp Vault server. Defaults to VAULT_ADDR environment variable.
    kms_project_id : Optional[str]
        The Google Cloud Project ID for KMS. Defaults to GCP_PROJECT_ID environment variable.
    log_level : str
        The logging level for the application (e.g., 'INFO', 'DEBUG', 'WARNING').

    Returns
    -------
    Dict[str, Any]
        An initialized state dictionary containing configuration and client placeholders.

    Examples
    --------
    >>> state = create_app_state(vault_addr='http://localhost:8200', log_level='DEBUG')
    >>> assert 'vault_addr' in state
    >>> assert state['log_level'] == 'DEBUG'
    """
    logger.info("Initializing application state.")
    state = {
        "vault_addr": vault_addr if vault_addr else os.getenv("VAULT_ADDR"),
        "vault_token": os.getenv("VAULT_TOKEN"), # Assuming token is passed via env for simplicity in setup
        "kms_project_id": kms_project_id if kms_project_id else os.getenv("GCP_PROJECT_ID"),
        "kms_client": None,
        "vault_client": None,
        "log_level": log_level.upper(),
        "metrics": deque(maxlen=100), # For tracking performance/errors
        "retry_attempts": 0
    }
    logger.remove() # Remove default handler
    logger.add(sys.stderr, level=state["log_level"])
    logger.debug(f"Initial state created: {state}")
    return state

from typing import Dict, Any, Optional
import requests # Import requests for specific exceptions handling

def connect_to_vault(state: Dict[str, Any]) -> Dict[str, Any]:
    """
    Establishes a connection to HashiCorp Vault with retry logic.

    Parameters
    ----------
    state : Dict[str, Any]
        The current application state dictionary, expected to contain 'vault_addr' and 'vault_token'.

    Returns
    -------
    Dict[str, Any]
        The updated state dictionary with the 'vault_client' initialized upon successful connection.
        Returns the state with 'vault_client' as None if connection fails after retries.

    Examples
    --------
    >>> state = create_app_state(vault_addr='http://localhost:8200', vault_token='myroottoken')
    >>> updated_state = connect_to_vault(state)
    >>> assert updated_state['vault_client'] is not None
    """
    vault_addr = state.get("vault_addr")
    vault_token = state.get("vault_token")
    if not vault_addr:
        logger.error("Vault address not found in state or environment variables.")
        return state

    max_retries = 5
    base_delay = 1 # seconds
    for i in range(max_retries):
        try:
            logger.info(f"Attempting to connect to Vault at {vault_addr} (Attempt {i+1}/{max_retries})")
            client = hvac.Client(url=vault_addr, token=vault_token)
            if client.is_authenticated():
                state["vault_client"] = client
                logger.success("Successfully connected and authenticated to Vault.")
                state["retry_attempts"] = 0
                return state
            else:
                logger.warning("Vault client not authenticated. Check token or authentication method.")
        except hvac.exceptions.VaultError as e:
            logger.error(f"Vault connection error: {e}")
        except requests.exceptions.ConnectionError as e:
            logger.error(f"Network connection error to Vault: {e}")
        except Exception as e:
            logger.error(f"An unexpected error occurred during Vault connection: {e}")

        delay = base_delay * (2 ** i) + random.uniform(0, 0.5) # Exponential backoff with jitter
        logger.info(f"Retrying in {delay:.2f} seconds...")
        time.sleep(delay)
        state["retry_attempts"] += 1

    logger.error(f"Failed to connect to Vault after {max_retries} attempts.")
    state["vault_client"] = None # Ensure client is None if connection fails
    return state

from typing import Dict, Any, Optional

def connect_to_kms(state: Dict[str, Any]) -> Dict[str, Any]:
    """
    Initializes a Google Cloud KMS client with retry logic.

    Parameters
    ----------
    state : Dict[str, Any]
        The current application state dictionary, expected to contain 'kms_project_id'.

    Returns
    -------
    Dict[str, Any]
        The updated state dictionary with the 'kms_client' initialized upon successful connection.
        Returns the state with 'kms_client' as None if connection fails after retries.

    Examples
    --------
    >>> state = create_app_state(kms_project_id='my-gcp-project')
    >>> updated_state = connect_to_kms(state)
    >>> assert updated_state['kms_client'] is not None
    """
    kms_project_id = state.get("kms_project_id")
    if not kms_project_id:
        logger.error("GCP Project ID for KMS not found in state or environment variables.")
        return state

    max_retries = 5
    base_delay = 1 # seconds
    for i in range(max_retries):
        try:
            logger.info(f"Attempting to initialize KMS client for project {kms_project_id} (Attempt {i+1}/{max_retries})")
            client = kms_v1.KeyManagementServiceClient()
            # A simple way to verify client functionality without an actual KMS call
            # is to ensure the client object is created. More robust would be a dummy API call.
            state["kms_client"] = client
            logger.success("Successfully initialized Google Cloud KMS client.")
            state["retry_attempts"] = 0
            return state
        except Exception as e:
            logger.error(f"Error initializing Google Cloud KMS client: {e}")

        delay = base_delay * (2 ** i) + random.uniform(0, 0.5) # Exponential backoff with jitter
        logger.info(f"Retrying in {delay:.2f} seconds...")
        time.sleep(delay)
        state["retry_attempts"] += 1

    logger.error(f"Failed to initialize Google Cloud KMS client after {max_retries} attempts.")
    state["kms_client"] = None # Ensure client is None if connection fails
    return state

from typing import Dict, Any, Optional
import requests # Import requests for specific exceptions handling

def store_secret_in_vault(state: Dict[str, Any], path: str, secret_data: Dict[str, Any]) -> bool:
    """
    Stores a secret in HashiCorp Vault's KV secrets engine with retry logic.

    Parameters
    ----------
    state : Dict[str, Any]
        The current application state dictionary, expected to contain 'vault_client'.
    path : str
        The path within the Vault KV secrets engine (e.g., 'kv/data/my-app/db-creds').
    secret_data : Dict[str, Any]
        A dictionary containing the secret key-value pairs to store.

    Returns
    -------
    bool
        True if the secret was stored successfully, False otherwise.

    Examples
    --------
    >>> # Assume 'state' has an authenticated vault_client
    >>> # success = store_secret_in_vault(state, 'kv/data/my-app/test-secret', {'api_key': '12345', 'user': 'testuser'})
    >>> # assert success is True
    """
    vault_client = state.get("vault_client")
    if not vault_client:
        logger.error("Vault client not initialized. Cannot store secret.")
        return False

    max_retries = 3
    base_delay = 0.5 # seconds
    for i in range(max_retries):
        try:
            logger.info(f"Attempting to store secret at '{path}' in Vault (Attempt {i+1}/{max_retries})")
            vault_client.secrets.kv.v2.create_or_update_secret(path=path, secret=secret_data)
            logger.success(f"Secret successfully stored at '{path}' in Vault.")
            state["retry_attempts"] = 0
            return True
        except hvac.exceptions.VaultError as e:
            logger.error(f"Vault error while storing secret at '{path}': {e}")
        except requests.exceptions.ConnectionError as e:
            logger.error(f"Network error while connecting to Vault: {e}")
        except Exception as e:
            logger.error(f"An unexpected error occurred while storing secret: {e}")

        delay = base_delay * (2 ** i) + random.uniform(0, 0.2) # Exponential backoff with jitter
        logger.warning(f"Retrying secret storage in {delay:.2f} seconds...")
        time.sleep(delay)
        state["retry_attempts"] += 1

    logger.error(f"Failed to store secret at '{path}' after {max_retries} attempts.")
    return False

from typing import Dict, Any, Optional

def get_secret_from_vault(state: Dict[str, Any], path: str) -> Optional[Dict[str, Any]]:
    """
    Retrieves a secret from HashiCorp Vault's KV secrets engine with retry logic.

    Parameters
    ----------
    state : Dict[str, Any]
        The current application state dictionary, expected to contain 'vault_client'.
    path : str
        The path within the Vault KV secrets engine where the secret is stored.

    Returns
    -------
    Optional[Dict[str, Any]]
        The secret data if retrieved successfully, otherwise None.

    Examples
    --------
    >>> # Assume 'state' has an authenticated vault_client and a secret at 'kv/data/my-app/test-secret'
    >>> # secret = get_secret_from_vault(state, 'kv/data/my-app/test-secret')
    >>> # if secret: print(secret)
    """
    vault_client = state.get("vault_client")
    if not vault_client:
        logger.error("Vault client not initialized. Cannot retrieve secret.")
        return None

    max_retries = 3
    base_delay = 0.5 # seconds
    for i in range(max_retries):
        try:
            logger.info(f"Attempting to retrieve secret from '{path}' in Vault (Attempt {i+1}/{max_retries})")
            read_response = vault_client.secrets.kv.v2.read_secret_version(path=path)
            secret_data = read_response['data']['data']
            logger.success(f"Secret successfully retrieved from '{path}' in Vault.")
            state["retry_attempts"] = 0
            return secret_data
        except hvac.exceptions.VaultError as e:
            if "No key exists at" in str(e): # Specific check for secret not found
                logger.warning(f"Secret not found at path '{path}' in Vault: {e}")
                return None
            logger.error(f"Vault error while retrieving secret from '{path}': {e}")
        except requests.exceptions.ConnectionError as e:
            logger.error(f"Network error while connecting to Vault: {e}")
        except Exception as e:
            logger.error(f"An unexpected error occurred while retrieving secret: {e}")

        delay = base_delay * (2 ** i) + random.uniform(0, 0.2) # Exponential backoff with jitter
        logger.warning(f"Retrying secret retrieval in {delay:.2f} seconds...")
        time.sleep(delay)
        state["retry_attempts"] += 1

    logger.error(f"Failed to retrieve secret from '{path}' after {max_retries} attempts.")
    return None

from typing import Dict, Any, Optional
import base64
from google.api_core.exceptions import GoogleAPIError

def encrypt_with_kms(state: Dict[str, Any], key_name: str, plaintext: str) -> Optional[str]:
    """
    Encrypts plaintext data using Google Cloud KMS with retry logic.

    Parameters
    ----------
    state : Dict[str, Any]
        The current application state dictionary, expected to contain 'kms_client'.
    key_name : str
        The full resource name of the KMS key.
    plaintext : str
        The data to be encrypted.

    Returns
    -------
    Optional[str]
        The base64-encoded ciphertext if encryption is successful, otherwise None.

    Examples
    --------
    >>> # Assume 'state' has an initialized kms_client
    >>> # key_name = 'projects/my-project/locations/global/keyRings/my-keyring/cryptoKeys/my-key'
    >>> # encrypted_text = encrypt_with_kms(state, key_name, 'sensitive data')
    >>> # if encrypted_text: print(f"Encrypted: {encrypted_text}")
    """
    kms_client = state.get("kms_client")
    if not kms_client:
        logger.error("KMS client not initialized. Cannot encrypt data.")
        return None

    plaintext_bytes = plaintext.encode('utf-8')

    max_retries = 3
    base_delay = 0.5 # seconds
    for i in range(max_retries):
        try:
            logger.info(f"Attempting to encrypt data with KMS key '{key_name}' (Attempt {i+1}/{max_retries})")
            encrypt_response = kms_client.encrypt(request={'name': key_name, 'plaintext': plaintext_bytes})
            ciphertext = base64.b64encode(encrypt_response.ciphertext).decode('utf-8')
            logger.success("Data successfully encrypted with KMS.")
            state["retry_attempts"] = 0
            return ciphertext
        except GoogleAPIError as e:
            logger.error(f"Google Cloud KMS API error during encryption: {e}")
        except Exception as e:
            logger.error(f"An unexpected error occurred during KMS encryption: {e}")

        delay = base_delay * (2 ** i) + random.uniform(0, 0.2) # Exponential backoff with jitter
        logger.warning(f"Retrying KMS encryption in {delay:.2f} seconds...")
        time.sleep(delay)
        state["retry_attempts"] += 1

    logger.error(f"Failed to encrypt data with KMS key '{key_name}' after {max_retries} attempts.")
    return None

from typing import Dict, Any, Optional

def decrypt_with_kms(state: Dict[str, Any], key_name: str, ciphertext: str) -> Optional[str]:
    """
    Decrypts ciphertext data using Google Cloud KMS with retry logic.

    Parameters
    ----------
    state : Dict[str, Any]
        The current application state dictionary, expected to contain 'kms_client'.
    key_name : str
        The full resource name of the KMS key used for decryption.
    ciphertext : str
        The base64-encoded ciphertext to be decrypted.

    Returns
    -------
    Optional[str]
        The decrypted plaintext string if decryption is successful, otherwise None.

    Examples
    --------
    >>> # Assume 'state' has an initialized kms_client and 'encrypted_text'
    >>> # key_name = 'projects/my-project/locations/global/keyRings/my-keyring/cryptoKeys/my-key'
    >>> # decrypted_text = decrypt_with_kms(state, key_name, encrypted_text)
    >>> # if decrypted_text: print(f"Decrypted: {decrypted_text}")
    """
    kms_client = state.get("kms_client")
    if not kms_client:
        logger.error("KMS client not initialized. Cannot decrypt data.")
        return None

    try:
        ciphertext_bytes = base64.b64decode(ciphertext)
    except Exception as e:
        logger.error(f"Error base64 decoding ciphertext: {e}")
        return None

    max_retries = 3
    base_delay = 0.5 # seconds
    for i in range(max_retries):
        try:
            logger.info(f"Attempting to decrypt data with KMS key '{key_name}' (Attempt {i+1}/{max_retries})")
            decrypt_response = kms_client.decrypt(request={'name': key_name, 'ciphertext': ciphertext_bytes})
            plaintext = decrypt_response.plaintext.decode('utf-8')
            logger.success("Data successfully decrypted with KMS.")
            state["retry_attempts"] = 0
            return plaintext
        except GoogleAPIError as e:
            logger.error(f"Google Cloud KMS API error during decryption: {e}")
        except Exception as e:
            logger.error(f"An unexpected error occurred during KMS decryption: {e}")

        delay = base_delay * (2 ** i) + random.uniform(0, 0.2) # Exponential backoff with jitter
        logger.warning(f"Retrying KMS decryption in {delay:.2f} seconds...")
        time.sleep(delay)
        state["retry_attempts"] += 1

    logger.error(f"Failed to decrypt data with KMS key '{key_name}' after {max_retries} attempts.")
    return None

def encrypt_with_vault_transit(state: Dict[str, Any], key_name: str, plaintext: str) -> Optional[str]:
    """Encrypts data using Vault's Transit engine."""
    client = state.get("vault_client")
    if not client:
        return None

    try:
        # Vault expects base64 encoded plaintext for Transit
        import base64
        encoded_plaintext = base64.b64encode(plaintext.encode('utf-8')).decode('utf-8')
        response = client.secrets.transit.encrypt_data(name=key_name, plaintext=encoded_plaintext)
        return response['data']['ciphertext']
    except Exception as e:
        logger.error(f"Vault Transit encryption error: {e}")
        return None

def decrypt_with_vault_transit(state: Dict[str, Any], key_name: str, ciphertext: str) -> Optional[str]:
    """Decrypts data using Vault's Transit engine."""
    client = state.get("vault_client")
    if not client:
        return None

    try:
        response = client.secrets.transit.decrypt_data(name=key_name, ciphertext=ciphertext)
        import base64
        return base64.b64decode(response['data']['plaintext']).decode('utf-8')
    except Exception as e:
        logger.error(f"Vault Transit decryption error: {e}")
        return None

from typing import Dict, Any
import time

def track_metric(state: Dict[str, Any], metric_name: str, value: Any) -> Dict[str, Any]:
    """
    Tracks an operational metric in the application state's metrics deque.

    Parameters
    ----------
    state : Dict[str, Any]
        The current application state dictionary, expected to contain a 'metrics' deque.
    metric_name : str
        The name of the metric to track.
    value : Any
        The value associated with the metric.

    Returns
    -------
    Dict[str, Any]
        The updated state dictionary.

    Examples
    --------
    >>> state = {'metrics': deque(maxlen=10)}
    >>> state = track_metric(state, 'vault_read_latency', 0.05)
    >>> assert len(state['metrics']) == 1
    >>> assert state['metrics'][0]['name'] == 'vault_read_latency'
    """
    timestamp = time.time()
    state['metrics'].append({'timestamp': timestamp, 'name': metric_name, 'value': value})
    logger.debug(f"Metric tracked: {metric_name} = {value}")
    return state

from typing import Dict, Any, Optional
import requests # Import requests for specific exceptions handling

def delete_secret_from_vault(state: Dict[str, Any], path: str) -> bool:
    """
    Deletes a secret from HashiCorp Vault's KV secrets engine with retry logic.

    Parameters
    ----------
    state : Dict[str, Any]
        The current application state dictionary, expected to contain 'vault_client'.
    path : str
        The path within the Vault KV secrets engine where the secret is to be deleted.

    Returns
    -------
    bool
        True if the secret was successfully deleted, False otherwise.

    Examples
    --------
    >>> # Assume 'state' has an authenticated vault_client and a secret at 'kv/data/my-app/temp-secret'
    >>> # success = delete_secret_from_vault(state, 'kv/data/my-app/temp-secret')
    >>> # assert success is True
    """
    vault_client = state.get("vault_client")
    if not vault_client:
        logger.error("Vault client not initialized. Cannot delete secret.")
        return False

    max_retries = 3
    base_delay = 0.5 # seconds
    for i in range(max_retries):
        try:
            logger.info(f"Attempting to delete secret from '{path}' in Vault (Attempt {i+1}/{max_retries})")
            vault_client.secrets.kv.v2.delete_latest_version_of_secret(path=path)
            logger.success(f"Secret successfully deleted from '{path}' in Vault.")
            state["retry_attempts"] = 0
            return True
        except hvac.exceptions.VaultError as e:
            if "Code: 404" in str(e): # Secret not found
                logger.warning(f"Secret at '{path}' not found or already deleted: {e}")
                return True # Consider it successful if it's already gone
            logger.error(f"Vault error while deleting secret from '{path}': {e}")
        except requests.exceptions.ConnectionError as e:
            logger.error(f"Network error while connecting to Vault: {e}")
        except Exception as e:
            logger.error(f"An unexpected error occurred while deleting secret: {e}")

        delay = base_delay * (2 ** i) + random.uniform(0, 0.2) # Exponential backoff with jitter
        logger.warning(f"Retrying secret deletion in {delay:.2f} seconds...")
        time.sleep(delay)
        state["retry_attempts"] += 1

    logger.error(f"Failed to delete secret from '{path}' after {max_retries} attempts.")
    return False

#@title Initialize App State and Connect to Vault/KMS
import os # Added import
from loguru import logger # Added import

# Simulate environment variables (replace with your actual values for a real setup)
os.environ['VAULT_ADDR'] = os.getenv('VAULT_ADDR', 'http://127.0.0.1:8200') # Default to dev server
os.environ['VAULT_TOKEN'] = os.getenv('VAULT_TOKEN', 'myroottoken') # Replace with your dev token or real token
os.environ['GCP_PROJECT_ID'] = os.getenv('GCP_PROJECT_ID', 'your-gcp-project-id') # Replace with your GCP Project ID

# Initialize application state
app_state = create_app_state(log_level='DEBUG')
logger.info(f"Application state initialized with log level: {app_state['log_level']}")

# Connect to Vault
app_state = connect_to_vault(app_state)

# Connect to KMS
app_state = connect_to_kms(app_state)

# Display connection status
print("\n--- Connection Status ---")
print(f"Vault Client Connected: {app_state['vault_client'] is not None}")
print(f"KMS Client Connected: {app_state['kms_client'] is not None}")

#@title Vault Secret Operations
import time # Added import
from loguru import logger # Added import

# Define a test secret
VAULT_TEST_PATH = 'kv/data/my-app/test-db-creds'
TEST_SECRET_DATA = {
    'username': 'app_user',
    'password': 'super_secret_password_123',
    'connection_string': 'jdbc:postgresql://db.example.com:5432/production'
}

# 1. Store a secret
logger.info(f"\n--- Storing secret at {VAULT_TEST_PATH} ---")
start_time = time.perf_counter()
store_success = store_secret_in_vault(app_state, VAULT_TEST_PATH, TEST_SECRET_DATA)
end_time = time.perf_counter()
latency = (end_time - start_time) * 1000 # milliseconds
app_state = track_metric(app_state, 'vault_store_latency_ms', latency)
app_state = track_metric(app_state, 'vault_store_success', int(store_success))

print(f"Secret store successful: {store_success} (Latency: {latency:.2f} ms)")

# 2. Retrieve the secret
logger.info(f"\n--- Retrieving secret from {VAULT_TEST_PATH} ---")
start_time = time.perf_counter()
retrieved_secret = get_secret_from_vault(app_state, VAULT_TEST_PATH)
end_time = time.perf_counter()
latency = (end_time - start_time) * 1000 # milliseconds
app_state = track_metric(app_state, 'vault_retrieve_latency_ms', latency)
app_state = track_metric(app_state, 'vault_retrieve_success', int(retrieved_secret is not None))

if retrieved_secret:
    print("Secret retrieved successfully:")
    for k, v in retrieved_secret.items():
        print(f"  {k}: {'*' * len(v) if k == 'password' else v}") # Mask password
else:
    print("Failed to retrieve secret.")

# 3. Attempt to retrieve a non-existent secret (edge case)
logger.info(f"\n--- Attempting to retrieve non-existent secret ---")
non_existent_secret = get_secret_from_vault(app_state, 'kv/data/my-app/non-existent-creds')
print(f"Non-existent secret retrieval result: {non_existent_secret}")
app_state = track_metric(app_state, 'vault_retrieve_non_existent_success', int(non_existent_secret is not None))

# 4. Delete the secret
logger.info(f"\n--- Deleting secret from {VAULT_TEST_PATH} ---")
start_time = time.perf_counter()
del_success = delete_secret_from_vault(app_state, VAULT_TEST_PATH)
end_time = time.perf_counter()
latency = (end_time - start_time) * 1000 # milliseconds
app_state = track_metric(app_state, 'vault_delete_latency_ms', latency)
app_state = track_metric(app_state, 'vault_delete_success', int(del_success))

print(f"Secret deletion successful: {del_success} (Latency: {latency:.2f} ms)")

# 5. Verify deletion
logger.info(f"\n--- Verifying secret deletion from {VAULT_TEST_PATH} ---")
verified_deleted = get_secret_from_vault(app_state, VAULT_TEST_PATH)
print(f"Verification after deletion: {verified_deleted is None}")

LOCAL_KEY_NAME = 'my-key'
PLAINTEXT = 'This is local sensitive data.'

# Perform local encryption
logger.info("--- Starting Local Transit Encryption ---")
ciphertext = encrypt_with_vault_transit(app_state, LOCAL_KEY_NAME, PLAINTEXT)

if ciphertext:
    print(f"Local Ciphertext: {ciphertext}")

    # Track metrics for the local operation
    app_state = track_metric(app_state, 'local_transit_encrypt_success', 1)

    # Perform local decryption
    decrypted = decrypt_with_vault_transit(app_state, LOCAL_KEY_NAME, ciphertext)

    if decrypted:
        print(f"Decrypted: {decrypted}")
        print(f"Original matches decrypted: {PLAINTEXT == decrypted}")
        app_state = track_metric(app_state, 'local_transit_decrypt_success', 1)
    else:
        print("Local decryption failed.")
else:
    print("Local encryption failed.")

#@title Plotting Metrics
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from IPython.display import display

# Convert metrics deque to a Pandas DataFrame
metrics_df = pd.DataFrame(list(app_state['metrics']))

if not metrics_df.empty:
    metrics_df['timestamp'] = pd.to_datetime(metrics_df['timestamp'], unit='s')

    print("\n--- Summary of All Tracked Metrics ---")
    display(metrics_df)

    # Filter for success/failure metrics including the new local transit metrics
    success_metrics = metrics_df[metrics_df['name'].str.contains('_success')]
    if not success_metrics.empty:
        plt.figure(figsize=(10, 5))
        # Updated to resolve FutureWarning by assigning x to hue and setting legend=False
        sns.barplot(data=success_metrics, x='name', y='value', hue='name', legend=False, palette='viridis')
        plt.title('Vault Operation Success Indicators')
        plt.xlabel('Operation Type')
        plt.ylabel('Success (1=True, 0=False)')
        plt.xticks(rotation=45, ha='right')
        plt.ylim(0, 1.1)
        plt.grid(axis='y', linestyle='--', alpha=0.7)
        plt.tight_layout()
        plt.show()
    else:
        print("No success metrics to plot.")

    # Filter for latency metrics
    latency_metrics = metrics_df[metrics_df['name'].str.contains('_latency_ms')]
    if not latency_metrics.empty:
        plt.figure(figsize=(10, 5))
        sns.lineplot(data=latency_metrics, x='timestamp', y='value', hue='name', marker='o')
        plt.title('Vault Operation Latencies')
        plt.xlabel('Time')
        plt.ylabel('Latency (ms)')
        plt.grid(True)
        plt.tight_layout()
        plt.show()
else:
    print("No metrics have been tracked yet.")