#!/usr/bin/env python
import os
import re
import json
import asyncio
from tempfile import gettempdir
from math import floor
from telegram import Bot, Update, Message, Chat, User
from telegram.constants import MessageLimit
from telegram.ext import Application, CommandHandler, MessageHandler, filters, ContextTypes
import openai

from .openai_singleton import OpenAIClientSingleton
from .initialize import INIT_BOT_CONFIG, init_structure, init_bot_config, init_bot_prompt
from .conversation import Conversation
from .tokenGPT import TokenGPT
from .message_handlers import handle_greetings, handle_common_queries, handle_url_ask
from .utils import read_yaml, read_text, exact_word_match, generate_error_path, log_error

class TelegramBot:
    """
    The bridge between the Telegram interface and Large Langage Model (LLM) using OpenAI's GPT models.
    - Tokens help retain conversation messages between the Telegram bot assistant and the user.
    - URLs in [square brackets] can be supplied on how the bot should interpret it.
    - For more information, see README or: https://github.com/Digital-Heresy/TeLLMgramBot
    """
    async def _tele_info(self):
        """Fetch the Telegram bot's information before running application."""
        try:
            bot = await Bot(token=os.environ['TELLMGRAMBOT_TELEGRAM_API_KEY']).get_me()
            self.telegram['username'] = bot.username
            self.telegram['name'] = bot.full_name
        except Exception as e:
            # Fail fast if the bot's identity cannot be fetched to avoid
            # continuing with an uninitialized self.telegram['username'].
            log_error(e, error_type='Telegram', error_filename=self.error_log)
            raise

    async def tele_commands(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
        """Show all the commands available when the user sends `/help` to the bot."""
        await update.message.reply_text("TeLLMgramBot commands:\n"
            "/nick - Set your nickname (e.g. \"/nick B0b #2\").\n"
            "/forget - Clear all of your conversations.\n"
            "/help - Display this usage information.\n\n"
            "You must have a username to converse (for now).\n"
            "Chat history will still show in your Telegram client.\n\n"
            "Administrator-only commands:\n"
            "/start - Go online to receive new responses (default).\n"
            "/stop - Go offline to prevent new responses.\n"
        )

    async def tele_start_command(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
        """A bot_owner-only command to go online and handle any new messages."""
        uname = update.message.from_user.username
        if uname == self.telegram['owner']:
            self._online = True
            greeting_text = f"Oh, hello {update.message.from_user.first_name}! Let me get to work!"
            await update.message.reply_text(greeting_text)
        else:
            await update.message.reply_text("Sorry, I can't do that for you.")

    async def tele_stop_command(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
        """A bot_owner-only command to go offline and not handle any new messages."""
        uname = update.message.from_user.username
        if uname == self.telegram['owner']:
            self._online = False
            await update.message.reply_text("Sure thing boss, cutting out!")
        else:
            await update.message.reply_text("Sorry, I can't do that for you.")

    async def tele_nick_command(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
        """
        Let bot assistant know to call the user by a different name other than the Telegram username.
        Must follow the format of "/nick(@<bot name>)? <nickname>", checked by a regular expression.
        """
        validated = await self.tele_validate(update)
        if not validated:
            return
        (msg, chat, user) = validated

        # If it's a group conversation, need to handle commands directly to bot with @<bot name> also
        pattern = re.compile(fr"^\s*/nick(?:@{re.escape(self.telegram['username'])})?\s+(.+?)\s*$", re.IGNORECASE)
        result = pattern.search(msg.text or "")
        if result:
            prompt = f"Please refer to me by my nickname, {result.group(1).strip()}, rather than my user name."
            await msg.reply_text(await self.tele_handle_response(prompt, msg))
        else:
            await msg.reply_text("Please provide a valid nickname after the command like \"/nick B0b #2\".")

    async def tele_forget_command(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
        """Remove the whole conversation including past session logs for peace of mind."""
        uname = update.message.from_user.username
        if uname in self.conversations:
            self.conversations[uname].clear_interaction()
            del self.conversations[uname]  # Delete object as well
            await update.message.reply_text("My memories of all our conversations are wiped!")
        else:
            await update.message.reply_text(
                "My apologies, but I don't recall any conversations with you, or you"
                " already asked me to forget about you. Either way, nice to meet you!"
            )

    async def tele_handle_response(self, text: str, msg: Message) -> str:
        """
        Primary function for handling any response including Generative AI, ensuring:
        - The owner started up the bot assistant for user interactions.
        - User has conversed with the bot assistant before.
        - Depending on regular interaction or URL request, formulate message based
          on the Generative AI model (e.g. OpenAI GPT) and token count checks.
        """
        # Starting ensures we get some kind of user account details for logging
        if not self._online:
            return "I'd love to chat, but I am offline at the moment!"

        # Print information of the user sending message for conversation
        uname = msg.from_user.username
        print(f"User {uname} in {msg.chat.type} chat ID {str(msg.chat.id)}")

        # For a new session, track if the user has conversed with the Telegram bot before
        # Username is consistent across restarts and different conversation instances
        if uname not in self.conversations:
            self.conversations[uname] = Conversation(
                uname,
                self.telegram['username'],
                self.chatgpt['prompt'],
                self.chatgpt['chat_model']
            )
            # If there are past conversations via logs, load by 50% threshold of tokens
            self.conversations[uname].get_past_interaction(
                floor(self.chatgpt['prune_threshold'] / 2)
            )

        # Add the user's message to our conversation
        self.conversations[uname].add_user_message(text)

        # Check if the user is asking about a [URL]
        url_match = re.search(r'\[http(s)?://\S+]', text)

        # Form the assistant's message based on low level easy stuff or send to GPT
        # OpenAI's Responses API relies on the maximum number of tokens the selected model can support
        reply = "Sorry, I couldn't process your message! Please contact my creator."
        if handle_greetings(text):
            reply = handle_greetings(text)
        elif handle_common_queries(text):
            reply = handle_common_queries(text)
        elif url_match:
            # URL content is passed into another model to summarize (GPT-4 preferred)
            await msg.reply_text("Sure, give me a moment to look at that URL...")
            reply = await handle_url_ask(text, self.chatgpt['url_model'])
        elif self._online:
            # This is essentially the transition point between quick Telegram replies and GPT
            response = await self.gpt_completion(uname)
            reply = response.choices[0].message.content

        # Calculate the total token count of our conversation messages via tiktoken
        token_count = self.conversations[uname].get_message_token_count()

        # If the user is getting closer to the bot's token threshold, warn the first time
        if token_count > self.chatgpt['prune_back_to'] and uname not in self.token_warning:
            reply += ("\n\n"
                "By the way, our conversation will soon reach my token limit, so I may"
                " start forgetting some of our older exchanges. Would you like me to"
                " summarize our conversation so far to keep the main points alive?"
            )
            self.token_warning[uname] = True

        # Add assistant's message to the user's conversation
        self.conversations[uname].add_assistant_message(reply)

        # Truncate older messages if the conversation is above the bot's token threshold
        if token_count > self.chatgpt['prune_threshold']:
            self.conversations[uname].prune_conversation(self.chatgpt['prune_back_to'])

        return reply

    async def tele_handle_message(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
        """Handles the Telegram side of the message, discerning between Private and Group conversation."""
        validated = await self.tele_validate(update)
        if not validated:
            return
        (msg, chat, user) = validated

        # If it's a group text, only reply if the bot is named
        # The real magic of how the bot behaves is in tele_handle_response()
        response = "Sorry, I couldn't process your message! Please contact my creator."
        if chat.type == 'supergroup' or chat.type == 'group':
            if exact_word_match(self.telegram['username'], msg.text):
                pattern = r'@?\b' + re.escape(self.telegram['username']) + r'\b'
                new_text = re.sub(pattern, '', msg.text).strip()
                response = await self.tele_handle_response(new_text, msg)
            elif (
                exact_word_match(self.telegram['nickname'], msg.text) or
                exact_word_match(self.telegram['initials'], msg.text)
            ):
                response = await self.tele_handle_response(msg.text, msg)
            else:
                return
        elif chat.type == 'private':
            response = await self.tele_handle_response(msg.text, msg)
        else:
            return

        # Split into smaller chunks since Telegram messages have a maximum text length (likely 4096)
        chunk_length = MessageLimit.MAX_TEXT_LENGTH - 1
        for chunk in [response[i:i+chunk_length] for i in range(0, len(response), chunk_length)]:
            await msg.reply_text(chunk)
            await asyncio.sleep(0.12)  # small pause to reduce risk of rate limiting

    async def tele_validate(self, update: Update) -> tuple[Message, Chat, User] | None:
        """
        Check if any update through Telegram has a complete message, chat, and user.
        This uses python-telegram-bot's Update `effective` methods for all three.
        """
        if update is None:
            return None

        # Any python-telegram-bot 'effective' method helps ensure there are no "None" types
        # Non-user updates like Inline queries or chat member updates makes either one undefined
        msg  = update.effective_message
        chat = update.effective_chat
        user = update.effective_user

        # If there is no chat or valid message of the unique conversation, no reason to reply
        if chat is None or msg is None:
            return None
        # chat.type => 'private', 'group', 'supergroup', 'channel'
        # chat.id = user.id if 'private'

        # Handle the user sending the message
        if user:
            # This ensures conversations don't overlap (for now, until user.id is implemented)
            uname = getattr(user, "username", None)
            if uname is None:
                await msg.reply_text('Add a username to your Telegram account so that I can talk to you!')
                return None
        else:
            # If not the user chatting directly to the bot, don't bother
            return None

        return (msg, chat, user)

    async def tele_error(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
        """Handle any conversation error caused on the Telegram side."""
        log_error(context.error, error_type='Telegram', error_filename=self.error_log)
        msg = update.effective_message if update else None
        if msg:
            await msg.reply_text("Sorry, I ran into an error! Please contact my creator.")

    async def gpt_completion(self, uname: str):
        """Get the OpenAI response based on bot configuration."""
        try:
            client = OpenAIClientSingleton.get_instance()
            response = await client.chat.completions.create(
                model=self.chatgpt['chat_model'],
                messages=self.conversations[uname].messages,
            )
            return response
        except openai.AuthenticationError as e:
            # Handle authentication error
            log_error(e, error_type='OpenAI-Authentication', error_filename=self.error_log)
        except openai.BadRequestError as e:
            # Handle invalid request error
            if re.search(r'maximum context.+reduce the length', str(e)):
                # Remove older messages and try again since the model's maximum token limit reached
                print(f"Response to {uname} reached maximum context length, pruning conversation")
                self.conversations[uname].prune_conversation(self.chatgpt['prune_back_to'])
                return await self.gpt_completion(uname)
            else:
                # Another error is actually invalid to investigate
                log_error(e, error_type='OpenAI-BadRequest', error_filename=self.error_log)
        except openai.APIConnectionError as e:
            # Handle API connection error
            log_error(e, error_type='OpenAI-APIConnection', error_filename=self.error_log)
        except Exception as e:
            # Catch any other unexpected exceptions
            log_error(e, error_type='Other', error_filename=self.error_log)

    def start_polling(self):
        """The main polling "loop" the user interacts with via Telegram."""
        print(f"TeLLMgramBot {self.telegram['username']} polling...")
        self.telegram['app'].run_polling(poll_interval=self.telegram['pollinterval'])
        print(f"TeLLMgramBot {self.telegram['username']} polling ended.")

    # Initialization
    def __init__(self,
        bot_owner      = INIT_BOT_CONFIG['bot_owner'],
        bot_nickname   = INIT_BOT_CONFIG['bot_nickname'],
        bot_initials   = INIT_BOT_CONFIG['bot_initials'],
        chat_model     = INIT_BOT_CONFIG['chat_model'],
        url_model      = INIT_BOT_CONFIG['url_model'],
        token_limit    = INIT_BOT_CONFIG['token_limit'],
        persona_temp   = INIT_BOT_CONFIG['persona_temp'],
        persona_prompt = INIT_BOT_CONFIG['persona_prompt']
    ):
        # Starting to initialize, not online yet
        self._online = False

        # First provide the main structure if not already there
        init_structure()

        # Initialize some variables
        self.error_log = generate_error_path()
        self.token_warning = {}  # Determines whether user has reached token limit by AI model
        self.conversations = {}  # Provides Conversation class per user based on bot response
        self.telegram = {
            'owner'        : bot_owner,
            'nickname'     : bot_nickname,
            'initials'     : bot_initials,
            'pollinterval' : 3
        }
        # Check for event running loops before getting the bot's information
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            # No running event loop; safe to run synchronously
            asyncio.run(self._tele_info())
        else:
            # Already in an event loop; schedule as a background task
            loop.create_task(self._tele_info())

        # Build our application with handlers for Commands, Messages, and Errors
        self.telegram['app'] = Application.builder().token(os.environ['TELLMGRAMBOT_TELEGRAM_API_KEY']).build()
        self.telegram['app'].add_handler(CommandHandler('help', self.tele_commands))
        self.telegram['app'].add_handler(CommandHandler('start', self.tele_start_command))
        self.telegram['app'].add_handler(CommandHandler('stop', self.tele_stop_command))
        self.telegram['app'].add_handler(CommandHandler('nick', self.tele_nick_command))
        self.telegram['app'].add_handler(CommandHandler('forget', self.tele_forget_command))
        self.telegram['app'].add_handler(MessageHandler(filters.TEXT, self.tele_handle_message))
        self.telegram['app'].add_error_handler(self.tele_error)

        # Get our LLM spun up with defaults if not defined by user input
        # Tokens as integers measure the length of conversation messages
        self.chatgpt = {
            'prompt'      : persona_prompt,
            'chat_model'  : chat_model,
            'url_model'   : url_model,
            'token_limit' : token_limit or TokenGPT(chat_model).max_tokens(),
            'temperature' : persona_temp or 1.0,
            'top_p'       : 0.9
        }
        # Set a rounded-down integer to prune a lengthy conversation by 500 tokens
        # Note if the upper limit is below 500, the lower limit is set to 0
        self.chatgpt['prune_threshold'] = floor(0.95 * self.chatgpt['token_limit'])
        self.chatgpt['prune_back_to'] = max(0, self.chatgpt['prune_threshold'] - 500)

        # Bot is now ready and active by default
        self._online = True 

    def set(config_file='config.yaml', prompt_file='test_personality.prmpt'):
        """Set TeLLMgramBot object based on its YAML configuration and prompt files."""
        # First provide the main structure if not already there
        init_structure()

        # Ensure both bot configuration and prompt files are defined and readable
        config = read_yaml(init_bot_config(config_file))
        prompt = read_text(init_bot_prompt(prompt_file))

        # Check any configuration values missing and apply default values:
        for parameter, value in INIT_BOT_CONFIG.items():
            if parameter == 'persona_prompt':
                # Apply initial prompt if not defined
                if not prompt:
                    prompt = value
                    print(f"File '{prompt_file}' is empty, set default prompt '{prompt}'")
            elif parameter not in config:
                # Apply initial configuration parameter with default if not defined
                config[parameter] = value
                if value:
                    print(f"Configuration '{parameter}' not defined, set to '{value}'")

        # Apply parameters to bot:
        return TelegramBot(
            bot_owner      = config['bot_owner'],
            bot_nickname   = config['bot_nickname'],
            bot_initials   = config['bot_initials'],
            chat_model     = config['chat_model'],
            url_model      = config['url_model'],
            token_limit    = config['token_limit'],
            persona_temp   = config['persona_temp'],
            persona_prompt = prompt
        )
