import logging
import asyncio, nest_asyncio
from telegram import Bot, Update
from telegram.ext import Application, CommandHandler, MessageHandler, filters

# import "gopkg.in/telebot.v3/middleware"

import requests
import json
import base64
from selenium import webdriver
from selenium.webdriver.chrome.options import Options
from io import BytesIO
from PIL import Image

# Apply nest_asyncio
nest_asyncio.apply()

# Set up logging
logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', level=logging.INFO)
logger = logging.getLogger(__name__)

# Telegram Bot Token
# t.me/kevin_ai_robot
TOKEN = '6805059978:AAHNJKuOeazMSJHc3-BXRCsFfEVyFHeFnjw'
# t.me/artitherobot 6749075936:AAHUHiPTDEIu6JH7S2fQdibwsu6JVG3FNG0


# LLM API Endpoint
LLM_ENDPOINT = "http://192.168.0.11:11434/api/chat"

#! Selenium WebDriver setup for screenshots
#chrome_options = Options()
#chrome_options.add_argument("--headless")
#driver = webdriver.Chrome(options=chrome_options)

async def start(update: Update, context):
    await context.bot.send_message(chat_id=update.effective_chat.id, text="Hi! I'm your bot.")

async def echo(update: Update, context):
    await context.bot.send_message(chat_id=update.effective_chat.id, text=update.message.text)

async def query_llm(user_message):
    """Query the LLM with the user's message."""
    data = {
        "model": "llama2",
        "messages": [{"role": "user", "content": user_message}]
    }
    response = requests.post(LLM_ENDPOINT, json=data)
    if response.status_code == 200:
        response_data = response.json()
        return response_data.get('message', {}).get('content', 'No response')
    else:
        return "Error: Unable to reach LLM"

async def screenshot(update, context):
    """Take a screenshot of a webpage."""
    url = ' '.join(context.args)
    update.message.reply_text('This will noramlly get a screenshot from: '.url)# url.', but currently under development')


    # driver.get(url)
    # screenshot = driver.get_screenshot_as_png()
    # image_stream = BytesIO(screenshot)
    # image_stream.seek(0)
    # image = Image.open(image_stream)
    # image_stream.close()
    # image.save('screenshot.png')
    # update.message.reply_photo(photo=open('screenshot.png', 'rb'))

async def error(update, context):
    """Log Errors caused by Updates."""
    logger.warning(f'Update "{update}" caused error "{context.error}"')


async def main():
    """Start the bot."""
    # Create an Application instance
    application = Application.builder().token(TOKEN).build()

    # Add handlers to the application
    application.add_handler(CommandHandler("start", start))
    application.add_handler(MessageHandler(filters.TEXT, echo))
    application.add_handler(CommandHandler("screenshot", screenshot)) # Ensure screenshot function is async
    application.add_handler(CommandHandler("ai",query_llm))
    application.add_error_handler(error) # Ensure error function is async

    # Run the bot
    await application.run_polling()


if __name__ == '__main__':
    loop = asyncio.get_event_loop()
    if loop.is_running():
        loop.create_task(main())
    else:
        asyncio.run(main())