# Install necessary libraries
!pip install --quiet google-cloud-logging tenacity pandas numpy matplotlib seaborn
# Added a comment to trigger re-execution and ensure imports are refreshed.

import os
import time
import random
import logging
from collections import deque
from datetime import datetime, timedelta

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from tenacity import retry, wait_exponential, stop_after_attempt, after_log, wait_random_exponential
from google.cloud import logging as cloud_logging # Mock for local development

# Configure basic logging for the notebook
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def create_infrastructure_state(project_id: str, region: str, log_level: str = 'INFO') -> dict:
    """
    Initializes the infrastructure state dictionary.

    Parameters
    ----------
    project_id : str
        The identifier for the cloud project.
    region : str
        The geographical region for resource deployment.
    log_level : str, optional
        The desired logging level (e.g., 'INFO', 'DEBUG'). Defaults to 'INFO'.

    Returns
    -------
    dict
        An initial infrastructure state dictionary.
    """
    state = {
        'config': {
            'project_id': project_id,
            'region': region,
            'log_level': log_level,
            'api_keys': {},
            'service_endpoints': {}
        },
        'resources': {
            'vms': {},
            'databases': {},
            'networks': {}
        },
        'monitoring': {
            'resource_health': {},
            'events_log': deque(maxlen=100) # Rolling window for recent events
        },
        'metrics': {
            'provisioning_times': [],
            'deallocation_times': [],
            'errors': []
        }
    }
    logger.info(f"Infrastructure state initialized for project '{project_id}' in region '{region}'.")
    return state

def setup_cloud_logger(state: dict, log_name: str) -> dict:
    """
    Simulates setting up a cloud-based logging client.

    Parameters
    ----------
    state : dict
        The current infrastructure state dictionary.
    log_name : str
        The name of the log to write to.

    Returns
    -------
    dict
        The updated state dictionary with the logging client.
    """
    try:
        # Mock client for demonstration. In a real scenario, this would be an actual cloud client.
        # client = cloud_logging.Client(project=state['config']['project_id'])
        # logger = client.logger(log_name)
        # For this example, we'll just use a placeholder to simulate cloud logger existence.
        state['config']['cloud_logger_client'] = f"MockCloudLoggerClient-{log_name}"
        state['config']['cloud_logger_name'] = log_name
        logger.info(f"Cloud logger '{log_name}' simulated setup successfully.")
        state['monitoring']['events_log'].append({'timestamp': datetime.now(), 'event': f"Cloud logger '{log_name}' initialized", 'level': 'INFO'})
    except Exception as e:
        logger.error(f"Failed to set up cloud logger: {e}")
        state['metrics']['errors'].append({'timestamp': datetime.now(), 'error': str(e), 'function': 'setup_cloud_logger'})
    return state

@retry(wait=wait_random_exponential(multiplier=1, min=1, max=10), stop=stop_after_attempt(5), after=after_log(logger, logging.WARNING))
def _provision_resource_with_retry(state: dict, resource_type: str, resource_name: str, specs: dict) -> dict:
    # Simulate a network or API failure 20% of the time
    if random.random() < 0.2:
        raise ConnectionError(f"Simulated transient failure during {resource_type} provisioning for {resource_name}")

    provision_start_time = time.time()
    time.sleep(1 + random.uniform(0.1, 0.5)) # Simulate provisioning time with jitter

    resource_id = f"{resource_type}-{resource_name}-{random.randint(1000, 9999)}"
    resource_info = {
        'id': resource_id,
        'type': resource_type,
        'name': resource_name,
        'status': 'RUNNING',
        'created_at': datetime.now().isoformat(),
        'specs': specs
    }
    if resource_type not in state['resources']:
        state['resources'][resource_type] = {}
    state['resources'][resource_type][resource_id] = resource_info

    provision_end_time = time.time()
    state['metrics']['provisioning_times'].append(provision_end_time - provision_start_time)
    logger.info(f"Successfully provisioned {resource_type} '{resource_name}' with ID '{resource_id}'.")
    state['monitoring']['events_log'].append({'timestamp': datetime.now(), 'event': f"Resource '{resource_name}' ({resource_type}) provisioned", 'level': 'INFO', 'resource_id': resource_id})
    return state

