"""
S3 Log Manager for Optimizer Factory
Handles all S3 operations for log file storage and retrieval.

This module provides a centralized interface for:
1. Uploading log archives to S3
2. Downloading individual log files from S3
3. Listing archived log sessions for repositories
4. Managing S3 object keys and directory structure
"""

import os
import boto3
import logging
from botocore.exceptions import ClientError
from datetime import datetime
from typing import Optional, List, Dict, Any

logger = logging.getLogger(__name__)

class S3LogManager:
    """
    Manages log file operations with an AWS S3 bucket.
    Handles uploading, downloading, and listing of log files.
    """
    
    def __init__(self, bucket_name: str, region_name: str):
        """
        Initializes the S3LogManager.

        Args:
            bucket_name: The name of the S3 bucket.
            region_name: The AWS region of the bucket.
        """
        if not bucket_name:
            raise ValueError("S3 bucket name cannot be empty.")
        
        self.bucket_name = bucket_name
        
        # Configure S3 client with optimized connection settings
        from botocore.config import Config
        config = Config(
            max_pool_connections=20,  # Increase connection pool size
            retries={'max_attempts': 3, 'mode': 'adaptive'}
        )
        
        self.s3_client = boto3.client("s3", region_name=region_name, config=config)
        
        logger.info(f"S3LogManager initialized for bucket '{bucket_name}' in region '{region_name}'")

    def _generate_s3_key(self, repo_url: str, filename: str) -> str:
        """
        Generates a structured S3 key based on the desired naming convention.
        Format: <repo-slug>/<YYYY-MM-DD_HH-MM-SS>/<filename>
        
        Args:
            repo_url: The repository URL to generate slug from
            filename: The filename to append to the key
            
        Returns:
            S3 key string in format: org-repo/2025-01-15_14-30-25/filename.log
        """
        # Create a URL-safe slug from the repo URL
        try:
            parts = (repo_url or "").rstrip("/").split("/")
            if len(parts) >= 2:
                org = parts[-2]
                name = parts[-1].replace(".git", "")
                repo_slug = f"{org}-{name}"
            else:
                repo_slug = "unknown-repo"
        except (IndexError, AttributeError):
            repo_slug = "unknown-repo"
        
        # Get current timestamp for the archive folder
        timestamp = datetime.utcnow().strftime('%Y-%m-%d_%H-%M-%S')

        return f"{repo_slug}/{timestamp}/{filename}"

    def upload_log_archive(self, local_dir_path: str, repo_url: str) -> Optional[str]:
        """
        Uploads all files from a local directory to a structured path in S3.

        Args:
            local_dir_path: The local directory containing log files.
            repo_url: The repository URL to structure the S3 path.
        
        Returns:
            The S3 prefix (folder path) where the logs were uploaded, or None if failed.
        """
        uploaded_files = []
        s3_prefix = ""
        
        logger.info(f"Starting upload of log archive from {local_dir_path} for repo {repo_url}")
        
        for root, _, files in os.walk(local_dir_path):
            for filename in files:
                local_path = os.path.join(root, filename)
                
                # Generate prefix once based on the first file
                if not s3_prefix:
                    s3_prefix = os.path.dirname(self._generate_s3_key(repo_url, filename))
                
                s3_key = f"{s3_prefix}/{filename}"
                
                try:
                    # Check if file exists before attempting upload
                    if not os.path.exists(local_path):
                        logger.warning(f"Local file not found for upload: {local_path}")
                        continue
                    
                    self.s3_client.upload_file(local_path, self.bucket_name, s3_key)
                    logger.info(f"Successfully uploaded {local_path} to s3://{self.bucket_name}/{s3_key}")
                    uploaded_files.append(s3_key)
                except ClientError as e:
                    logger.error(f"Failed to upload {local_path} to S3: {e}")
                except FileNotFoundError:
                    logger.warning(f"Local file not found for upload: {local_path}")
        
        if uploaded_files:
            logger.info(f"Upload completed: {len(uploaded_files)} files uploaded to s3://{self.bucket_name}/{s3_prefix}")
            return s3_prefix
        else:
            logger.error("No files were successfully uploaded")
            return None

    def download_log_stream(self, s3_key: str):
        """
        Downloads a log file from S3 and returns a streaming response object.
        
        Args:
            s3_key: The S3 object key to download
            
        Returns:
            S3 response body object for streaming, or None if not found
        """
        try:
            response = self.s3_client.get_object(Bucket=self.bucket_name, Key=s3_key)
            logger.info(f"Successfully initiated download of s3://{self.bucket_name}/{s3_key}")
            return response['Body']
        except ClientError as e:
            if e.response['Error']['Code'] == 'NoSuchKey':
                logger.error(f"S3 key not found: {s3_key}")
            else:
                logger.error(f"Error downloading from S3: {e}")
            return None

    def list_log_archives(self, repo_url: str) -> List[Dict[str, Any]]:
        """
        Lists all archived log sessions for a given repository.
        A "session" is a unique folder under the repository's slug.
        
        Args:
            repo_url: The repository URL to list archives for
            
        Returns:
            List of session dictionaries with session_id and files
        """
        try:
            parts = (repo_url or "").rstrip("/").split("/")
            if len(parts) >= 2:
                org = parts[-2]
                name = parts[-1].replace(".git", "")
                repo_slug = f"{org}-{name}"
            else:
                logger.warning(f"Could not parse repo URL for slug generation: {repo_url}")
                return []
        except (IndexError, AttributeError):
            logger.warning(f"Invalid repo URL format: {repo_url}")
            return []
            
        logger.info(f"Listing log archives for repo slug: {repo_slug}")
        
        try:
            paginator = self.s3_client.get_paginator('list_objects_v2')
            pages = paginator.paginate(Bucket=self.bucket_name, Prefix=f"{repo_slug}/", Delimiter='/')
            
            sessions = []
            for page in pages:
                for prefix in page.get('CommonPrefixes', []):
                    session_id = prefix.get('Prefix').strip('/')
                    files = self._list_files_in_session(session_id)
                    sessions.append({
                        "session_id": session_id, 
                        "files": files
                    })
            
            # Sort sessions by name (timestamp) descending
            sessions.sort(key=lambda x: x['session_id'], reverse=True)
            
            logger.info(f"Found {len(sessions)} log sessions for {repo_slug}")
            return sessions
            
        except ClientError as e:
            logger.error(f"Error listing S3 objects for {repo_slug}: {e}")
            return []

    def _list_files_in_session(self, session_prefix: str) -> List[Dict[str, Any]]:
        """
        Helper to list all files within a specific session prefix.
        
        Args:
            session_prefix: The S3 prefix for the session (e.g., "org-repo/2025-01-15_14-30-25")
            
        Returns:
            List of file dictionaries with metadata
        """
        files = []
        try:
            response = self.s3_client.list_objects_v2(Bucket=self.bucket_name, Prefix=f"{session_prefix}/")
            for obj in response.get('Contents', []):
                files.append({
                    "filename": os.path.basename(obj['Key']),
                    "path": obj['Key'],  # The full S3 key is now the path
                    "size_bytes": obj['Size'],
                    "modified_iso": obj['LastModified'].isoformat(),
                })
            
            logger.debug(f"Found {len(files)} files in session {session_prefix}")
            return files
            
        except ClientError as e:
            logger.error(f"Error listing files in session {session_prefix}: {e}")
            return []

    def list_session_files(self, session_id: str) -> List[Dict[str, Any]]:
        """
        Lists all files in a specific log session for downloading.
        
        Args:
            session_id: The session identifier (e.g., "org-repo/2025-01-15_14-30-25")
            
        Returns:
            List of file dictionaries with s3_key for downloading
        """
        files = []
        try:
            response = self.s3_client.list_objects_v2(Bucket=self.bucket_name, Prefix=f"{session_id}/")
            for obj in response.get('Contents', []):
                files.append({
                    "filename": os.path.basename(obj['Key']),
                    "s3_key": obj['Key'],  # Full S3 key for downloading
                    "size_bytes": obj['Size'],
                    "modified_iso": obj['LastModified'].isoformat(),
                })
            
            logger.info(f"Found {len(files)} files in session {session_id}")
            return files
            
        except ClientError as e:
            logger.error(f"Error listing files in session {session_id}: {e}")
            return []
