"""
Azure Blob Storage Integration - Handles Azure Blob Storage communication.
Manages file uploads, SAS token generation, and storage operations.
"""

from azure.storage.blob import BlobServiceClient, generate_blob_sas, BlobSasPermissions
from azure.core.exceptions import ResourceNotFoundError
import os
import logging
import time
from datetime import datetime, timedelta

from core.logger import get_logger

DEFAULT_CONTAINER_NAME = "test-screenshots"
SAS_TOKEN_VALIDITY_DAYS = 7
MAX_UPLOAD_RETRIES = 3
RETRY_BASE_DELAY_SECONDS = 2


def _suppress_azure_logging():
    """Suppress verbose Azure SDK logging."""
    azure_loggers = [
        'azure.core.pipeline.policies.http_logging_policy',
        'azure.storage.blob',
        'azure.core.pipeline',
        'urllib3.connectionpool',
    ]
    for logger_name in azure_loggers:
        logging.getLogger(logger_name).setLevel(logging.WARNING)


_suppress_azure_logging()


class AzureBlobService:
    """
    Integration service for Azure Blob Storage.
    Handles file uploads, SAS token generation, and container management.
    """

    def __init__(self, logger=None):
        """Initialize Azure Blob Storage service."""
        self.logger = logger or get_logger(__name__)
        connection_string = os.getenv("AZURE_STORAGE_CONNECTION_STRING")
        self.connection_string = connection_string
        try:
            self.blob_service_client = BlobServiceClient.from_connection_string(connection_string)
            self.container_name = DEFAULT_CONTAINER_NAME
            self._ensure_container_exists()
        except Exception as e:
            self.logger.error(f"Failed to initialize blob storage: {str(e)}")
            raise

    def _ensure_container_exists(self):
        """Ensure the blob container exists, creating it if necessary."""
        try:
            container_client = self.blob_service_client.get_container_client(self.container_name)
            try:
                container_client.get_container_properties()
                self.logger.info(f"Container '{self.container_name}' exists")
            except ResourceNotFoundError:
                container_client.create_container()
                self.logger.info(f"Created container '{self.container_name}'")
        except Exception as e:
            self.logger.error(f"Error ensuring container exists: {str(e)}")
            raise

    def _generate_sas_token(self, blob_name: str) -> str:
        """Generate a SAS token for the specified blob."""
        try:
            account_name = self.blob_service_client.account_name
            account_key = self.blob_service_client.credential.account_key
            sas_token = generate_blob_sas(
                account_name=account_name,
                container_name=self.container_name,
                blob_name=blob_name,
                account_key=account_key,
                permission=BlobSasPermissions(read=True),
                expiry=datetime.utcnow() + timedelta(days=SAS_TOKEN_VALIDITY_DAYS)
            )
            return sas_token
        except Exception as e:
            self.logger.error(f"Error generating SAS token: {str(e)}")
            raise

    def upload_screenshot(self, local_file_path: str, test_run_id: str) -> str:
        """Upload a screenshot file to blob storage with retry logic."""
        if not os.path.exists(local_file_path):
            raise FileNotFoundError(f"Screenshot file not found: {local_file_path}")

        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = os.path.basename(local_file_path)
        blob_name = f"{test_run_id}/{timestamp}_{filename}"
        return self._upload_file_with_retry(local_file_path, blob_name, "screenshot")

    def upload_html_report(self, local_file_path: str, test_run_id: str) -> str:
        """Upload an HTML report file to blob storage with retry logic."""
        if not os.path.exists(local_file_path):
            raise FileNotFoundError(f"HTML report file not found: {local_file_path}")

        filename = os.path.basename(local_file_path)
        blob_name = f"reports/{test_run_id}/{filename}"
        return self._upload_file_with_retry(local_file_path, blob_name, "HTML report")

    def _upload_file_with_retry(self, local_file_path: str, blob_name: str, file_type: str) -> str:
        """Upload a file to blob storage with exponential backoff retry logic."""
        container_client = self.blob_service_client.get_container_client(self.container_name)
        blob_client = container_client.get_blob_client(blob_name)
        retry_count = 0
        last_error = None

        while retry_count < MAX_UPLOAD_RETRIES:
            try:
                with open(local_file_path, "rb") as file_data:
                    blob_client.upload_blob(file_data, overwrite=True)
                sas_token = self._generate_sas_token(blob_name)
                return f"{blob_client.url}?{sas_token}"
            except Exception as e:
                last_error = e
                retry_count += 1
                if retry_count < MAX_UPLOAD_RETRIES:
                    delay = RETRY_BASE_DELAY_SECONDS ** retry_count
                    self.logger.warning(f"Upload attempt {retry_count} failed for {file_type}: {str(e)}. Retrying in {delay}s...")
                    time.sleep(delay)
                else:
                    self.logger.error(f"Failed to upload {file_type} after {MAX_UPLOAD_RETRIES} attempts: {str(e)}")
                    raise last_error


# Backward compatibility - maintain existing API
_blob_service_instance = AzureBlobService()


class BlobStorageManager:
    """Backward compatibility wrapper for existing code."""
    
    def __init__(self):
        """Initialize using singleton instance."""
        self._service = _blob_service_instance
    
    def upload_screenshot(self, local_file_path: str, test_run_id: str) -> str:
        """Upload a screenshot file to blob storage."""
        return self._service.upload_screenshot(local_file_path, test_run_id)
    
    def upload_html_report(self, local_file_path: str, test_run_id: str) -> str:
        """Upload an HTML report file to blob storage."""
        return self._service.upload_html_report(local_file_path, test_run_id)


__all__ = ["AzureBlobService", "BlobStorageManager"]
