#chat.py
"""
Handles the chat widget
"""

import gi
from gi.repository import Gtk, Gio, Adw, Gdk, GLib
import logging, os, datetime, random, json, threading, re, importlib.util
from ..constants import SAMPLE_PROMPTS, cache_dir
from ..sql_manager import generate_uuid, prettify_model_name, generate_numbered_name, Instance as SQL
from . import dialog, voice, models, blocks
from .message import Message

logger = logging.getLogger(__name__)


@Gtk.Template(resource_path='/com/jeffser/Alpaca/widgets/chat/folder.ui')
class Folder(Adw.NavigationPage):
    __gtype_name__ = 'AlpacaFolder'

    bottom_bar = Gtk.Template.Child()
    list_stack = Gtk.Template.Child()
    scrolled_window = Gtk.Template.Child()
    folder_list_box = Gtk.Template.Child()
    separator = Gtk.Template.Child()
    chat_list_box = Gtk.Template.Child()
    top_indicator = Gtk.Template.Child()
    bottom_indicator = Gtk.Template.Child()

    def __init__(self, folder_id:str=None, folder_name:str=_('Root'), folder_color:str=None, show_bar:bool=True):
        self.folder_id = folder_id
        self._scroll_timeout_id = None
        super().__init__()
        indicator_top_drop_target = Gtk.DropTarget.new(Gtk.ListBoxRow, Gdk.DragAction.COPY)
        indicator_top_drop_target.connect("accept", lambda *_: self.start_scrolling(-1))
        indicator_top_drop_target.connect("leave", lambda *_: self.stop_scrolling())
        self.top_indicator.add_controller(indicator_top_drop_target)

        indicator_bottom_drop_target = Gtk.DropTarget.new(Gtk.ListBoxRow, Gdk.DragAction.COPY)
        indicator_bottom_drop_target.connect("accept", lambda *_: self.start_scrolling(1))
        indicator_bottom_drop_target.connect("leave", lambda *_: self.stop_scrolling())
        self.bottom_indicator.add_controller(indicator_bottom_drop_target)

        drop_target_folder = Gtk.DropTarget.new(FolderRow, Gdk.DragAction.MOVE)
        drop_target_folder.connect("drop", self.on_drop_folder)
        self.bottom_bar.add_controller(drop_target_folder)
        drop_target_chat = Gtk.DropTarget.new(ChatRow, Gdk.DragAction.MOVE)
        drop_target_chat.connect("drop", self.on_drop_chat)
        self.bottom_bar.add_controller(drop_target_chat)

        self.get_child().set_reveal_bottom_bars(show_bar)
        self.set_title(folder_name)
        if folder_color:
            self.add_css_class('folder-{}'.format(folder_color))

    def start_scrolling(self, direction):
        if self._scroll_timeout_id:
            GLib.source_remove(self._scroll_timeout_id)
        self._scroll_timeout_id = GLib.timeout_add(30, self.do_scroll, direction)

    def stop_scrolling(self):
        if self._scroll_timeout_id:
            GLib.source_remove(self._scroll_timeout_id)
            self._scroll_timeout_id = None

    def do_scroll(self, direction):
        adj = self.scrolled_window.get_vadjustment()
        new_value = adj.get_value() + direction * 15 # scroll speed
        adj.set_value(max(0, min(new_value, adj.get_upper() - adj.get_page_size())))

        self.top_indicator.set_visible(new_value >= 0) # Hide the top indicator so it doesn't get in the way of the top folder
        return True

    def on_drop_folder(self, target, row, x, y):
        folder_page = self.get_root().chat_list_navigationview.get_previous_page(self)
        if row.folder_id != folder_page.folder_id:
            SQL.move_folder_to_folder(row.folder_id, folder_page.folder_id)
            row.get_parent().remove(row)
            folder_page.folder_list_box.prepend(row)
            row.set_visible(True)
            self.update_visibility()
            return True

    def on_drop_chat(self, target, row, x, y):
        folder_page = self.get_root().chat_list_navigationview.get_previous_page(self)
        row.chat.folder_id = folder_page.folder_id
        SQL.insert_or_update_chat(row.chat)
        row.get_parent().remove(row)
        folder_page.chat_list_box.prepend(row)
        row.set_visible(True)
        self.update_visibility()
        return True

    def update_visibility(self, searching:bool=False):
        folder_visible = False
        chat_visible = False

        for row in list(self.folder_list_box):
            folder_visible = folder_visible or row.get_visible()
            if folder_visible:
                break

        for row in list(self.chat_list_box):
            chat_visible = chat_visible or row.get_visible()
            if chat_visible:
                break

        self.folder_list_box.set_visible(folder_visible)
        self.chat_list_box.set_visible(chat_visible)
        self.separator.set_visible(folder_visible and chat_visible)

        if searching:
            self.list_stack.set_visible_child_name('content' if folder_visible or chat_visible else 'no-results')
        else:
            self.list_stack.set_visible_child_name('content' if folder_visible or chat_visible else 'empty')

    def on_search(self, raw_query:str):
        query = re.escape(raw_query)
        root = self.get_root()
        include_messages = root.settings.get_value('folder-search-mode').unpack()

        if include_messages:
            root.searchentry_messages.set_text(raw_query)

        if len(list(self.folder_list_box)) + len(list(self.chat_list_box)) == 0:
            self.list_stack.set_visible_child_name('empty')
            return

        for row in list(self.folder_list_box):
            row.set_visible(re.search(query, row.get_name(), re.IGNORECASE))

            if row.get_visible():
                if query:
                    query_escaped = GLib.markup_escape_text(query)
                    row_label_text = GLib.markup_escape_text(row.get_name())
                    highlighted_text = re.sub(f"({query_escaped})", r"<span background='yellow' bgalpha='30%'>\1</span>", row_label_text, flags=re.IGNORECASE)
                    row.label.set_markup(highlighted_text)
                else:
                    row.label.set_markup(row.get_name())

        for row in list(self.chat_list_box):
            title_match = re.search(query, row.get_name(), re.IGNORECASE)

            messages_str = ""
            message_match = False

            if include_messages:
                messages_str = '\n'.join([m.get('content') for m in row.chat.convert_to_ollama()]) or '\n'.join([m[4] for m in SQL.get_messages(row.chat)])
                message_match = re.search(query, messages_str, re.IGNORECASE)

            row.set_visible(title_match or message_match)
            if row.get_visible():
                if query:
                    query_escaped = GLib.markup_escape_text(query)
                    row_label_text = GLib.markup_escape_text(row.get_name())
                    highlighted_text = re.sub(f"({query_escaped})", r"<span background='yellow' bgalpha='30%'>\1</span>", row_label_text, flags=re.IGNORECASE)
                    row.label.set_markup(highlighted_text)
                else:
                    row.label.set_markup(row.get_name())

        self.update_visibility(True)

    def update(self):
        self.chat_list_box.remove_all()
        self.folder_list_box.remove_all()
        selected_chat = self.get_root().settings.get_value('default-chat').unpack()
        chats = SQL.get_chats_by_folder(self.folder_id)
        if len(chats) > 0:
            if selected_chat not in [row[0] for row in chats] and not self.folder_id:
                selected_chat = chats[0][0]

            for row in chats:
                self.add_chat(
                    chat_name=row[1],
                    chat_id=row[0],
                    is_template=row[2]==1,
                    mode=0
                )
                if row[0] == selected_chat and len(list(self.chat_list_box)) > 0:
                    self.chat_list_box.select_row(list(self.chat_list_box)[-1])

        if len(list(self.chat_list_box)) == 0 and not self.folder_id:
            self.chat_list_box.select_row(self.new_chat().row)

        if not self.chat_list_box.get_selected_row() and not self.folder_id:
            self.chat_list_box.select_row(list(self.chat_list_box)[0])

        folders = SQL.get_chat_folders(self.folder_id)
        for f in folders:
            row = FolderRow(f[0], f[1], f[2], f[3])
            self.folder_list_box.append(row)

        self.update_visibility()

    def add_chat(self, chat_name:str, chat_id:str, is_template:bool, mode:int): #mode = 0: append, mode = 1: prepend
        chat_name = chat_name.strip()
        if chat_name and mode in (0, 1):
            chat_name = generate_numbered_name(chat_name, [row.get_name() for row in list(self.chat_list_box)])
            chat = Chat(
                chat_id=chat_id,
                name=chat_name,
                folder_id=self.folder_id,
                is_template=is_template
            )

            if chat:
                if mode == 0:
                    self.chat_list_box.append(chat.row)
                else:
                    self.chat_list_box.prepend(chat.row)
                self.update_visibility()
                return chat

    def new_chat(self, chat_name:str=_('New Chat')):
        if not chat_name.strip():
            chat_name = _('New Chat')
        chat = self.add_chat(
            chat_name=chat_name,
            chat_id=generate_uuid(),
            is_template=False,
            mode=1
        )
        self.chat_list_box.select_row(chat.row)
        if chat:
            SQL.insert_or_update_chat(chat)
            return chat

    def new_folder(self, name:str, color:str):
        name = generate_numbered_name(name, [row.get_name() for row in list(self.folder_list_box)])
        row = FolderRow(
            generate_uuid(),
            name,
            color,
            self.folder_id
        )
        SQL.insert_or_update_folder(
            row.folder_id,
            row.folder_name,
            row.folder_color,
            row.folder_parent
        )
        self.folder_list_box.prepend(row)
        self.update_visibility()

    def prompt_new_folder(self):
        options = {
            _('Cancel'): {},
            _('Accept'): {
                'appearance': 'suggested',
                'callback': lambda name, toggle_group: self.new_folder(name, toggle_group.get_active_name()),
                'default': True
            }
        }

        d = dialog.Entry(
            _('New Folder'),
            '',
            list(options.keys())[0],
            options,
            {'placeholder': _('New Folder'), 'text': _('New Folder')}
        )
        color_group = Adw.ToggleGroup()
        color_names = ('blue', 'teal', 'green', 'yellow', 'orange', 'red', 'pink', 'purple', 'slate')
        for c in color_names:
            icon = Gtk.Image.new_from_icon_name('big-dot-symbolic')
            icon.add_css_class('button-{}'.format(c))
            icon.set_icon_size(2)
            toggle = Adw.Toggle(
                name=c,
                child=icon
            )
            color_group.add(toggle)

        color_group.set_active_name(random.choice(color_names))
        d.container.append(color_group)

        d.show(self.get_root())

    @Gtk.Template.Callback()
    def chat_changed(self, listbox, row):
        if not listbox.get_root() or (row and not row.get_root()):
            return

        if row and row.get_root():
            row.get_root().chat_page.set_title(row.get_name())

        last_chat_id = -1
        if listbox.get_root().chat_bin.get_child():
            last_chat_id = listbox.get_root().chat_bin.get_child().chat_id

        if row and row.chat.chat_id != last_chat_id:
            if listbox.get_root().chat_bin.get_child():
                list_box = listbox.get_root().chat_bin.get_child().row.get_parent()
                if list_box and list_box != self.chat_list_box:
                    listbox.get_root().chat_bin.get_child().row.get_parent().unselect_all()

            # Load New Chat
            new_chat = row.chat
            if new_chat.get_parent():
                new_chat.get_parent().set_child(None)
            if new_chat.busy:
                self.get_root().global_footer.toggle_action_button(False)
            else:
                self.get_root().global_footer.toggle_action_button(True)

            if len(list(new_chat.container)) == 0:
                new_chat.load_messages()

            # Show New Stack Page
            root = self.get_root()
            root.chat_bin.set_child(new_chat)
            chat_search_query = root.searchentry_messages.get_text()
            if chat_search_query:
                GLib.idle_add(new_chat.on_search, chat_search_query)

            # Select Model
            GLib.idle_add(self.auto_select_model)

    def auto_select_model(self):
        def find_model_index(model_name:str) -> int:
            text_models = models.text.list_from_selector()

            if len(text_models) == 0 or not model_name:
                return -1
            detected_models = [i for i, m in enumerate(list(text_models.keys())) if m == model_name]
            if len(detected_models) > 0:
                return detected_models[0]
            return -1

        chat = self.get_root().chat_bin.get_child()
        if chat:
            model_index = -1
            if len(list(chat.container)) > 0:
                model_index = find_model_index(list(chat.container)[-1].get_model())
            if model_index == -1:
                model_index = find_model_index(self.get_root().get_current_instance().properties.get('default_model'))

            if model_index and model_index != -1:
                self.get_root().global_footer.model_selector.set_selected(model_index)

