btgo-playwright/notifications.py

"""
Notification module for BTGO Scraper
Handles email and Discord notifications with file attachments
"""

import smtplib
import logging
import zipfile
import os
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText
from email.mime.base import MIMEBase
from email import encoders
from pathlib import Path
from datetime import datetime
from typing import List
import requests


class EmailNotifier:
    """Handles email notifications via SMTP"""

    def __init__(self, config):
        self.config = config
        self.enabled = config.EMAIL_ENABLED

    def send(self, files: List[str], accounts: list) -> bool:
        """
        Send email with CSV attachments

        Args:
            files: List of file paths to attach
            accounts: List of account data with balances

        Returns:
            True if successful, False otherwise
        """
        if not self.enabled:
            logging.info("Email notifications disabled")
            return False

        try:
            # Validate config
            if not all([self.config.SMTP_SERVER, self.config.EMAIL_FROM, self.config.EMAIL_TO]):
                logging.error("Email configuration incomplete")
                return False

            # Create message
            msg = MIMEMultipart()
            msg['From'] = self.config.EMAIL_FROM
            msg['To'] = self.config.EMAIL_TO
            msg['Subject'] = f'BTGO Scraper Results - {datetime.now().strftime("%Y-%m-%d %H:%M")}'

            # Email body
            body = self._create_email_body(files, accounts)
            msg.attach(MIMEText(body, 'html'))

            # Attach files
            total_size = 0
            for file_path in files:
                if not Path(file_path).exists():
                    logging.warning(f"File not found for email: {file_path}")
                    continue

                file_size = Path(file_path).stat().st_size
                total_size += file_size

                # Check size limit (10MB typical SMTP limit)
                if total_size > 10 * 1024 * 1024:
                    logging.warning(f"Email attachments exceed 10MB, creating ZIP archive")
                    return self._send_with_zip(files, accounts)

                with open(file_path, 'rb') as f:
                    part = MIMEBase('application', 'octet-stream')
                    part.set_payload(f.read())
                    encoders.encode_base64(part)
                    part.add_header(
                        'Content-Disposition',
                        f'attachment; filename={Path(file_path).name}'
                    )
                    msg.attach(part)

            # Send email
            logging.info(f"Sending email to {self.config.EMAIL_TO}...")
            with smtplib.SMTP(self.config.SMTP_SERVER, self.config.SMTP_PORT) as server:
                server.starttls()
                if self.config.SMTP_USERNAME and self.config.SMTP_PASSWORD:
                    server.login(self.config.SMTP_USERNAME, self.config.SMTP_PASSWORD)
                server.send_message(msg)

            logging.info(f"✓ Email sent successfully to {self.config.EMAIL_TO}")
            return True

        except Exception as e:
            logging.error(f"Failed to send email: {e}")
            return False

    def _send_with_zip(self, files: List[str], accounts: list) -> bool:
        """Send email with files compressed as ZIP"""
        try:
            # Create ZIP archive
            timestamp = datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
            zip_path = Path(self.config.OUTPUT_DIR) / f'btgo_export_{timestamp}.zip'

            with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
                for file_path in files:
                    if Path(file_path).exists():
                        zipf.write(file_path, Path(file_path).name)

            logging.info(f"Created ZIP archive: {zip_path}")

            # Send with ZIP attachment instead
            msg = MIMEMultipart()
            msg['From'] = self.config.EMAIL_FROM
            msg['To'] = self.config.EMAIL_TO
            msg['Subject'] = f'BTGO Scraper Results (ZIP) - {datetime.now().strftime("%Y-%m-%d %H:%M")}'

            body = self._create_email_body([str(zip_path)], accounts, is_zip=True)
            msg.attach(MIMEText(body, 'html'))

            with open(zip_path, 'rb') as f:
                part = MIMEBase('application', 'zip')
                part.set_payload(f.read())
                encoders.encode_base64(part)
                part.add_header('Content-Disposition', f'attachment; filename={zip_path.name}')
                msg.attach(part)

            with smtplib.SMTP(self.config.SMTP_SERVER, self.config.SMTP_PORT) as server:
                server.starttls()
                if self.config.SMTP_USERNAME and self.config.SMTP_PASSWORD:
                    server.login(self.config.SMTP_USERNAME, self.config.SMTP_PASSWORD)
                server.send_message(msg)

            logging.info(f"✓ Email with ZIP sent successfully to {self.config.EMAIL_TO}")
            return True

        except Exception as e:
            logging.error(f"Failed to send email with ZIP: {e}")
            return False

    def _create_email_body(self, files: List[str], accounts: list, is_zip: bool = False) -> str:
        """Create HTML email body"""
        file_list = '<br>'.join([f'• {Path(f).name}' for f in files])
        file_count = 1 if is_zip else len(files)

        # Calculate total balance
        total_ron = sum(acc['sold'] for acc in accounts if acc.get('moneda') == 'RON')

        # Build account balance list
        account_balances = ""
        for acc in accounts:
            nume = acc['nume_cont']
            sold = acc['sold']
            moneda = acc['moneda']
            account_balances += f'<tr><td style="padding: 8px; border-bottom: 1px solid #ecf0f1;">{nume}</td><td style="padding: 8px; text-align: right; border-bottom: 1px solid #ecf0f1;"><strong>{sold:,.2f} {moneda}</strong></td></tr>'

        return f"""
        <html>
        <body style="font-family: Arial, sans-serif;">
            <h2 style="color: #2c3e50;">BTGO Scraper Results</h2>
            <p><strong>Execution time:</strong> {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}</p>
            <p><strong>Accounts processed:</strong> {len(accounts)}</p>
            <p><strong>Files attached:</strong> {file_count}</p>
            <hr>
            <h3 style="color: #27ae60;">Solduri:</h3>
            <table style="width: 100%; border-collapse: collapse; margin-bottom: 20px;">
                {account_balances}
                <tr style="background-color: #f8f9fa;">
                    <td style="padding: 12px; font-weight: bold;">TOTAL</td>
                    <td style="padding: 12px; text-align: right; font-weight: bold; color: #27ae60; font-size: 1.2em;">{total_ron:,.2f} RON</td>
                </tr>
            </table>
            <hr>
            <h3>{'Archive contents:' if is_zip else 'Attached files:'}</h3>
            {file_list}
            <hr>
            <p style="color: #7f8c8d; font-size: 12px;">
                This email was automatically generated by BTGO Scraper.<br>
                For issues, check the logs at your deployment location.
            </p>
        </body>
        </html>
        """


