import os
import logging
import sys
import asyncio
import time
from typing import TypedDict, Union, TypeAlias, List # Added List for type hint
from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup
from telegram.ext import Application, CommandHandler, MessageHandler, filters, ContextTypes, CallbackQueryHandler
from browse_command import browse_command, button_callback

class MessageHandlerLogicResult(TypedDict):
    success: bool
    response_text: Union[str, None]
    error_message: Union[str, None]

LogicResult: TypeAlias = MessageHandlerLogicResult

class TelegramHelper:
    CLAUDE_REBOOT_TARGET = 'claude'
    HTML_QUOTE_BLOCK_START = '<blockquote expandable><b>Thinking...</b>'
    HTML_QUOTE_BLOCK_END = '</blockquote>'
    DEFAULT_REBOOT_CLAUDE_FILE = '.reboot_claude'
    DEFAULT_REBOOT_FILE = '.doreboot'
    CHUNK_MESSAGE_SLEEP_DURATION = 0.1

    def __init__(self, bot,
                 reboot_claude_file_path: str | None = None,
                 reboot_file_path: str | None = None,
                 chunk_message_sleep_duration: float | None = None):
        self.bot = bot
        self.telegram_bot_token = os.getenv('TELEGRAM_BOT_TOKEN')
        self.start_time = time.time()
        self.reboot_claude_file = reboot_claude_file_path or self.DEFAULT_REBOOT_CLAUDE_FILE
        self.reboot_file = reboot_file_path or self.DEFAULT_REBOOT_FILE
        self.chunk_message_sleep_duration = chunk_message_sleep_duration if chunk_message_sleep_duration is not None else self.CHUNK_MESSAGE_SLEEP_DURATION

    async def _start_logic(self) -> str:
        await self.bot.start()
        return "Hello! I'm your AI assistant. How can I help you today?"

    async def start(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
        response_message = await self._start_logic()
        await update.message.reply_text(response_message)

    async def _clear_logic(self, user_id: int) -> str:
        self.bot.clear_conversation_history(user_id)
        return "Conversation history cleared. Let's start fresh!"

    async def clear(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
        user_id = update.effective_user.id
        response_message = await self._clear_logic(user_id)
        await update.message.reply_text(response_message)

    async def _status_logic(self) -> str:
        return await self.bot.get_bot_status()

    async def status(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
        response_message = await self._status_logic()
        await update.message.reply_text(response_message)

    async def _switch_logic(self) -> str:
        if hasattr(self.bot, 'switch_model'):
            return await self.bot.switch_model()
        else:
            return "Model switching is not supported for this bot."

    async def switch(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
        response_message = await self._switch_logic()
        await update.message.reply_text(response_message)

    async def update_status_message(self, context: ContextTypes.DEFAULT_TYPE, chat_id: int, message_id: int, status: str):
        keyboard = [
            [InlineKeyboardButton("Abort", callback_data='abort')]
        ]
        reply_markup = InlineKeyboardMarkup(keyboard)
        await context.bot.edit_message_text(
            chat_id=chat_id,
            message_id=message_id,
            text=f"Current status: {status}",
            reply_markup=reply_markup
        )

    async def _handle_message_logic(self, user_id: int, user_message: str) -> LogicResult:
        try:
            response = await self.bot.handle_message(user_id, user_message)
            processed_response = response.replace("<think>", self.HTML_QUOTE_BLOCK_START).replace("</think>", self.HTML_QUOTE_BLOCK_END)
            return LogicResult(success=True, response_text=processed_response, error_message=None)
        except Exception as e:
            logging.error(f"Error in _handle_message_logic for user {user_id}: {str(e)}")
            return LogicResult(success=False, response_text=None, error_message=str(e))

    async def handle_message(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
        user_id = update.effective_user.id
        user_message = update.message.text
        chat_id = update.effective_chat.id
        status_message_obj = None

        try:
            status_message_obj = await update.message.reply_text(
                "Processing your request...",
                reply_markup=InlineKeyboardMarkup([[InlineKeyboardButton("Abort", callback_data='abort')]])
            )
            self.bot.set_processing_status(user_id, status_message_obj.message_id)

            logic_result = await self._handle_message_logic(user_id, user_message)

            if status_message_obj:
                try:
                    await context.bot.delete_message(chat_id=chat_id, message_id=status_message_obj.message_id)
                except Exception as e_del:
                    logging.warning(f"Failed to delete status message: {e_del}")
            self.bot.clear_processing_status(user_id)

            if logic_result["success"]:
                response_text = logic_result["response_text"]
                if response_text:
                    if len(response_text) > 4096:
                        chunks = [response_text[i:i + 4096] for i in range(0, len(response_text), 4096)]
                        for chunk in chunks:
                            await update.message.reply_text(chunk)
                            await asyncio.sleep(self.chunk_message_sleep_duration)
                    else:
                        await update.message.reply_text(response_text)
                else:
                    logging.warning("Successful logic result but no response text.")
                    await update.message.reply_text("Something went unexpectedly well, but I have nothing to say.")
            else:
                await update.message.reply_text("Sorry, an error occurred while processing your request.")

        except Exception as e:
            logging.error(f"Outer error in handle_message for user {user_id}: {str(e)}")
            if status_message_obj and self.bot.processing_status.get(user_id):
                self.bot.clear_processing_status(user_id)
            try:
                await update.message.reply_text("Sorry, an unexpected error occurred with the bot.")
            except Exception as e_reply:
                logging.error(f"Failed to send error reply: {e_reply}")

    async def _abort_processing_logic(self, user_id: int) -> str:
        return await self.bot.abort_processing(user_id)

    async def abort_processing(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
        query = update.callback_query
        await query.answer()
        user_id = query.from_user.id
        response_text = await self._abort_processing_logic(user_id)
        await query.edit_message_text(text=response_text)

    # --- Reboot Command --- 
    def _reboot_logic(self, user_message_parts: List[str], chat_id_to_write: str) -> None:
        """Handles the logic for creating reboot files."""
        if len(user_message_parts) > 1 and user_message_parts[1].lower() == self.CLAUDE_REBOOT_TARGET:
            try:
                with open(self.reboot_claude_file, 'w') as f:
                    f.write("") # Create/truncate the file
                logging.info(f"Created/truncated Claude reboot file: {self.reboot_claude_file}")
            except IOError as e:
                logging.error(f"Failed to create/truncate Claude reboot file {self.reboot_claude_file}: {e}")

        # Create the main reboot file if it doesn't exist
        if not os.path.exists(self.reboot_file):
            try:
                with open(self.reboot_file, 'w') as f:
                    f.write(chat_id_to_write)
                logging.info(f"Created main reboot file: {self.reboot_file} with chat_id.")
            except IOError as e:
                logging.error(f"Failed to create main reboot file {self.reboot_file}: {e}")
        else:
            logging.info(f"Main reboot file {self.reboot_file} already exists. Not overwriting chat_id.")

    async def reboot(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
        """Handles the /reboot command, triggers file creation and exits."""
        user_message_parts = update.message.text.split()
        chat_id_str = str(update.effective_chat.id) if update and update.effective_chat else ""
        
        self._reboot_logic(user_message_parts, chat_id_str)

        if update:
            try:
                await update.message.reply_text("Rebooting the bot...")
            except Exception as e_reply:
                logging.error(f"Failed to send reboot reply: {e_reply}")
        
        logging.info("Initiating shutdown for reboot...")
        sys.exit(0) # This part is not directly testable for completion in unit tests

    # --- Check Doreboot File --- 
    async def _check_doreboot_file_logic(self) -> Union[str, None]:
        """Checks for the reboot file, reads chat_id, removes file, and returns chat_id."""
        if os.path.exists(self.reboot_file):
            chat_id = None
            try:
                with open(self.reboot_file, 'r') as f:
                    chat_id = f.read().strip()
                # Attempt to remove the file after reading
                try:
                    os.remove(self.reboot_file)
                    logging.info(f"Successfully read and removed reboot file: {self.reboot_file}")
                except OSError as e_remove:
                    logging.error(f"Failed to remove reboot file {self.reboot_file} after reading: {e_remove}")
                    # Still return chat_id if read was successful, to attempt notification
                return chat_id
            except IOError as e_read:
                logging.error(f"Error reading reboot file {self.reboot_file}: {e_read}")
                # If reading failed, attempt to remove anyway if it exists, to prevent stale files
                if os.path.exists(self.reboot_file):
                    try:
                        os.remove(self.reboot_file)
                        logging.warning(f"Removed reboot file {self.reboot_file} after a read error.")
                    except OSError as e_remove_after_fail:
                        logging.error(f"Failed to remove reboot file {self.reboot_file} even after a read error: {e_remove_after_fail}")
                return None # Reading failed
        return None # File does not exist

    async def check_doreboot_file(self, application: Application) -> None:
        """Checks for reboot file using logic method and sends notification if applicable."""
        chat_id = await self._check_doreboot_file_logic()
        if chat_id:
            try:
                await application.bot.send_message(chat_id=chat_id, text="The application has finished initializing.")
                logging.info(f"Sent reboot initialization notification to chat_id: {chat_id}")
            except Exception as e:
                logging.error(f"Failed to send reboot initialization notification to chat_id {chat_id}: {e}")

    async def browse(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
        await browse_command(update, context, self.bot)

    def run(self):
        application = Application.builder().token(self.telegram_bot_token).build()

        application.add_handler(CommandHandler("start", self.start))
        application.add_handler(CommandHandler("clear", self.clear))
        application.add_handler(CommandHandler("switch", self.switch))
        application.add_handler(CommandHandler("status", self.status))
        application.add_handler(CommandHandler("reboot", self.reboot))
        application.add_handler(CommandHandler("browse", self.browse))
        application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, self.handle_message))
        application.add_handler(CallbackQueryHandler(self.abort_processing, pattern='^abort$'))
        application.add_handler(CallbackQueryHandler(button_callback, pattern='^(browse|file):'))

        logging.info("Bot is running...")

        loop = asyncio.get_event_loop()
        if loop.is_running(): # pragma: no cover
            loop.create_task(self.check_doreboot_file(application))
        else: # pragma: no cover
            asyncio.run(self.check_doreboot_file(application))

        application.run_polling()