def simulate_resource_provisioning(state: dict, resource_type: str, resource_name: str, specs: dict = None) -> dict:
    """
    Simulates the provisioning of a cloud resource with retry logic.

    Parameters
    ----------
    state : dict
        The current infrastructure state dictionary.
    resource_type : str
        The type of resource to provision (e.g., 'vm', 'database').
    resource_name : str
        A unique name for the resource.
    specs : dict, optional
        Specifications for the resource (e.g., 'machine_type', 'disk_size'). Defaults to an empty dictionary.

    Returns
    -------
    dict
        The updated state dictionary with the newly provisioned resource.
    """
    if specs is None:
        specs = {}
    try:
        state = _provision_resource_with_retry(state, resource_type, resource_name, specs)
    except Exception as e:
        logger.error(f"Failed to provision {resource_type} '{resource_name}' after multiple retries: {e}")
        state['metrics']['errors'].append({'timestamp': datetime.now(), 'error': str(e), 'function': 'simulate_resource_provisioning', 'resource_name': resource_name})
    return state

def monitor_resource_health(state: dict, resource_id: str) -> dict:
    """
    Simulates monitoring the health status of a specific resource.

    Parameters
    ----------
    state : dict
        The current infrastructure state dictionary.
    resource_id : str
        The ID of the resource to monitor.

    Returns
    -------
    dict
        The updated state dictionary with the latest health status of the resource.
    """
    health_statuses = ['HEALTHY', 'HEALTHY', 'HEALTHY', 'DEGRADED', 'UNHEALTHY']
    simulated_health = random.choice(health_statuses)

    if resource_id not in state['monitoring']['resource_health']:
        state['monitoring']['resource_health'][resource_id] = deque(maxlen=5) # Rolling window for health history

    state['monitoring']['resource_health'][resource_id].append({'timestamp': datetime.now(), 'status': simulated_health})

    if simulated_health != 'HEALTHY':
        logger.warning(f"Resource '{resource_id}' is reported as {simulated_health}.")
        state['monitoring']['events_log'].append({'timestamp': datetime.now(), 'event': f"Resource '{resource_id}' is {simulated_health}", 'level': 'WARNING', 'resource_id': resource_id})
    else:
        logger.debug(f"Resource '{resource_id}' is healthy.")

    return state

def summarize_resource_health(state: dict, status_filter: str = None) -> pd.DataFrame:
    """
    Generates a summary of the health status for all monitored resources.

    Parameters
    ----------
    state : dict
        The current infrastructure state dictionary.
    status_filter : str, optional
        Filter by a specific health status (e.g., 'UNHEALTHY'). Defaults to None.

    Returns
    -------
    pd.DataFrame
        A DataFrame summarizing the health status of resources.
    """
    summary_data = []
    for resource_id, health_history in state['monitoring']['resource_health'].items():
        if health_history:
            latest_entry = health_history[-1]
            resource_type = resource_id.split('-')[0] if '-' in resource_id else 'unknown'
            resource_name = '-'.join(resource_id.split('-')[1:-1]) if '-' in resource_id else 'unknown'
            summary_data.append({
                'resource_id': resource_id,
                'resource_type': resource_type,
                'resource_name': resource_name,
                'latest_status': latest_entry['status'],
                'last_checked': latest_entry['timestamp']
            })

    df_summary = pd.DataFrame(summary_data)

    if status_filter:
        df_summary = df_summary[df_summary['latest_status'] == status_filter]

    if not df_summary.empty:
        logger.info(f"Resource health summary generated. Filter: {status_filter if status_filter else 'None'}")
    else:
        logger.info(f"No resources found matching health summary criteria. Filter: {status_filter if status_filter else 'None'}")

    return df_summary

@retry(wait=wait_random_exponential(multiplier=1, min=1, max=10), stop=stop_after_attempt(5), after=after_log(logger, logging.WARNING))
def _deallocate_resource_with_retry(state: dict, resource_id: str) -> dict:
    # Simulate a network or API failure 15% of the time
    if random.random() < 0.15:
        raise ConnectionError(f"Simulated transient failure during deallocation for {resource_id}")

    deallocate_start_time = time.time()
    time.sleep(0.5 + random.uniform(0.1, 0.3)) # Simulate deallocation time with jitter

    resource_found = False
    for r_type in state['resources']:
        if resource_id in state['resources'][r_type]:
            resource_info = state['resources'][r_type].pop(resource_id) # Remove resource
            resource_found = True
            break

    if not resource_found:
        logger.warning(f"Resource '{resource_id}' not found for deallocation.")
        state['monitoring']['events_log'].append({'timestamp': datetime.now(), 'event': f"Attempted to deallocate non-existent resource '{resource_id}'", 'level': 'WARNING'})
        return state

    deallocate_end_time = time.time()
    state['metrics']['deallocation_times'].append(deallocate_end_time - deallocate_start_time)
    logger.info(f"Successfully deallocated resource '{resource_id}'.")
    state['monitoring']['events_log'].append({'timestamp': datetime.now(), 'event': f"Resource '{resource_id}' deallocated", 'level': 'INFO', 'resource_id': resource_id})
    return state