class TelegramNotifier:
    """Handles Telegram notifications via Bot API"""

    def __init__(self, config):
        self.config = config
        self.enabled = config.TELEGRAM_ENABLED
        self.bot_token = config.TELEGRAM_BOT_TOKEN
        self.chat_id = config.TELEGRAM_CHAT_ID
        self.max_file_size = 50 * 1024 * 1024  # 50MB Telegram limit

    def send(self, files: List[str], accounts: list, telegram_message_id: str = None, telegram_chat_id: str = None) -> bool:
        """
        Send Telegram message with file attachments

        Args:
            files: List of file paths to attach
            accounts: List of account data with balances
            telegram_message_id: Message ID to edit (for progress updates)
            telegram_chat_id: Chat ID to edit (for progress updates)

        Returns:
            True if successful, False otherwise
        """
        if not self.enabled:
            logging.info("Telegram notifications disabled")
            return False

        try:
            if not self.bot_token or not self.chat_id:
                logging.error("Telegram bot token or chat ID not configured")
                return False

            # Store message IDs for editing
            self.progress_message_id = telegram_message_id
            self.progress_chat_id = telegram_chat_id or self.chat_id
            logging.info(f"Received telegram_message_id: {telegram_message_id}, telegram_chat_id: {telegram_chat_id}")
            logging.info(f"Stored progress_message_id: {self.progress_message_id}, progress_chat_id: {self.progress_chat_id}")

            # Check if SEND_AS_ZIP flag is set (from Telegram bot /scrape_zip command)
            send_as_zip = os.getenv('SEND_AS_ZIP', 'false').lower() == 'true'

            if send_as_zip:
                logging.info("SEND_AS_ZIP flag detected - sending as ZIP archive")
                return self._send_with_zip(files, accounts)

            # Check total file size
            total_size = sum(Path(f).stat().st_size for f in files if Path(f).exists())

            if total_size > self.max_file_size:
                logging.warning(f"Files exceed Telegram 50MB limit, creating ZIP archive")
                return self._send_with_zip(files, accounts)
            else:
                return self._send_files(files, accounts)

        except Exception as e:
            logging.error(f"Failed to send Telegram notification: {e}")
            return False

    def _send_message(self, text: str) -> bool:
        """Send text message to Telegram"""
        try:
            url = f"https://api.telegram.org/bot{self.bot_token}/sendMessage"
            payload = {
                'chat_id': self.chat_id,
                'text': text,
                'parse_mode': 'HTML'
            }
            response = requests.post(url, json=payload)
            return response.status_code == 200

        except Exception as e:
            logging.error(f"Failed to send Telegram message: {e}")
            return False

    def _edit_message(self, chat_id: str, message_id: str, text: str) -> bool:
        """Edit existing Telegram message (for progress updates)"""
        try:
            url = f"https://api.telegram.org/bot{self.bot_token}/editMessageText"
            payload = {
                'chat_id': chat_id,
                'message_id': message_id,
                'text': text,
                'parse_mode': 'Markdown'
            }
            response = requests.post(url, json=payload)
            return response.status_code == 200

        except Exception as e:
            # Silent fail pentru progress updates - nu blochez scraper-ul
            logging.debug(f"Progress update failed: {e}")
            return False

    def _send_files(self, files: List[str], accounts: list) -> bool:
        """Send files directly to Telegram"""
        try:
            # Calculate total balance
            total_ron = sum(acc['sold'] for acc in accounts if acc.get('moneda') == 'RON')

            # Build message with account balances
            message = f"<b>BTGO SCRAPER - FINALIZAT CU SUCCES</b>\n\n"
            message += f"Timp: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"
            message += f"Conturi procesate: {len(accounts)}\n"
            message += f"Fisiere: {len(files)}\n\n"

            message += "<b>SOLDURI:</b>\n"
            for acc in accounts:
                nume = acc['nume_cont']
                sold = acc['sold']
                moneda = acc['moneda']
                message += f"{nume}: {sold:,.2f} {moneda}\n"
            message += f"\n<b>TOTAL: {total_ron:,.2f} RON</b>\n"

            logging.info(f"Sending Telegram notification to chat {self.progress_chat_id}...")
            logging.info(f"Progress message_id: {self.progress_message_id}, chat_id: {self.progress_chat_id}")

            # Edit initial message or send new one
            if self.progress_message_id and self.progress_chat_id:
                # Edit the initial progress message
                logging.info(f"Editing initial message {self.progress_message_id} in chat {self.progress_chat_id}")
                url = f"https://api.telegram.org/bot{self.bot_token}/editMessageText"
                payload = {
                    'chat_id': self.progress_chat_id,
                    'message_id': self.progress_message_id,
                    'text': message,
                    'parse_mode': 'HTML'
                }
                response = requests.post(url, json=payload)
                if response.status_code == 200:
                    logging.info("✓ Message edited successfully with balances")
                else:
                    logging.error(f"Failed to edit message: {response.status_code} - {response.text}")
            else:
                # Send new message if no message_id provided
                logging.warning("No progress_message_id provided, sending new message")
                self._send_message(message)

            # Send each file as document
            url = f"https://api.telegram.org/bot{self.bot_token}/sendDocument"
            success_count = 0

            for file_path in files:
                if not Path(file_path).exists():
                    logging.warning(f"File not found: {file_path}")
                    continue

                file_size = Path(file_path).stat().st_size
                if file_size > self.max_file_size:
                    logging.warning(f"File too large for Telegram: {Path(file_path).name} ({file_size / 1024 / 1024:.2f} MB)")
                    continue

                with open(file_path, 'rb') as f:
                    files_data = {'document': f}
                    data = {
                        'chat_id': self.chat_id,
                        'caption': f"{Path(file_path).name}"
                    }
                    response = requests.post(url, data=data, files=files_data)

                    if response.status_code == 200:
                        logging.info(f"✓ Sent to Telegram: {Path(file_path).name}")
                        success_count += 1
                    else:
                        logging.error(f"Failed to send {Path(file_path).name}: {response.text}")

            if success_count > 0:
                logging.info(f"✓ Telegram notification sent successfully ({success_count}/{len(files)} files)")
                return True
            else:
                logging.error("No files were sent to Telegram")
                return False

        except Exception as e:
            logging.error(f"Failed to send Telegram files: {e}")
            return False

    def _send_with_zip(self, files: List[str], accounts: list) -> bool:
        """Send files as ZIP archive to Telegram"""
        try:
            # Create ZIP archive
            timestamp = datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
            zip_path = Path(self.config.OUTPUT_DIR) / f'btgo_export_{timestamp}.zip'

            with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
                for file_path in files:
                    if Path(file_path).exists():
                        zipf.write(file_path, Path(file_path).name)

            zip_size = Path(zip_path).stat().st_size
            logging.info(f"Created ZIP archive: {zip_path} ({zip_size / 1024 / 1024:.2f} MB)")

            # Check if ZIP is still too large
            if zip_size > self.max_file_size:
                logging.error(f"ZIP archive exceeds Telegram 50MB limit ({zip_size / 1024 / 1024:.2f} MB)")
                # Send message without attachment
                return self._send_message_only(accounts, files)

            # Calculate total balance
            total_ron = sum(acc['sold'] for acc in accounts if acc.get('moneda') == 'RON')

            # Build message with balances
            message = f"<b>BTGO SCRAPER - FINALIZAT CU SUCCES (ARHIVA ZIP)</b>\n\n"
            message += f"Timp: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"
            message += f"Conturi procesate: {len(accounts)}\n"
            message += f"Dimensiune arhiva: {zip_size / 1024 / 1024:.2f} MB\n"
            message += f"Fisiere in arhiva: {len(files)}\n\n"

            message += "<b>SOLDURI:</b>\n"
            for acc in accounts:
                nume = acc['nume_cont']
                sold = acc['sold']
                moneda = acc['moneda']
                message += f"{nume}: {sold:,.2f} {moneda}\n"
            message += f"\n<b>TOTAL: {total_ron:,.2f} RON</b>\n"

            # Edit initial message or send new one
            if self.progress_message_id and self.progress_chat_id:
                url = f"https://api.telegram.org/bot{self.bot_token}/editMessageText"
                payload = {
                    'chat_id': self.progress_chat_id,
                    'message_id': self.progress_message_id,
                    'text': message,
                    'parse_mode': 'HTML'
                }
                response = requests.post(url, json=payload)
                if response.status_code != 200:
                    logging.error(f"Failed to edit message: {response.text}")
            else:
                self._send_message(message)

            # Send ZIP file
            url = f"https://api.telegram.org/bot{self.bot_token}/sendDocument"
            with open(zip_path, 'rb') as f:
                files_data = {'document': f}
                data = {
                    'chat_id': self.chat_id,
                    'caption': f"BTGO Export Archive - {timestamp}"
                }
                response = requests.post(url, data=data, files=files_data)

            if response.status_code == 200:
                logging.info(f"✓ Telegram notification with ZIP sent successfully")
                return True
            else:
                logging.error(f"Telegram send failed: {response.status_code} - {response.text}")
                return False

        except Exception as e:
            logging.error(f"Failed to send Telegram ZIP: {e}")
            return False

    def _send_message_only(self, accounts: list, files: List[str]) -> bool:
        """Send Telegram message without files (when too large)"""
        try:
            file_list = '\n'.join([f'{Path(f).name} ({Path(f).stat().st_size / 1024:.1f} KB)' for f in files if Path(f).exists()])

            # Calculate total balance
            total_ron = sum(acc['sold'] for acc in accounts if acc.get('moneda') == 'RON')

            message = f"<b>BTGO SCRAPER - FINALIZAT CU SUCCES</b>\n\n"
            message += f"Timp: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"
            message += f"Conturi procesate: {len(accounts)}\n\n"

            message += "<b>SOLDURI:</b>\n"
            for acc in accounts:
                nume = acc['nume_cont']
                sold = acc['sold']
                moneda = acc['moneda']
                message += f"{nume}: {sold:,.2f} {moneda}\n"
            message += f"\n<b>TOTAL: {total_ron:,.2f} RON</b>\n\n"

            message += f"<b>ATENTIE:</b> Fisiere prea mari pentru Telegram\n\n"
            message += f"<b>Fisiere generate:</b>\n{file_list}\n\n"
            message += f"Verifica locatia de deployment pentru fisiere."

            # Edit initial message or send new one
            if self.progress_message_id and self.progress_chat_id:
                url = f"https://api.telegram.org/bot{self.bot_token}/editMessageText"
                payload = {
                    'chat_id': self.progress_chat_id,
                    'message_id': self.progress_message_id,
                    'text': message,
                    'parse_mode': 'HTML'
                }
                response = requests.post(url, json=payload)
                return response.status_code == 200
            else:
                return self._send_message(message)

        except Exception as e:
            logging.error(f"Failed to send Telegram message: {e}")
            return False