@Gtk.Template(resource_path='/com/jeffser/Alpaca/widgets/chat/template_selector.ui')
class TemplateSelector(Adw.Dialog):
    __gtype_name__ = 'AlpacaTemplateSelector'

    main_stack = Gtk.Template.Child()
    search_button = Gtk.Template.Child()

    def __init__(self, chat):
        super().__init__()
        self.chat = chat
        flow_box = self.main_stack.get_child_by_name('content')
        flow_box.remove_all()
        templates = SQL.get_templates()
        for t in templates:
            template_title = GLib.markup_escape_text(t[1])
            button = Gtk.Button(
                label=template_title,
                tooltip_text=_("Use '{}'").format(template_title)
            )
            button.connect('clicked', lambda *_, tem_name=t[1], tem_id=t[0]: self.selected_item(tem_name, tem_id))
            flow_box.append(button)
        flow_box.set_max_children_per_line(min(max(len(templates), 1), 3))
        self.update_visibility()

    @Gtk.Template.Callback()
    def on_search(self, entry:Gtk.SearchEntry):
        query = entry.get_text()
        flow_box = self.main_stack.get_child_by_name('content')
        for child in list(flow_box):
            child.set_visible(re.search(query, child.get_child().get_label(), re.IGNORECASE))
        self.update_visibility(True)

    def update_visibility(self, searching:bool=False):
        flow_box = self.main_stack.get_child_by_name('content')
        for child in list(flow_box):
            if child.get_visible():
                self.main_stack.set_visible_child_name('content')
                return

        self.main_stack.set_visible_child_name('no-results' if searching else 'empty')
        self.search_button.set_visible(searching)

    def selected_item(self, template_name:str, template_id:str):
        SQL.duplicate_chat(template_id, self.chat)
        self.chat.load_messages()
        self.close()

