import logging
import os
import asyncio
from telethon import TelegramClient, events, sync
from llm import generate_llm  # Assuming these are your custom functions
from sd import generate_sd    # Replace with actual imports as needed
import socks

# Load environment variables
def load_env():
    global TELEGRAM_API_ID, TELEGRAM_API_HASH, TELEGRAM_BOT_TOKEN
    TELEGRAM_API_ID = int(os.getenv("TELEGRAM_API_ID"))
    TELEGRAM_API_HASH = os.getenv("TELEGRAM_API_HASH")
    TELEGRAM_BOT_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN")
    if not all([TELEGRAM_API_ID, TELEGRAM_API_HASH, TELEGRAM_BOT_TOKEN]):
        raise ValueError("One or more required environment variables are not set")

load_env()

# Configure logging
LOG_FILE = "bot.log"
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler(LOG_FILE),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

proxy_server = '142.93.68.63'
proxy_port = 2434
proxy_secret = 'ee32b920dffb51643028e2f6b878d4eac1666172616b61762e636f6d'
proxy_dc_id = 2  # This is usually 2 for MTProto proxies

proxy = (
    socks.SOCKS5,
    proxy_server,
    proxy_port,
    True,
    'vpn',
    'unlimited'
)

async def main():
    # Initialize Telethon client
    client = TelegramClient('bot', TELEGRAM_API_ID, TELEGRAM_API_HASH, proxy=proxy).start(TELEGRAM_BOT_TOKEN)

    @client.on(events.NewMessage)
    async def handle_message(event):
        chat_id = event.chat_id
        user_message = event.raw_text
        logger.info(f"Chat ID: {chat_id} - Received Message: {user_message}")

        # Handle /imagine command
        if user_message.startswith('/imagine'):
            prompt = user_message.replace('/imagine', '').strip()
            if not prompt:
                await event.reply("Please provide a prompt after /imagine.")
                logger.info(f"Chat ID: {chat_id} - Sent Message: Please provide a prompt after /imagine.")
                return

            generating_message = await event.reply("Generating...")
            logger.info(f"Chat ID: {chat_id} - Sent Message: Generating...")

            try:
                # Example function call to generate image data
                image_data, image_path = generate_sd(prompt)
                await client.delete_messages(chat_id, [generating_message.id])

                if image_data:
                    await client.send_file(chat_id, image_data)
                    logger.info(f"Chat ID: {chat_id} - Sent Image {image_path}")
                else:
                    await event.reply("Failed to generate image. Please try again later.")
                    logger.error(f"Chat ID: {chat_id} - Failed to generate image.")
            except Exception as e:
                logger.error(f"Chat ID: {chat_id} - Error in handle_image_generation: {e}")
                await event.reply("There was an error generating the image. Please try again later.")

        else:
            # Default response handling
            try:
                msg = generate_llm(user_message)  # Example function call to generate a response
                await event.reply(msg)
                logger.info(f"Chat ID: {chat_id} - Sent Message: {msg}")
            except Exception as e:
                logger.error(f"Chat ID: {chat_id} - Error in response_text: {e}")
                await event.reply("There was an error processing your request.")

    await client.run_until_disconnected()

# Run the main function
if __name__ == '__main__':
    asyncio.run(main())