class NotificationService:
    """Orchestrates all notification channels"""

    def __init__(self, config):
        self.config = config
        self.email = EmailNotifier(config)
        self.telegram = TelegramNotifier(config)

    def send_all(self, files: List[str], accounts: list, telegram_message_id: str = None, telegram_chat_id: str = None) -> dict:
        """
        Send notifications via all enabled channels

        Args:
            files: List of file paths to send
            accounts: List of account data with balances
            telegram_message_id: Message ID to edit (for Telegram progress updates)
            telegram_chat_id: Chat ID to edit (for Telegram progress updates)

        Returns:
            Dictionary with status of each channel
        """
        results = {
            'email': False,
            'telegram': False
        }

        account_count = len(accounts)
        logging.info(f"Sending notifications for {len(files)} files...")

        # Send via email
        if self.config.EMAIL_ENABLED:
            results['email'] = self.email.send(files, accounts)

        # Send via Telegram
        if self.config.TELEGRAM_ENABLED:
            results['telegram'] = self.telegram.send(
                files,
                accounts,
                telegram_message_id=telegram_message_id,
                telegram_chat_id=telegram_chat_id
            )

        # Log summary
        success_count = sum(results.values())
        total_enabled = sum([self.config.EMAIL_ENABLED, self.config.TELEGRAM_ENABLED])

        if success_count == total_enabled:
            logging.info(f"✓ All notifications sent successfully ({success_count}/{total_enabled})")
        elif success_count > 0:
            logging.warning(f"Partial notification success ({success_count}/{total_enabled})")
        else:
            logging.error("All notifications failed")

        return results