# src/utils/google_drive_service.py
from google.oauth2 import service_account
from googleapiclient.discovery import build
from googleapiclient.http import MediaIoBaseDownload
import io
from typing import List, Dict, Any
import logging

class GoogleDriveService:
    def __init__(self, credentials_path: str):
        """
        Initialize Google Drive service
        
        Args:
            credentials_path (str): Path to service account credentials file
        """
        self.credentials = service_account.Credentials.from_service_account_file(
            credentials_path,
            scopes=['https://www.googleapis.com/auth/drive.readonly']
        )
        self.service = build('drive', 'v3', credentials=self.credentials)

    def get_folder_contents(self, folder_id: str, include_subfolders: bool = False) -> List[Dict[str, Any]]:
        """
        Get contents of a Drive folder including subfolders if specified
        
        Args:
            folder_id (str): ID of the folder to process
            include_subfolders (bool): Whether to include contents of subfolders (default: False)
            
        Returns:
            List[Dict]: List of file metadata
        """
        all_files = []
        try:
            # Get all items in the current folder
            query = f"'{folder_id}' in parents and trashed=false"
            results = self.service.files().list(
                q=query,
                fields="files(id, name, mimeType, modifiedTime, parents)",
                supportsAllDrives=True,
                includeItemsFromAllDrives=True
            ).execute()
            
            items = results.get('files', [])
            
            for item in items:
                if item['mimeType'] == 'application/vnd.google-apps.folder' and include_subfolders:
                    # Recursively get contents of subfolder
                    try:
                        subfolder_files = self.get_folder_contents(
                            item['id'],
                            include_subfolders=True
                        )
                        # Add folder path information to each file
                        for file in subfolder_files:
                            if not file.get('folder_path'):
                                file['folder_path'] = []
                            file['folder_path'].insert(0, {
                                'id': item['id'],
                                'name': item['name']
                            })
                        all_files.extend(subfolder_files)
                    except Exception as e:
                        logging.error(f"Error processing subfolder {item['name']}: {str(e)}")
                        continue
                else:
                    # For backward compatibility, maintain original structure
                    # but add folder path information
                    item['folder_path'] = []
                    all_files.append(item)
                    
            return all_files
            
        except Exception as e:
            logging.error(f"Error getting folder contents for folder {folder_id}: {str(e)}")
            return []  # Return empty list for backward compatibility

    def download_file(self, file_id: str) -> bytes:
        """
        Download a file from Drive
        
        Args:
            file_id (str): ID of the file to download
            
        Returns:
            bytes: File content
        """
        request = self.service.files().get_media(fileId=file_id)
        content = io.BytesIO()
        downloader = MediaIoBaseDownload(content, request)
        
        done = False
        while not done:
            _, done = downloader.next_chunk()
            
        content.seek(0)
        return content.read()

    def export_file(self, file_id: str, mime_type: str) -> bytes:
        """
        Export a Google Workspace file to a different format
        
        Args:
            file_id (str): ID of the file to export
            mime_type (str): MIME type to export to
            
        Returns:
            bytes: Exported file content
        """
        request = self.service.files().export_media(
            fileId=file_id,
            mimeType=mime_type
        )
        content = io.BytesIO()
        downloader = MediaIoBaseDownload(content, request)
        
        done = False
        while not done:
            _, done = downloader.next_chunk()
            
        content.seek(0)
        return content.read()