import logging
from azure.storage.blob import BlobServiceClient, generate_blob_sas, BlobSasPermissions
from azure.data.tables import TableServiceClient, generate_table_sas, TableSasPermissions
from datetime import datetime, timezone, timedelta

class AzureStorageConnector:
    def __init__(self, connection_string=None, account_name=None, account_key=None):
        self.connection_string = connection_string
        self.account_name = account_name
        self.account_key = account_key
        self.logger = logging.getLogger('AzureStorageConnector')

    def generate_blob_sas_key(self, container_name, permissions=BlobSasPermissions(read=True), expiry=None):
        """
        Generate a SAS token for Azure Blob Storage.
        """
        expiry = expiry or datetime.now(timezone.utc) + timedelta(hours=1)  
        sas_token = generate_blob_sas(
            account_name=self.account_name,
            container_name=container_name,
            account_key=self.account_key,
            permission=permissions,
            expiry=expiry
        )
        return sas_token

    def generate_table_sas_key(self, table_name, permissions=TableSasPermissions(read=True), expiry=None):
        """
        Generate a SAS token for Azure Table Storage.
        """
        expiry = expiry or datetime.now(timezone.utc) + timedelta(hours=1)  
        sas_token = generate_table_sas(
            account_name=self.account_name,
            table_name=table_name,
            account_key=self.account_key,
            permission=permissions,
            expiry=expiry
        )
        return sas_token

    def connect_blob_service(self, container_name, blob_permissions=BlobSasPermissions(read=True), blob_expiry=None):
        """
        Establish a connection to the Azure Blob Service.
        """
        try:
            if self.connection_string:
                self.blob_service_client = BlobServiceClient.from_connection_string(self.connection_string)
                self.logger.info("Connected to Azure Blob Service using connection string")
            elif self.account_name and self.account_key:
                sas_key = self.generate_blob_sas_key(container_name, permissions=blob_permissions, expiry=blob_expiry)
                self.blob_service_client =  BlobServiceClient(sas_key)

                self.logger.info("Connected to Azure Blob Service using SAS Key")
            else:
                raise ValueError("Either connection string or both account name and account key must be provided")
        except Exception as e:
            self.logger.error(f"Error connecting to Azure Blob Service: {e}")

    def connect_table_service(self, table_name, table_permissions=TableSasPermissions(read=True), table_expiry=None):
        """
        Establish a connection to the Azure Table Service.
        """
        try:
            if self.connection_string:
                self.table_service_client = TableServiceClient.from_connection_string(self.connection_string)
                self.logger.info("Connected to Azure Table Service using connection string")
            elif self.account_name and self.account_key:
                sas_key = self.generate_table_sas_key(table_name, permissions=table_permissions, expiry=table_expiry)
                self.table_service_client = TableServiceClient(sas_key)
                self.logger.info("Connected to Azure Table Service using SAS token")
            else:
                raise ValueError("Either connection string or both account name and account key must be provided")
        except Exception as e:
            self.logger.error(f"Error connecting to Azure Table Service: {e}")