@Gtk.Template(resource_path='/com/jeffser/Alpaca/widgets/chat/character_greeting_selector.ui')
class CharacterGreetingSelector(Adw.Dialog):
    __gtype_name__ = 'AlpacaCharacterGreetingSelector'

    carousel = Gtk.Template.Child()

    def __init__(self, greetings:list, callback:callable):
        super().__init__()
        self.callback = callback
        for i, greeting in enumerate(greetings):
            msg = Message(
                dt=datetime.datetime.now(),
                message_id=generate_uuid(),
                mode=1,
                author=None
            )
            msg.header_options_button.unparent()
            msg.header_label.set_label(_('Greeting #{}').format(i+1))
            msg.block_container.set_content(greeting)

            scrolled_window = Gtk.ScrolledWindow(
                propagate_natural_height=True,
                propagate_natural_width=True,
                vexpand=True,
                css_classes=["card", "p0"],
                overflow=1,
                child=Gtk.Label(
                    wrap=True,
                    wrap_mode=2,
                    label=greeting,
                    css_classes=["body", "p10"]
                )
            )

            button = Gtk.Button(
                css_classes=["suggested-action", "pill"],
                halign=3,
                label=_("Use Greeting")
            )
            button.connect('clicked', lambda button, text=greeting: self.greeting_selected(text))

            container = Gtk.Box(
                orientation=1,
                spacing=10,
                css_classes=["p10"]
            )
            container.append(scrolled_window)
            container.append(button)

            self.carousel.append(container)

    def scroll_n_pages(self, number:int):
        index = self.carousel.get_position() + number
        if index < 0:
            index = self.carousel.get_n_pages() -1
        elif index > self.carousel.get_n_pages() -1:
            index = 0
        self.carousel.scroll_to(self.carousel.get_nth_page(index), True)

    def greeting_selected(self, text:str):
        self.callback(text)
        self.close()

    @Gtk.Template.Callback()
    def previous_page(self, button=None):
        self.scroll_n_pages(-1)

    @Gtk.Template.Callback()
    def next_page(self, button=None):
        self.scroll_n_pages(1)


