#!/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
from .providers.factory import get_provider
from .providers.base import ContextLengthExceededError, ProviderAuthError, ProviderConnectionError
from .initialize import INIT_BOT_CONFIG, ApiKeyStatus, init_structure, init_bot_config, init_bot_prompt
from .conversation import Conversation
from .database import get_private_mode, set_private_mode, delete_messages_for_user, delete_messages_for_chat
from .models import TokenLimits
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 a Large Language Model (LLM).
    - 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['bot_id'] = bot.id
            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"
            "/private - Toggle private mode (private chats only).\n"
            "/help - Display this usage information.\n\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 conversation history from the database (behavior depends on chat type).

        In private chats: deletes all messages in the chat (including bot replies), fully clearing the conversation.
        In group chats: deletes only this user's rows; other participants' messages remain.
        Also removes the chat's Conversation object from this bot session's memory if it exists.
        """
        user_id = update.message.from_user.id
        chat_id = update.message.chat.id
        chat_type = update.message.chat.type
        # Private chats: delete all messages in the chat (including bot replies)
        # Group chats: delete only the user's rows (other participants' messages remain)
        if chat_type == 'private':
            await delete_messages_for_chat(chat_id)
        else:
            await delete_messages_for_user(user_id)
        if chat_id in self.conversations:
            del self.conversations[chat_id]
        await update.message.reply_text("My memories of our conversations are wiped!")

    async def tele_private_command(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
        """
        Toggle private mode for the requesting user (private chats only).

        When private mode is ON, messages sent in private chats are flagged as is_private=1
        in the database and excluded when building context for group conversations. This lets
        users keep sensitive private exchanges out of shared group contexts.

        Usage:
            /private       — toggles private mode on (if off) or off (if on)
            /private off   — explicitly disables private mode

        Only available in private chats; sending the command in a group returns an error.
        """
        validated = await self.tele_validate(update)
        if not validated:
            return
        (msg, chat, user) = validated

        if chat.type != 'private':
            await msg.reply_text("Private mode is only available in private chats.")
            return

        user_id = user.id
        args = (msg.text or "").strip().split()
        force_off = len(args) > 1 and args[1].lower() == 'off'
        current = await get_private_mode(user_id)
        new_state = False if (force_off or current) else True
        await set_private_mode(user_id, new_state)

        if new_state:
            await msg.reply_text(
                "Private mode is ON. Future messages in this private chat will be excluded "
                "from group conversation contexts. Use /private off to disable."
            )
        else:
            await msg.reply_text(
                "Private mode is OFF. Future messages in this private chat will be included in group conversation contexts."
            )

    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.

        Conversations are keyed by chat_id (Telegram chat ID) and track both user_id
        and bot_id to support multi-participant group conversations. Respects per-user
        private mode settings when loading past context.

        Feature gating: If chat_enabled or url_analysis_enabled are False (per key_status),
        those features return a generic "couldn't process" error rather than attempting
        to call disabled LLM providers.

        Args:
            text: The user's message text.
            msg: The Telegram Message object containing chat/user context.

        Returns:
            The bot's response string. If no handler matches or features are disabled,
            returns a generic error message.

        Side Effects:
            - Adds user message to conversation history (keyed by chat_id).
            - Stores user_id, username (display only), and private mode flag.
            - Loads past interactions if this is the first message in this chat_id.
            - Generates a token warning if conversation nears the limit.
            - Prunes conversation if token count exceeds the threshold.
            - Stores bot's response with bot_id and bot_username for tracking.
        """
        # 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!"

        # Extract identity and context from the message
        # Username is now optional and only for display
        user_id = msg.from_user.id
        username = msg.from_user.username
        chat_id = msg.chat.id
        chat_type = msg.chat.type
        print(f"User {user_id} ({username}) in {chat_type} Chat {chat_id}")

        # For a new session, create a Conversation keyed by chat_id and load past history
        if chat_id not in self.conversations:
            self.conversations[chat_id] = Conversation(
                chat_id,
                chat_type,
                self.chatgpt['prompt'],
                self.chatgpt['chat_model']
            )
            await self.conversations[chat_id].get_past_interaction(
                floor(self.chatgpt['prune_threshold'] / 2), user_id
            )

        # Add the user's message to our conversation, respecting private mode
        # Private mode only applies in private chats — group messages are never flagged private
        is_private = (chat_type == 'private') and await get_private_mode(user_id)
        await self.conversations[chat_id].add_user_message(text, user_id, username, is_private)

        # 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 the LLM
        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 and self.key_status.url_analysis_enabled:
            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 and self.key_status.chat_enabled:
            # This is the transition point between quick Telegram replies and the LLM
            reply = await self.llm_completion(chat_id)

        # Check token count before storing to determine if warning should be appended
        token_count = await self.conversations[chat_id].get_message_token_count()
        if token_count > self.chatgpt['prune_back_to'] and chat_id 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[chat_id] = True

        # Add the full reply (including any warning) to the conversation
        await self.conversations[chat_id].add_assistant_message(
            reply, self.telegram['bot_id'], self.telegram['username']
        )

        # Recompute token count after full exchange is stored, then prune if needed
        token_count = await self.conversations[chat_id].get_message_token_count()
        if token_count > self.chatgpt['prune_threshold']:
            await self.conversations[chat_id].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.
        Requires a valid message and chat; user.id must be present (username is optional).
        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; user.id is always present
        if not user:
            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 llm_completion(self, chat_id: int, _retry: bool = True) -> str:
        """
        Get the LLM response for the conversation in the given chat.

        Args:
            chat_id: Telegram chat ID (negative for groups, positive for private).
            _retry: Internal flag to prevent infinite retry loops on context overflow.

        Returns:
            The LLM's response string, or a user-facing error message on failure.

        Selects the provider (OpenAI or Anthropic) automatically based on the configured model.
        On context length errors, prunes the conversation and retries once. Other errors are
        logged and return a user-facing fallback string.
        """
        provider = get_provider(self.chatgpt['chat_model'])
        try:
            return await provider.complete(self.chatgpt['chat_model'], self.conversations[chat_id].messages)
        except ContextLengthExceededError:
            if not _retry:
                return "Sorry, your conversation is too long for me to respond to. Try /forget to start fresh."
            print(f"Chat {chat_id} reached maximum context length, pruning conversation")
            await self.conversations[chat_id].prune_conversation(self.chatgpt['prune_back_to'])
            return await self.llm_completion(chat_id, _retry=False)
        except ProviderAuthError as e:
            log_error(e, error_type='Authentication', error_filename=self.error_log)
            return "Sorry, I'm having trouble authenticating with the AI provider. Please contact the bot owner."
        except ProviderConnectionError as e:
            log_error(e, error_type='Connection', error_filename=self.error_log)
            return "Sorry, I couldn't reach the AI provider. Please try again in a moment."
        except Exception as e:
            log_error(e, error_type='Other', error_filename=self.error_log)
            return "Sorry, something went wrong on my end. Please try again."

    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'],
        key_status: ApiKeyStatus | None = None
    ):
        """
        Initialize the Telegram bot with LLM configuration and API keys.

        Args:
            bot_owner: Telegram username of the bot owner (required, no `@`).
            bot_nickname: Nickname the bot responds to in group chats.
            bot_initials: Initials the bot responds to in group chats.
            chat_model: LLM model for conversation (e.g., 'gpt-5-mini', 'claude-sonnet-4-6').
            url_model: LLM model for URL analysis (can differ from chat_model).
            token_limit: Maximum tokens for conversations. If None, uses the chat_model default.
            persona_temp: LLM temperature (0.0-2.0). If None, defaults to 1.0.
            persona_prompt: System prompt defining the bot's behavior and personality.
            key_status: ApiKeyStatus object indicating available features. If None, calls init_structure().

        Side Effects:
            - Fetches bot metadata (username, full_name, bot_id) from Telegram API via _tele_info().
            - Initializes command and message handlers for the Telegram application.
            - Sets self._online to True once initialization completes (ready for polling).
        """
        # Starting to initialize, not online yet
        self._online = False

        # First provide the main structure if not already there; capture API key status
        self.key_status = key_status if key_status is not None else 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 = {
            'bot_id'       : 0,  # overwritten by _tele_info(); 0 is a safe sentinel
            '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(CommandHandler('private', self.tele_private_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 TokenLimits(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'):
        """
        Instantiate a TelegramBot from configuration and prompt files.

        Reads bot settings from the specified YAML configuration file and prompt file,
        applies defaults for any missing values, and returns a fully initialized TelegramBot.
        Also calls init_structure() to bootstrap directories and API keys, and returns
        the resulting ApiKeyStatus to gate features.

        Args:
            config_file: Path to bot configuration YAML file (default: 'config.yaml').
            prompt_file: Path to bot persona prompt file (default: 'test_personality.prmpt').

        Returns:
            A fully initialized TelegramBot instance.

        Side Effects:
            - Calls init_structure() which may create directories and check API keys.
            - May print warning messages if configuration values are missing or prompt file is empty.
        """
        # First provide the main structure if not already there; capture API key status
        key_status = 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,
            key_status    = key_status
        )