def simulate_resource_deallocation(state: dict, resource_id: str) -> dict:
    """
    Simulates the deallocation of a previously provisioned cloud resource with retry logic.

    Parameters
    ----------
    state : dict
        The current infrastructure state dictionary.
    resource_id : str
        The ID of the resource to deallocate.

    Returns
    -------
    dict
        The updated state dictionary with the resource removed or marked as deallocated.
    """
    try:
        state = _deallocate_resource_with_retry(state, resource_id)
    except Exception as e:
        logger.error(f"Failed to deallocate resource '{resource_id}' after multiple retries: {e}")
        state['metrics']['errors'].append({'timestamp': datetime.now(), 'error': str(e), 'function': 'simulate_resource_deallocation', 'resource_id': resource_id})
    return state

# Initialize the state
initial_state = create_infrastructure_state(project_id='trading-bot-prod', region='us-central1', log_level='DEBUG')

# Setup mock cloud logger
initial_state = setup_cloud_logger(initial_state, log_name='trading-bot-logs')

print("Current Infrastructure State (after initialization and logger setup):")
# Use json.dumps for pretty printing if the dictionary gets complex, but for now direct print is fine
import json
print(json.dumps(initial_state['config'], indent=2))
print(f"Number of events logged: {len(initial_state['monitoring']['events_log'])}")

current_state = initial_state.copy()

# Provision multiple VMs
vm_specs_1 = {'machine_type': 'e2-standard-2', 'disk_size_gb': 50, 'os_image': 'debian-cloud'}
current_state = simulate_resource_provisioning(current_state, resource_type='vm', resource_name='bot-executor-01', specs=vm_specs_1)

vm_specs_2 = {'machine_type': 'e2-medium', 'disk_size_gb': 30, 'os_image': 'ubuntu-pro'}
current_state = simulate_resource_provisioning(current_state, resource_type='vm', resource_name='data-feeder-01', specs=vm_specs_2)

# Provision a database
db_specs = {'db_type': 'PostgreSQL', 'version': '13', 'tier': 'db-f1-micro', 'storage_gb': 20}
current_state = simulate_resource_provisioning(current_state, resource_type='database', resource_name='trade-db-prod', specs=db_specs)

# Provision another VM, potentially encountering an error
vm_specs_3 = {'machine_type': 'n1-standard-1', 'disk_size_gb': 40, 'os_image': 'centos'}
current_state = simulate_resource_provisioning(current_state, resource_type='vm', resource_name='backtest-worker-01', specs=vm_specs_3)

print("\nResources after provisioning:")
for r_type, resources in current_state['resources'].items():
    print(f"  {r_type.upper()}:")
    for res_id, res_info in resources.items():
        print(f"    - {res_id} (Status: {res_info['status']})")

print(f"\nTotal provisioning times recorded: {len(current_state['metrics']['provisioning_times'])}.")
print(f"Total errors recorded during provisioning: {len(current_state['metrics']['errors'])}")

provisioning_times = current_state['metrics']['provisioning_times']

if provisioning_times:
    plt.figure(figsize=(10, 6))
    sns.histplot(provisioning_times, kde=True, bins=5)
    plt.title('Distribution of Simulated Resource Provisioning Times')
    plt.xlabel('Time (seconds)')
    plt.ylabel('Frequency')
    plt.grid(axis='y', linestyle='--', alpha=0.7)
    plt.axvline(np.mean(provisioning_times), color='r', linestyle='--', label=f'Mean: {np.mean(provisioning_times):.2f}s')
    plt.legend()
    plt.show()
else:
    print("No provisioning times recorded to visualize.")

print("\n--- Monitoring Resource Health ---")
all_resource_ids = []
for r_type, resources in current_state['resources'].items():
    all_resource_ids.extend(resources.keys())