@Gtk.Template(resource_path='/com/jeffser/Alpaca/widgets/chat/chat.ui')
class Chat(Gtk.Stack):
    __gtype_name__ = 'AlpacaChat'

    scrolledwindow = Gtk.Template.Child()
    container = Gtk.Template.Child()
    prompt_container = Gtk.Template.Child()
    welcome_screen = Gtk.Template.Child()
    use_template_button = Gtk.Template.Child()
    use_character_button = Gtk.Template.Child()

    def __init__(self, chat_id:str=None, name:str=_("New Chat"), folder_id:str=None, is_template=False):
        super().__init__()
        self.set_name(name)
        self.busy = False
        self.chat_id = chat_id
        self.folder_id = folder_id
        self.is_template = is_template
        self.row = ChatRow(self)
        self.use_template_button.set_visible(bool(self.chat_id))
        GLib.idle_add(self.update_prompts)
        GLib.idle_add(self.connect_model_selector)
        #self.connect('notify::root', lambda *_: self.connect_model_selector())

    def connect_model_selector(self):
        if self.get_root():
            global_footer = self.get_root().global_footer
            if global_footer:
                global_footer.model_selector.selector.connect('notify::selected', lambda selector, ud: self.on_model_change(selector))
                GLib.idle_add(self.on_model_change, global_footer.model_selector)

    def on_model_change(self, model_selector):
        char_dict = {}
        selected_item = model_selector.get_selected_item()
        if selected_item:
            char_dict = SQL.get_model_preferences(selected_item.model.get_name()).get('character', {})
        use_character = char_dict.get('data', {}).get('extensions', {}).get('com.jeffser.Alpaca', {}).get('enabled', False)
        self.use_character_button.set_visible(use_character)

    def on_search(self, query:str):
        query = re.escape(query)
        for m in list(self.container):
            content = m.get_content()
            if content:
                string_search = re.search(query, content, re.IGNORECASE)
                m.set_visible(string_search)
                for block in list(m.block_container):
                    if isinstance(block, blocks.text.Text):
                        if query:
                            query_escaped = GLib.markup_escape_text(query)
                            block_content = GLib.markup_escape_text(block.get_content())
                            highlighted_text = re.sub(f"({query_escaped})", r"<span background='yellow' bgalpha='30%'>\1</span>", block_content, flags=re.IGNORECASE)
                            block.set_markup(highlighted_text)
                        else:
                            block.set_content(block.get_content())

        GLib.idle_add(self.update_visibility, True)

    def update_visibility(self, searching:bool=False):
        for m in list(self.container):
            if m.get_visible():
                self.set_visible_child_name('content')
                return
        self.set_visible_child_name('no-results' if searching and len(list(self.container)) > 0 else 'welcome-screen')

    def stop_message(self):
        self.busy = False
        root = self.get_root() or self.row.get_root()
        if root:
            root.global_footer.toggle_action_button(True)

    def unload_messages(self):
        for widget in list(self.container):
            GLib.idle_add(widget.unparent)
            GLib.idle_add(widget.unrealize)
        #self.set_visible_child_name('loading')

    def add_message(self, message):
        self.container.append(message)
        GLib.idle_add(self.update_visibility)

    def load_messages(self):
        def load_message(message_element, content):
            attachments = SQL.get_attachments(message_element)
            message_element.block_container.set_content(content)

            for attachment in attachments:
                message_element.add_attachment(
                    file_id=attachment[0],
                    name=attachment[2],
                    attachment_type=attachment[1],
                    content=attachment[3]
                )

        messages = SQL.get_messages(self)
        for message in messages:
            message_element = Message(
                dt=datetime.datetime.strptime(message[3] + (":00" if message[3].count(":") == 1 else ""), '%Y/%m/%d %H:%M:%S'),
                message_id=message[0],
                mode=('user', 'assistant', 'system').index(message[1]),
                author=message[2]
            )
            threading.Thread(target=load_message, args=(message_element, message[4])).start()
            self.container.append(message_element)
        GLib.idle_add(self.update_visibility)

    def convert_to_ollama(self) -> list:
        messages = []
        for message in list(self.container):
            if message.get_content() and message.dt:
                message_data = {
                    'role': ('user', 'assistant', 'system')[message.mode],
                    'content': ''
                }

                for image in message.image_attachment_container.get_content():
                    if 'images' not in message_data:
                        message_data['images'] = []

                    message_data['images'].append(image['content'])

                for attachment in message.attachment_container.get_content():
                    if attachment.get('type') not in ('thought', 'metadata'):
                        message_data['content'] += '```{} ({})\n{}\n```\n\n'.format(attachment.get('name'), attachment.get('type'), attachment.get('content'))
                message_data['content'] += message.get_content()
                messages.append(message_data)
        return messages

    def convert_to_json(self, include_metadata:bool=False) -> list:
        messages = []
        for message in list(self.container):
            if message.get_content() and message.dt:
                message_data = {
                    'role': ('user', 'assistant', 'system')[message.mode],
                    'content': []
                }
                for image in message.image_attachment_container.get_content():
                    message_data['content'].append({
                        'type': 'image_url',
                        'image_url': f'data:image/png;base64,{image.get("content")}'
                    })
                message_data['content'].append({
                    'type': 'text',
                    'text': ''
                })
                for attachment in message.attachment_container.get_content():
                    if attachment.get('type') == 'thought':
                        message_data['thinking'] = attachment.get('content')
                    elif attachment.get('type') != 'metadata':
                        message_data['content'][0]['text'] += '```{} ({})\n{}\n```\n\n'.format(attachment.get('name'), attachment.get('type'), attachment.get('content'))
                message_data['content'][0 if ("text" in message_data.get("content", [''])[0]) else 1]['text'] += message.get_content()
                if include_metadata:
                    message_data['date'] = message.dt.strftime("%Y/%m/%d %H:%M:%S")
                    message_data['model'] = message.get_model()
                messages.append(message_data)
        return messages

    @Gtk.Template.Callback()
    def update_prompts(self, button=None):
        for el in list(self.prompt_container):
            el.unparent()
        prompts = random.sample(SAMPLE_PROMPTS, 3)

        for prompt in prompts:
            prompt_button = Gtk.Button(
                child=Gtk.Label(
                    label=prompt,
                    justify=2,
                    wrap=True
                ),
                tooltip_text=_("Send prompt: '{}'").format(prompt),
                name=prompt
            )
            prompt_button.connect('clicked', lambda *_, prompt=prompt : self.selected_prompt(prompt))
            self.prompt_container.append(prompt_button)

    @Gtk.Template.Callback()
    def show_template_selector(self, button):
        TemplateSelector(self).present(self.get_root())

    @Gtk.Template.Callback()
    def use_character(self, button:Gtk.Widget=None, ignore_greetings:bool=False, model_name:str=""):
        character_dict = {}

        if not model_name:
            selected_item = self.get_root().global_footer.model_selector.get_selected_item()
            if selected_item:
                model_name = selected_item.model.get_name()

        if model_name:
            character_dict = SQL.get_model_preferences(model_name).get('character', {})

        use_character = character_dict.get('data', {}).get('extensions', {}).get('com.jeffser.Alpaca', {}).get('enabled', False)
        if character_dict == {} or not use_character:
            if button:
                button.set_visible(False)
            return

        if len(character_dict.keys()) > 0:
            character_data = character_dict.get('data', {})

            description_content = character_data.get('description')
            if description_content:
                message_element = Message(
                    dt=datetime.datetime.now(),
                    message_id=generate_uuid(),
                    mode=2
                )
                attachment = message_element.add_attachment(
                    file_id = generate_uuid(),
                    name = _("Description"),
                    attachment_type = 'plain_text',
                    content = description_content
                )
                SQL.insert_or_update_attachment(message_element, attachment)
                self.add_message(message_element)
                message_element.block_container.set_content('')
                GLib.idle_add(message_element.save)

            first_message_content = character_data.get('first_mes')
            if first_message_content:
                message_element = Message(
                    dt=datetime.datetime.now(),
                    message_id=generate_uuid(),
                    mode=1,
                    author=model_name
                )
                self.add_message(message_element)
                message_element.block_container.set_content(first_message_content)
                GLib.idle_add(message_element.save)

            if not ignore_greetings:
                alternate_greetings = character_data.get('alternate_greetings', [])
                if len(alternate_greetings) > 0:
                    CharacterGreetingSelector(alternate_greetings, self.selected_prompt).present(self.get_root())

            GLib.idle_add(self.update_visibility)
        else:
            dialog.show_toast(
                message = _("The selected model does not have a character card"),
                root_widget = self.get_root()
            )

    def selected_prompt(self, prompt:str):
        if self.get_root().get_name() == 'AlpacaWindow':
            if len(list(self.get_root().model_manager.added_model_flowbox)) == 0:
                if self.get_root().get_current_instance().instance_type == 'empty':
                    self.get_root().get_application().lookup_action('instance_manager').activate()
                else:
                    self.get_root().get_application().lookup_action('model_manager').activate()
        buffer = self.get_root().global_footer.get_buffer()
        buffer.delete(buffer.get_start_iter(), buffer.get_end_iter())
        buffer.insert(buffer.get_start_iter(), prompt, len(prompt.encode('utf-8')))
        GLib.idle_add(self.get_root().send_message)

