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 AI bot. Ask me aything with /ask")

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


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"

def ask(update, context):
    user_message = ' '.join(context.args)
    llm_response = query_llm(user_message)
    update.message.reply_text(llm_response)

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_handler(update: Update): #context: CallbackContext
    """Handle errors occurred during the execution of commands."""
    # Log the error before we do anything else
    # logger.error(f"An error occurred: {context.error}")
    logger.error(f"An error occurred:")

    # Send a message to the user
    # error_message = "Sorry, an unexpected error occurred. Please try again later."
    # if update.effective_message:
    #     await context.bot.send_message(chat_id=update.effective_chat.id, text=error_message)


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

    # Add handlers to the application
    # Command handlers should be registered before the generic message handler
    application.add_handler(CommandHandler("start", start))
    application.add_handler(CommandHandler("screenshot", screenshot)) # Ensure screenshot function is async
    application.add_handler(CommandHandler("ai", ask))
    application.add_handler(CommandHandler("ask", ask))
    application.add_handler(CommandHandler("llm", ask))
    
    # This handler should be last as it's the most generic
    application.add_handler(MessageHandler(filters.TEXT, echo))

    # Register the error handler
    application.add_error_handler(error_handler)

    # 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())