kimport os
import asyncio
from aiosmtpd.controller import Controller
from email.parser import BytesParser
from email.policy import default
import telegram
import re

# Telegram configuration
TELEGRAM_API_KEY = os.getenv('TELEGRAM_API_KEY')
DEFAULT_CHAT_ID = os.getenv('TELEGRAM_CHAT_ID')  # Fallback chat ID

# SMTP configuration
SMTP_HOST = '0.0.0.0'  # Listen on all interfaces
SMTP_PORT = 2525  # Alternative SMTP port

class SMTPHandler:
    async def handle_RCPT(self, server, session, envelope, address, rcpt_options):
        envelope.rcpt_tos.append(address)
        return '250 OK'

    async def handle_DATA(self, server, session, envelope):
        parser = BytesParser(policy=default)
        message = parser.parsebytes(envelope.content)
        
        subject = message.get('subject', 'No subject')
        body = message.get_body(preferencelist=('plain', 'html')).get_content()

        telegram_message = f"Subject: {subject}\n\n{body}"
        
        # Extract chat ID from the recipient email address
        chat_id = self.extract_chat_id(envelope.rcpt_tos[0])
        
        bot = telegram.Bot(TELEGRAM_API_KEY)
        try:
            await bot.send_message(chat_id=chat_id, text=telegram_message)
            print(f"Message sent to Telegram chat ID: {chat_id}")
        except telegram.error.BadRequest as e:
            print(f"Failed to send message to chat ID: {chat_id}. Error: {str(e)}")
            print("Attempting to send to default chat ID.")
            await bot.send_message(chat_id=DEFAULT_CHAT_ID, text=telegram_message)

        return '250 Message accepted for delivery'

    def extract_chat_id(self, email):
        match = re.match(r'(-?\d+)@telegram-relay\.local', email)
        if match:
            return match.group(1)
        print(f"Invalid email format: {email}. Using default chat ID.")
        return DEFAULT_CHAT_ID

async def start_smtp_server():
    controller = Controller(SMTPHandler(), hostname=SMTP_HOST, port=SMTP_PORT)
    controller.start()

if __name__ == '__main__':
    loop = asyncio.get_event_loop()
    loop.create_task(start_smtp_server())
    print(f"SMTP server started on {SMTP_HOST}:{SMTP_PORT}")
    loop.run_forever()