import logging
from telegram.ext import Updater, CommandHandler, MessageHandler, filters, CallbackContext


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

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

def start(update, context):
    """Send a message when the command /start is issued."""
    update.message.reply_text('Hi I am ! Send me a message, and I will interact with LLM.')

def echo(update, context):
    """Echo the user message."""
    user_message = update.message.text
    response = query_llm(user_message)
    update.message.reply_text(response)

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 screenshot(update, context):
    """Take a screenshot of a webpage."""
    url = ' '.join(context.args)
    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'))

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

def main():
    """Start the bot."""
    updater = Updater(TOKEN, use_context=True)

    dp = updater.dispatcher
    dp.add_handler(CommandHandler("start", start))
    dp.add_handler(MessageHandler(Filters.text, echo))
    dp.add_handler(CommandHandler("screenshot", screenshot, pass_args=True))
    dp.add_error_handler(error)

    updater.start_polling()
    updater.idle()

if __name__ == '__main__':
    main()