import json
from google.oauth2 import service_account
from googleapiclient.discovery import build
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart
import base64


class GmailClient:
    """Gmail API client using service account with domain-wide delegation."""

    def __init__(self, key, impersonator):
        """
        Initialize Gmail client.

        Args:
            key: str or dict : Service account credentials (JSON string or dict)
            impersonator: str : Email address of user to impersonate
        """
        try:
            service_account_key = json.loads(key) if isinstance(key, str) else key
        except json.JSONDecodeError:
            raise ValueError("Invalid JSON key provided.")

        self.impersonator = impersonator
        self.scopes = ['https://www.googleapis.com/auth/gmail.send']
        self.credentials = service_account.Credentials.from_service_account_info(
            service_account_key, scopes=self.scopes
        ).with_subject(self.impersonator)
        self.service = build('gmail', 'v1', credentials=self.credentials)

    def create_message(self, sender, to, subject, message_text, message_html):
        """
        Create an email message with both plain text and HTML parts.

        Args:
            sender: str : Sender email address
            to: str : Recipient email address
            subject: str : Email subject
            message_text: str : Plain text body
            message_html: str : HTML body

        Returns:
            dict : Message ready to be sent via Gmail API
        """
        message = MIMEMultipart('alternative')
        message['to'] = to
        message['from'] = sender
        message['subject'] = subject

        text_part = MIMEText(message_text, 'plain')
        message.attach(text_part)

        html_part = MIMEText(message_html, 'html')
        message.attach(html_part)

        raw_message = {'raw': base64.urlsafe_b64encode(message.as_bytes()).decode()}
        return raw_message

    def send_message(self, message):
        """
        Send an email message via Gmail API.

        Args:
            message: dict : Message object from create_message()

        Returns:
            dict : Sent message response or None on error
        """
        try:
            message = self.service.users().messages().send(userId='me', body=message).execute()
            print(f'Message sent. Message ID: {message["id"]}')
            return message
        except Exception as e:
            print(f'An error occurred: {e}')
            return None