# Simulate monitoring multiple times
for i in range(5):
    print(f"\nMonitoring round {i+1}:")
    for res_id in all_resource_ids:
        current_state = monitor_resource_health(current_state, res_id)
    time.sleep(0.1) # Small delay between monitoring rounds

# Summarize current health status
health_summary_df = summarize_resource_health(current_state)
print("\nComprehensive Health Summary:")
display(health_summary_df)

# Check for unhealthy resources (edge case testing)
unhealthy_resources_df = summarize_resource_health(current_state, status_filter='UNHEALTHY')
print("\nUnhealthy Resources (if any):")
display(unhealthy_resources_df)

print(f"Total events in rolling log: {len(current_state['monitoring']['events_log'])}")

events_df = pd.DataFrame(list(current_state['monitoring']['events_log']))

if not events_df.empty:
    events_df['timestamp'] = pd.to_datetime(events_df['timestamp'])
    events_df = events_df.sort_values('timestamp')

    plt.figure(figsize=(15, 7))
    # Create a scatter plot for events, colored by level
    sns.scatterplot(data=events_df, x='timestamp', y='event', hue='level', style='level', s=100)
    plt.title('Simulated Infrastructure Events Over Time')
    plt.xlabel('Time')
    plt.ylabel('Event Description')
    plt.xticks(rotation=45, ha='right')
    plt.tight_layout()
    plt.grid(axis='x', linestyle='--', alpha=0.5)
    plt.legend(title='Event Level')
    plt.show()
else:
    print("No events in the log to visualize.")

print("\n--- Deallocating Resources ---")

resources_to_deallocate = []
for r_type, resources in list(current_state['resources'].items()): # Iterate over a copy to allow modification
    resources_to_deallocate.extend(resources.keys())

for res_id in resources_to_deallocate:
    current_state = simulate_resource_deallocation(current_state, res_id)
    time.sleep(0.05) # Small delay between deallocations

# Attempt to deallocate a non-existent resource (edge case)
current_state = simulate_resource_deallocation(current_state, 'non-existent-resource-123')

print("\nResources remaining after deallocation:")
found_remaining = False
for r_type, resources in current_state['resources'].items():
    if resources:
        found_remaining = True
        print(f"  {r_type.upper()}: {list(resources.keys())}")
if not found_remaining:
    print("  All resources successfully deallocated.")

print(f"\nTotal deallocation times recorded: {len(current_state['metrics']['deallocation_times'])}")

deallocation_times = current_state['metrics']['deallocation_times']
errors = current_state['metrics']['errors']

plt.figure(figsize=(14, 6))

# Subplot 1: Deallocation Times
plt.subplot(1, 2, 1)
if deallocation_times:
    sns.histplot(deallocation_times, kde=True, bins=3)
    plt.title('Distribution of Simulated Resource Deallocation Times')
    plt.xlabel('Time (seconds)')
    plt.ylabel('Frequency')
    plt.grid(axis='y', linestyle='--', alpha=0.7)
    plt.axvline(np.mean(deallocation_times), color='r', linestyle='--', label=f'Mean: {np.mean(deallocation_times):.2f}s')
    plt.legend()
else:
    plt.text(0.5, 0.5, 'No deallocation times recorded.', horizontalalignment='center', verticalalignment='center', transform=plt.gca().transAxes)
    plt.title('Simulated Resource Deallocation Times')
    plt.xlabel('Time (seconds)')
    plt.ylabel('Frequency')

# Subplot 2: Errors Summary
plt.subplot(1, 2, 2)
if errors:
    errors_df = pd.DataFrame(errors)
    error_counts = errors_df['function'].value_counts()
    sns.barplot(x=error_counts.index, y=error_counts.values, palette='viridis')
    plt.title('Errors Encountered by Function')
    plt.xlabel('Function')
    plt.ylabel('Error Count')
    plt.xticks(rotation=45, ha='right')
    plt.grid(axis='y', linestyle='--', alpha=0.7)
else:
    plt.text(0.5, 0.5, 'No errors recorded.', horizontalalignment='center', verticalalignment='center', transform=plt.gca().transAxes)
    plt.title('Errors Encountered by Function')
    plt.xlabel('Function')
    plt.ylabel('Error Count')

plt.tight_layout()
plt.show()

print("\nSummary of all recorded errors:")
display(pd.DataFrame(errors))