ui.py

import shutil
import os

from rich.console import Console
from rich.table import Table
from rich.panel import Panel
from rich.rule import Rule
from rich.markup import escape

from prompt_toolkit.key_binding import KeyBindings
from prompt_toolkit.widgets import TextArea
from prompt_toolkit.layout import Layout
from prompt_toolkit.styles import Style
from prompt_toolkit.application import Application
from prompt_toolkit import prompt

from utils import format_date

console = Console()


def print_inbox(emails, label="INBOX", labels=None, page=1, total_pages=1, is_sent=False):
    terminal_width = shutil.get_terminal_size().columns
    
    if total_pages > 1:
        title = f"{label} - Page {page}/{total_pages}"
    else:
        title = label

    table = Table(show_header=True, header_style="bold cyan", expand=True, padding=(0, 1))

    table.add_column("ID", style="dim", width=4, justify="right")
    table.add_column("★", width=1)
    if is_sent:
        table.add_column("To", ratio=1, overflow="ellipsis", no_wrap=True)
    else:
        table.add_column("From", ratio=1, overflow="ellipsis", no_wrap=True)
    table.add_column("Subject", ratio=3, overflow="ellipsis", no_wrap=True)
    table.add_column("Date", width=7, justify="right", style="dim")

    for mail in emails:
        id = mail["id"]

        if is_sent:
            from_header = mail["to"]
        else:
            from_header = mail["from"]

        if "<" in from_header:
            from_header = from_header.split("<")[0].strip().strip('"')

        from_header = escape(from_header.replace("\r", "").replace("\n", " ").strip())

        subject_header = escape(mail["subject"] or "")
        subject_header = " ".join(subject_header.splitlines()).strip()

        date_header = format_date(mail["date"])

        if "\\Flagged" in mail["flags"]:
            id = f"[yellow]{id}[/]"
            star = "[yellow]★[/]"
            subject_header = f"[yellow]{escape(subject_header)}[/]"
        else:
            star = " "

        is_read = "\\Seen" in mail["flags"]
        if not is_read:
            table.add_row(f"[bold]{id}[/]", star, f"[bold]{from_header}[/]", f"[bold]{subject_header}[/]", f"[bold]{date_header}[/]")
        else:
            table.add_row(id, star, from_header, subject_header, date_header)

    if terminal_width >= 80:
        sidebar = Table(show_header=False, box=None, expand=True, padding=(0, 1))
        sidebar.add_column("label", no_wrap=True, overflow="ellipsis")

        for la in labels:
            if la["name"] == label or la["display"] == label:
                sidebar.add_row(f"[reverse bold cyan]{la['display']}[/]")

            else:
                sidebar.add_row(la["display"])

        sidebar_height = len(labels) + 2
        table_height = len(emails) + 6
        height = max(sidebar_height, table_height)

        layout = Table.grid(expand=True)
        layout.add_column(width=22)
        layout.add_column(ratio=1)
        layout.add_row(Panel(sidebar, title="Labels", border_style="white", height=height), Panel(table, title=title, border_style="red", height=height))
        
        console.print(layout)

    else:
        console.print(Panel(table, title=title, border_style="red"))


def print_email(mail):
    if isinstance(mail, dict):
        mail = [mail]
    
    num = len(mail)
    subject = " ".join(mail[-1]["subject"].splitlines()).strip()
    console.print(Rule(f"[bold cyan]{subject}[/] ({num} emails)", style="cyan"))
    console.print()

    for item, mail in enumerate(mail):
        last = item == num - 1
        if last:
            border = "cyan"
        else:
            border = "white"

        from_header = f"[bold]From:[/bold] {mail['from']}\n"
        to_header = f"[bold]To:[/bold] {mail['to']}\n"
        id_header = f"[bold]ID:[/bold] {mail['id']}\n\n"
        date_header = f"[bold]Date:[/bold] {mail['date']}\n"
        subject_header = f"[bold]Subject:[/bold] {mail['subject']}"
        subject_header = " ".join(subject_header.splitlines()).strip()

        console.print(Panel(from_header + to_header + id_header + date_header + subject_header, border_style=border))

        attachments = mail["attachments"]
        if attachments:
            lines = []
            for number, attachment in enumerate(attachments, 1):
                filename = attachment["filename"]
                size_kb = attachment["size"] / 1024

                lines.append(f"{number}. {filename} ({size_kb:.1f} KB)")

            attachments_text = "\n".join(lines)

            console.print(Panel(attachments_text, title="Attachments", border_style="yellow"))


        mail_body = mail["body"]
        if mail_body:
            console.print(Panel(mail_body, border_style="green"))
        else:
            console.print(Panel("[italic dim]sybau[/]", border_style="green"))

        if not last:
            console.print()

    console.print("\n")