@Gtk.Template(resource_path='/com/jeffser/Alpaca/widgets/chat/folder_row.ui')
class FolderRow(Gtk.ListBoxRow):
    __gtype_name__ = 'AlpacaFolderRow'

    label = Gtk.Template.Child()

    def __init__(self, folder_id:str=None, folder_name:str=_('Root'), folder_color:str=None, folder_parent:str=None):
        super().__init__(
            name=folder_name
        )
        self.folder_id = folder_id
        self.folder_name = folder_name
        self.folder_color = folder_color
        self.folder_parent = folder_parent

        self.label.set_label(folder_name)
        self.get_child().set_tooltip_text(folder_name)
        if folder_color:
            self.add_css_class('folder-{}'.format(folder_color))

        drop_target_folder = Gtk.DropTarget.new(FolderRow, Gdk.DragAction.MOVE)
        drop_target_folder.connect("drop", self.on_drop_folder)
        self.add_controller(drop_target_folder)
        drop_target_chat = Gtk.DropTarget.new(ChatRow, Gdk.DragAction.MOVE)
        drop_target_chat.connect("drop", self.on_drop_chat)
        self.add_controller(drop_target_chat)

        drag_source = Gtk.DragSource()
        drag_source.set_actions(Gdk.DragAction.MOVE)
        drag_source.connect("drag-cancel", lambda *_: self.set_visible(True))
        drag_source.connect('prepare', lambda *_: Gdk.ContentProvider.new_for_value(self))
        drag_source.connect("drag-begin", lambda s,d: self.on_drag_begin(s,d,self.get_ancestor(Adw.NavigationPage)))
        drag_source.connect("drag-end", lambda s,d,r: self.on_drag_end(s,d,r,self.get_ancestor(Adw.NavigationPage)))
        self.add_controller(drag_source)

    def on_drop_folder(self, target, row, x, y):
        if row.folder_id != self.folder_id:
            SQL.move_folder_to_folder(row.folder_id, self.folder_id)
            row.get_parent().remove(row)
            folder_page = self.get_root().chat_list_navigationview.find_page(self.folder_id)
            if folder_page:
                folder_page.folder_list_box.prepend(row)
                row.set_visible(True)
                folder_page.update_visibility()
            return True

    def on_drop_chat(self, target, row, x, y):
        row.chat.folder_id = self.folder_id
        SQL.insert_or_update_chat(row.chat)
        row.get_parent().remove(row)
        folder_page = self.get_root().chat_list_navigationview.find_page(self.folder_id)
        if folder_page:
            folder_page.chat_list_box.prepend(row)
            row.set_visible(True)
            folder_page.update_visibility()
        return True

    def on_drag_begin(self, source, drag, page):
        page.top_indicator.set_visible(True)
        page.bottom_indicator.set_visible(True)
        snapshot = Gtk.Snapshot()
        self.snapshot_child(self.get_child(), snapshot)
        paintable = snapshot.to_paintable()
        source.set_icon(paintable, 0, 0)
        self.set_visible(False)
        GLib.idle_add(self.get_root().get_chat_list_page().scrolled_window.get_vadjustment().set_value, 0)

    def on_drag_end(self, source, drag, res, page):
        page.top_indicator.set_visible(False)
        page.bottom_indicator.set_visible(False)
        if page._scroll_timeout_id:
            GLib.source_remove(page._scroll_timeout_id)
            page._scroll_timeout_id = None

    @Gtk.Template.Callback()
    def show_popup(self, *args):
        rect = Gdk.Rectangle()
        if len(args) == 4:
            rect.x, rect.y = args[2], args[3]
        else:
            rect.x, rect.y = args[1], args[2]

        actions = [
            [
                {
                    'label': _('Edit Folder'),
                    'callback': self.prompt_edit,
                    'icon': 'document-edit-symbolic'
                }
            ],
            [
                {
                    'label': _('Delete Folder'),
                    'callback': self.prompt_delete,
                    'icon': 'user-trash-symbolic'
                }
            ]
        ]
        popup = dialog.Popover(actions)
        popup.set_parent(self)
        popup.set_pointing_to(rect)
        popup.popup()

    @Gtk.Template.Callback()
    def open_folder(self, button=None):
        if not self.get_root().chat_list_navigationview.find_page(self.folder_id):
            folder_page = Folder(self.folder_id, self.folder_name, self.folder_color, True)
            folder_page.set_tag(self.folder_id)
            self.get_root().chat_list_navigationview.add(folder_page)
            folder_page.update()
        self.get_root().chat_list_navigationview.push_by_tag(self.folder_id)

    def edit(self, new_name:str, new_color:str):
        if not new_name:
            new_name = _('New Folder')
        if new_name != self.folder_name:
            new_name = generate_numbered_name(new_name, [row.get_name() for row in list(self.get_parent())])
        self.folder_name = new_name
        self.label.set_label(new_name)
        self.label.set_tooltip_text(new_name)
        self.set_name(new_name)

        self.remove_css_class('folder-{}'.format(self.folder_color))
        self.folder_color = new_color
        self.add_css_class('folder-{}'.format(self.folder_color))
        SQL.insert_or_update_folder(self.folder_id, self.folder_name, self.folder_color, self.folder_parent)

        folder_page = self.get_root().chat_list_navigationview.find_page(self.folder_id)
        if folder_page:
            folder_page.set_title(new_name)
            folder_page.set_css_classes(['folder-{}'.format(self.folder_color)])

    def prompt_edit(self):
        options = {
            _('Cancel'): {},
            _('Accept'): {
                'appearance': 'suggested',
                'callback': lambda name, toggle_group: self.edit(name, toggle_group.get_active_name()),
                'default': True
            }
        }

        d = dialog.Entry(
            _('Edit Folder'),
            _("Editing '{}'").format(self.get_name()),
            list(options.keys())[0],
            options,
            {'placeholder': _('New Folder'), 'text': self.folder_name}
        )
        color_group = Adw.ToggleGroup()
        for c in ('blue', 'teal', 'green', 'yellow', 'orange', 'red', 'pink', 'purple', 'slate'):
            icon = Gtk.Image.new_from_icon_name('big-dot-symbolic')
            icon.add_css_class('button-{}'.format(c))
            icon.set_icon_size(2)
            toggle = Adw.Toggle(
                name=c,
                child=icon
            )
            color_group.add(toggle)
            if self.folder_color == c:
                color_group.set_active_name(c)
        d.container.append(color_group)

        d.show(self.get_root())

    def delete(self):
        if len(list(self.get_parent())) == 1:
            self.get_parent().set_visible(False)
            list(self.get_parent().get_parent())[1].set_visible(False)
        self.get_parent().remove(self)
        SQL.remove_folder(self.folder_id)

    def prompt_delete(self):
        dialog.simple(
            parent = self.get_root(),
            heading = _('Delete Folder'),
            body = _("Are you sure you want to delete '{}' and all it's sub-folders and chats?").format(self.get_name()),
            callback = lambda: self.delete(),
            button_name = _('Delete'),
            button_appearance = 'destructive'
        )

@Gtk.Template(resource_path='/com/jeffser/Alpaca/widgets/chat/chat_row.ui')
class ChatRow(Gtk.ListBoxRow):
    __gtype_name__ = 'AlpacaChatRow'

    label = Gtk.Template.Child()
    spinner = Gtk.Template.Child()
    indicator = Gtk.Template.Child()

    def __init__(self, chat:Chat):
        super().__init__()
        self.chat = chat
        self.set_name(self.chat.get_name())
        self.set_tooltip_text(self.chat.get_name())
        self.label.set_label(self.chat.get_name())

        drag_source = Gtk.DragSource()
        drag_source.set_actions(Gdk.DragAction.MOVE)
        drag_source.connect("drag-cancel", lambda *_: self.set_visible(True))
        drag_source.connect('prepare', lambda *_: Gdk.ContentProvider.new_for_value(self))
        drag_source.connect("drag-begin", lambda s,d: self.on_drag_begin(s,d,self.get_ancestor(Adw.NavigationPage)))
        drag_source.connect("drag-end", lambda s,d,r: self.on_drag_end(s,d,r,self.get_ancestor(Adw.NavigationPage)))
        self.add_controller(drag_source)

    def on_drag_begin(self, source, drag, page):
        page.top_indicator.set_visible(True)
        page.bottom_indicator.set_visible(True)
        snapshot = Gtk.Snapshot()
        self.snapshot_child(self.get_child(), snapshot)
        paintable = snapshot.to_paintable()
        source.set_icon(paintable, 0, 0)
        self.set_visible(False)
        GLib.idle_add(self.get_root().get_chat_list_page().scrolled_window.get_vadjustment().set_value, 0)

    def on_drag_end(self, source, drag, res, page):
        page.top_indicator.set_visible(False)
        page.bottom_indicator.set_visible(False)
        if page._scroll_timeout_id:
            GLib.source_remove(page._scroll_timeout_id)
            page._scroll_timeout_id = None

    @Gtk.Template.Callback()
    def show_popup(self, *args):
        rect = Gdk.Rectangle()
        if len(args) == 4:
            rect.x, rect.y = args[2], args[3]
        else:
            rect.x, rect.y = args[1], args[2]

        actions = [
            [
                {
                    'label': _('Edit Chat'),
                    'callback': self.prompt_edit,
                    'icon': 'document-edit-symbolic'
                },
                {
                    'label': _('Duplicate Chat'),
                    'callback': self.duplicate,
                    'icon': 'edit-copy-symbolic'
                },
                {
                    'label': _('Export Chat'),
                    'callback': self.prompt_export,
                    'icon': 'folder-download-symbolic'
                }
            ],
            [
                {
                    'label': _('Delete Chat'),
                    'callback': self.prompt_delete,
                    'icon': 'user-trash-symbolic'
                }
            ]
        ]
        if importlib.util.find_spec('kokoro'):
            actions[0].append({
                'label': _('Generate Podcast'),
                'callback': lambda: voice.PodcastDialog(self.chat).present(self.get_root()),
                'icon': 'audio-input-microphone-symbolic'
            })
        popup = dialog.Popover(actions)
        popup.set_parent(self)
        popup.set_pointing_to(rect)
        popup.popup()

    # Call in different thread
    def update_profile_pictures(self):
        for msg in list(self.chat.container):
            GLib.idle_add(msg.update_profile_picture)

    def edit(self, new_name:str, is_template:bool):
        if not new_name:
            new_name = _('New Chat')
        if new_name != self.chat.get_name():
            new_name = generate_numbered_name(new_name, [row.get_name() for row in list(self.get_parent())])
            self.label.set_label(new_name)
            self.label.set_tooltip_text(new_name)
            self.chat.set_name(new_name)
            self.set_name(new_name)
        self.chat.is_template = is_template
        SQL.insert_or_update_chat(self.chat)

        if self.get_parent().get_selected_row() == self:
            self.get_root().chat_page.set_title(self.get_name())

    def prompt_edit(self):
        options = {
            _('Cancel'): {},
            _('Accept'): {
                'appearance': 'suggested',
                'callback': lambda name, template_switch: self.edit(name, template_switch.get_active()),
                'default': True
            }
        }

        d = dialog.Entry(
            _('Edit Chat'),
            _("Editing '{}'").format(self.get_name()),
            list(options.keys())[0],
            options,
            {'placeholder': _('New Chat'), 'text': self.get_name()}
        )
        template_switch = Adw.SwitchRow(
            title=_('Use as Template'),
            active=self.chat.is_template,
            activatable=False
        )

        d.container.append(template_switch)

        d.show(self.get_root())

    def delete(self):
        window = self.get_root()
        list_box = self.get_parent()
        list_box.remove(self)
        SQL.delete_chat(self.chat)
        if len(list(list_box)) == 0:
            chat_list_page = window.get_chat_list_page()
            if chat_list_page.folder_id:
                previous_page = window.chat_list_navigationview.get_previous_page(chat_list_page)
                previous_page.chat_list_box.select_row(list(previous_page.chat_list_box)[0])
                previous_page.update_visibility()
            else:
                chat_list_page.new_chat()
            chat_list_page.update_visibility()
        if not list_box.get_selected_row() or list_box.get_selected_row() == self:
            list_box.select_row(list_box.get_row_at_index(0))
        if voice.message_dictated and voice.message_dictated.chat.chat_id == self.chat.chat_id:
            voice.message_dictated.popup.tts_button.set_active(False)

    def prompt_delete(self):
        dialog.simple(
            parent = self.get_root(),
            heading = _('Delete Chat'),
            body = _("Are you sure you want to delete '{}'?").format(self.get_name()),
            callback = lambda: self.delete(),
            button_name = _('Delete'),
            button_appearance = 'destructive'
        )

    def duplicate(self):
        new_chat_name = _("Copy of {}".format(self.get_name()))
        new_chat_id = generate_uuid()
        new_chat = self.get_root().get_chat_list_page().add_chat(
            chat_name=new_chat_name,
            chat_id=new_chat_id,
            is_template=False,
            mode=1
        )
        SQL.duplicate_chat(self.chat.chat_id, new_chat)

    def on_export_successful(self, file, result):
        file.replace_contents_finish(result)
        dialog.show_toast(_("Chat exported successfully"), self.get_root())

    def on_export_chat(self, file_dialog, result, temp_path):
        file = file_dialog.save_finish(result)
        if file:
            with open(temp_path, "rb") as db:
                file.replace_contents_async(
                    db.read(),
                    etag=None,
                    make_backup=False,
                    flags=Gio.FileCreateFlags.NONE,
                    cancellable=None,
                    callback=self.on_export_successful
                )

    def export_md(self, obsidian:bool):
        logger.info("Exporting chat (MD)")
        markdown = []
        for message_element in list(self.chat.container):
            if message_element.get_content() and message_element.dt:
                message_author = _('User')
                if message_element.get_model():
                    message_author = prettify_model_name(message_element.get_model())
                if message_element.mode == 2:
                    message_author = _('System')

                markdown.append('### **{}** | {}'.format(message_author, message_element.dt.strftime("%Y/%m/%d %H:%M:%S")))
                markdown.append(message_element.get_content())
                for file in message_element.image_attachment_container.get_content():
                    markdown.append('![🖼️ {}](data:image/{};base64,{})'.format(file.get('name'), file.get('name').split('.')[1], file.get('content')))
                emojis = {
                    'plain_text': '📃',
                    'code': '💻',
                    'pdf': '📕',
                    'youtube': '📹',
                    'website': '🌐',
                    'thought': '🧠'
                }
                for file in message_element.attachment_container.get_content():
                    if obsidian:
                        file_block = "> [!quote]- {}\n".format(file.get('name'))
                        for line in file.get('content').split("\n"):
                            file_block += "> {}\n".format(line)
                        markdown.append(file_block)
                    else:
                        markdown.append('<details>\n\n<summary>{} {}</summary>\n\n```TXT\n{}\n```\n\n</details>'.format(emojis.get(file.get('type'), '📃'), file.get('name'), file.get('content')))
                markdown.append('----')
        markdown.append('Generated from [Alpaca](https://github.com/Jeffser/Alpaca)')
        with open(os.path.join(cache_dir, 'export.md'), 'w') as f:
            f.write('\n\n'.join(markdown))
        file_dialog = Gtk.FileDialog(initial_name=f"{self.get_name()}.md")
        file_dialog.save(parent=self.get_root(), cancellable=None, callback=lambda file_dialog, result, temp_path=os.path.join(cache_dir, 'export.md'): self.on_export_chat(file_dialog, result, temp_path))

    def export_db(self):
        logger.info("Exporting chat (DB)")
        if os.path.isfile(os.path.join(cache_dir, 'export.db')):
            os.remove(os.path.join(cache_dir, 'export.db'))
        SQL.export_db(self.chat, os.path.join(cache_dir, 'export.db'))
        file_dialog = Gtk.FileDialog(initial_name=f"{self.get_name()}.db")
        file_dialog.save(parent=self.get_root(), cancellable=None, callback=lambda file_dialog, result, temp_path=os.path.join(cache_dir, 'export.db'): self.on_export_chat(file_dialog, result, temp_path))

    def export_json(self, include_metadata:bool):
        logger.info("Exporting chat (JSON)")
        with open(os.path.join(cache_dir, 'export.json'), 'w') as f:
            f.write(json.dumps({self.get_name() if include_metadata else 'messages': self.chat.convert_to_json(include_metadata)}, indent=4))
        file_dialog = Gtk.FileDialog(initial_name=f"{self.get_name()}.json")
        file_dialog.save(parent=self.get_root(), cancellable=None, callback=lambda file_dialog, result, temp_path=os.path.join(cache_dir, 'export.json'): self.on_export_chat(file_dialog, result, temp_path))

    def prompt_export(self):
        options = {
            _("Importable (.db)"): self.export_db,
            _("Markdown"): lambda: self.export_md(False),
            _("Markdown (Obsidian Style)"): lambda: self.export_md(True),
            _("JSON"): lambda: self.export_json(False),
            _("JSON (Include Metadata)"): lambda: self.export_json(True)
        }
        dialog.simple_dropdown(
            parent = self.get_root(),
            heading = _("Export Chat"),
            body = _("Select a method to export the chat"),
            callback = lambda option, options=options: options[option](),
            items = options.keys()
        )