def compose_editor(to="", subject="", quote="", signature="", cc="", bcc="", body=""):
    console.print(Rule("New Email", style="cyan"))
    console.print()

    try:
        to = prompt("To: ", default=to or "").strip()
        if not to:
            console.print("\n[red]Recipient email is required.[/]")
            return None, None, None, None, None, [], False
        
        cc = prompt("CC (leave blank if none): ", default=cc or "").strip()
        bcc = prompt("BCC (leave blank if none): ", default=bcc or "").strip()

        subject = prompt("Subject: ", default=subject or "").strip()

    except KeyboardInterrupt:
        console.print("\n[red]Email composition cancelled.[/]")
        return None, None, None, None, None, [], False

    console.print()
    console.print(Rule(style="dim"))
    console.print("  [dim]Ctrl+S[/dim] [cyan]send[/cyan]    [dim]Ctrl+C[/dim] [cyan]cancel / save draft[/cyan]    [dim]Arrow keys[/dim] [cyan]move cursor[/cyan]")
    console.print(Rule(style="dim"))
    console.print()

    precompiled_stuff = []

    if signature:
        precompiled_stuff.append(f"\n\n-- \n{signature}")

    if quote:
        quoted_list = []
        for line in quote.splitlines():
            quoted_list.append(f"> {line}")

        quoted_text = "\n".join(quoted_list)

        precompiled_stuff.append(f"\n\n-----------------------------------------\n{quoted_text}")

    initial_email = (body or "") + "".join(precompiled_stuff)

    result = {"body": None, "cancelled": False}

    body_area = TextArea(
        text=initial_email,
        multiline=True,
        scrollbar=True,
        focus_on_click=True,
        wrap_lines=True,
        )
    
    body_area.buffer.cursor_position = 0

    keys = KeyBindings()

    @keys.add("c-s")
    def send_email(event):
        result["body"] = body_area.text
        event.app.exit()

    @keys.add("c-c")
    def cancel_email(event):
        result["body"] = body_area.text
        result["cancelled"] = True
        event.app.exit()


    style = Style.from_dict({
        "textarea": "bg:#1a1a1a #ffffff",
        "textarea focused": "bg:#1a1a1a #ffffff",
    })

    layout = Layout(body_area)

    app = Application(
        layout=layout,
        key_bindings=keys,
        style=style,
        mouse_support=True
    )
    
    app.run()

    if result["cancelled"]:
        body_for_now = result["body"] or ""
        console.print()
        answer = console.input("Save as draft? ([green]Y[/]/[red]n[/]): ").strip().lower()
        if answer == "y" or answer == "":
            return to, cc, bcc, subject, body_for_now, [], True
        
        else:
            console.print("\n[red]Cancelled.[/]")
            return None, None, None, None, None, [], False
    
    body = result["body"]
    
    attachments = []
    console.print("\nAdd attachments? Enter file paths or leave blank to skip or finish :D")
    while True:
        path = console.input("attach> ").strip().strip('"').strip("'")
        if not path:
            break

        if os.path.isfile(path):
            attachments.append(path)
            console.print(f"[green]Added:[/] {os.path.basename(path)}")
        else:
            console.print(f"[red]File not found:[/] {path}")

    return to, cc, bcc, subject, body, attachments, False


def confirm_send(to, cc, bcc, subject, body, attachments=None):
    console.print()

    to_cc_bcc = f"To: {to}\n"
    if cc:
        to_cc_bcc += f"CC: {cc}\n"
    if bcc:
        to_cc_bcc += f"BCC: {bcc}\n"

    attachments_text = ""
    if attachments:
        lines = []
        for number, attachment in enumerate(attachments, 1):
            filename = os.path.basename(attachment)
            size_kb = os.path.getsize(attachment) / 1024

            lines.append(f"{number}. {filename} ({size_kb:.1f} KB)\n")

        attachments_text = "\n\nAttachments:\n" + "".join(lines) + "\n"

    console.print(Panel(
        f"{to_cc_bcc}Subject: {subject}\n\n{body}{attachments_text}",
        title="Confirm Send",
        border_style="green"
    ))
    
    answer = console.input("Send this email? ([green]Y[/]/[red]n[/]): ").strip().lower()
    
    if answer == "y" or answer == "":
        return True
    else:
        return False
    

def confirm_remove_account(email):
    console.print()
    console.print(Panel(f"Are you sure you want to remove the account [bold red]{email}[/]?", border_style="red"))

    answer = console.input("Confirm removal? ([green]Y[/]/[red]n[/]): ").strip().lower()

    if answer == "y" or answer == "":
        return True
    else:
        return False
    

def print_search_inbox(emails, query, label="INBOX", labels=None):
    if not emails:
        console.print(Panel(f"No results found for '{query}' in {label}.", border_style="red"))
        return
    
    print_inbox(emails, label=label, labels=labels)


def print_settings(settings):
    sett = list(settings.items())

    table = Table(show_header=False, box=None, expand=True, padding=(0, 1))
    table.add_column("setting", no_wrap=True, style="cyan")
    table.add_column("value")
    table.add_column("description", style="dim")

    for key, value in sett:
        if key == "emails_per_page":
            table.add_row("emails_per_page", str(value), "Number of emails to show per page when listing inbox.")
        
        elif key == "download_folder":
            table.add_row("download_folder", value, "Default folder where attachments will be downloaded.")

        elif key == "signature":
            if value:
                table.add_row("signature", value, "Default signature added to the end of composed emails.")
            else:
                table.add_row("signature", "[dim]Not set[/]", "Default signature added to the end of composed emails.")

        elif key == "auto_mark_read":
            table.add_row("auto_mark_read", str(value), "Automatically mark emails as read when viewing them.")
        
        elif key == "default_label":
            table.add_row("default_label", value, "Label that will be opened by default on startup.")
    
    console.print(Panel(table, title="Settings", border_style="cyan"))


def print_aliases(aliases):
    if not aliases:
        console.print(Panel("No aliases set up yet.", border_style="red"))
        return
    
    table = Table(show_header=True, header_style="bold cyan", expand=True, padding=(0, 2))
    table.add_column("Alias", style="cyan", no_wrap=True)
    table.add_column("Command")
    
    for alias, command in aliases.items():
        table.add_row(alias, command)
    
    console.print(Panel(table, title="Aliases", border_style="cyan"))


def print_accounts(accounts, current_email):
    table = Table(show_header=True, header_style="bold cyan", expand=True, padding=(0, 2))
    table.add_column("Email")
    table.add_column("Currently using", width=10)

    for account in accounts:
        if account == current_email:
            table.add_row(f"[reverse bold cyan]{account}[/]", "this one")
        
        else:
            table.add_row(account, "")

    console.print(Panel(table, title="Accounts", border_style="cyan"))