#!/usr/bin/python3
import sys, os, locale, signal, re, argparse
import gi, subprocess, time, threading, pickle
from locale import gettext as _
gi.require_version('Gtk', '3.0')
from gi.repository import Gtk, GLib, Gdk
if os.path.exists('./ddlib.py'):
  from ddlib import *
else:
  from ddpy.ddlib import *

digest = []
PULSE = False
RUN = False
DEBUG = False
_threads = {}

DOMAIN = 'ddpy'
LOCALE_DIR = '/usr/share/locale'
locale.bindtextdomain(DOMAIN, LOCALE_DIR)
locale.textdomain(DOMAIN)

UI = '/usr/share/ddpy/ddpy.ui'
iconpath = '/usr/share/icons'
cfgdir = os.path.join(os.path.expanduser("~/.config/"), os.path.basename(sys.argv[0]))

if os.path.exists('ddpy.ui'):
  UI = './ddpy.ui'

ICON = os.path.join(iconpath, 'ddpy.svg')

def pickle_write():
    os.makedirs(cfgdir, exist_ok=True)
    with open(cfgdir + '/dump', 'wb') as f:
       pickle.dump(config, f)

def dialog(first='Warning!', second=None, mtype='info',
          btn=[Gtk.STOCK_OK, Gtk.ResponseType.OK,
          Gtk.STOCK_CANCEL, Gtk.ResponseType.CANCEL]):
    '''
    По умолчанию информационное окно с кнопками ОК, отмена
    Если нажать OK - вернет True
    Если CANCEL или закрыть окно - False
    Если свои кнопки вернет, то что назначено кнопке
    '''
    if mtype == 'info':
        mtype = Gtk.MessageType.INFO
    elif mtype == 'error':
        mtype = Gtk.MessageType.ERROR
    text = None
    parent = None
    if 'window' in globals():
        parent = window
    messagedialog = Gtk.MessageDialog(transient_for=parent,
            destroy_with_parent=True,
            modal=True,
            message_type=mtype,
            buttons=Gtk.ButtonsType.NONE,
            text=first)
    messagedialog.set_keep_above(True)
    messagedialog.set_skip_taskbar_hint(True)
    messagedialog.set_skip_pager_hint(True)
    messagedialog.set_decorated(False)

    if second:
        messagedialog.format_secondary_text(second)
    # кнопки
    buttons = []
    for i, (label, response_id) in enumerate(zip(btn[::2], btn[1::2])):
        button = messagedialog.add_button(label, response_id)
        buttons.append(button)
        if i == 0:
          button.set_can_default(True)
          button.grab_default()
    messagedialog.set_default_response(btn[1])
    response = messagedialog.run()
    messagedialog.show_all()
    messagedialog.destroy()
    if response == Gtk.ResponseType.OK:
        return True
    return False

def show_digest(parent_window, title, items):
    """
    :param parent_window: Родительское окно (Gtk.Window)
    :param title: Заголовок окна
    :param items: Список строк для отображения
    """
    dialog = Gtk.Dialog(transient_for=parent_window,
                      modal=True,
                      title=title)
    dialog.set_default_size(400, 300)
    dialog.set_border_width(0)
    header = Gtk.HeaderBar()
    header.set_show_close_button(False)
    header.set_title(title)
    dialog.set_titlebar(header)
    close_button = Gtk.Button.new_from_icon_name("window-close", Gtk.IconSize.BUTTON)
    close_button.connect("clicked", lambda _: dialog.destroy())
    header.pack_end(close_button)
    content_area = dialog.get_content_area()
    scrolled_window = Gtk.ScrolledWindow()
    scrolled_window.set_policy(Gtk.PolicyType.NEVER, Gtk.PolicyType.AUTOMATIC)
    content_area.pack_start(scrolled_window, True, True, 0)
    list_box = Gtk.ListBox()
    list_box.set_selection_mode(Gtk.SelectionMode.NONE)
    scrolled_window.add(list_box)
    for item in items:
        row = Gtk.ListBoxRow()
        label = Gtk.Label(label=item)
        label.set_xalign(0)
        label.set_margin_start(10)
        label.set_margin_end(10)
        label.set_margin_top(5)
        label.set_margin_bottom(5)
        row.add(label)
        list_box.add(row)
    button_box = Gtk.ButtonBox(orientation=Gtk.Orientation.HORIZONTAL)
    button_box.set_layout(Gtk.ButtonBoxStyle.END)
    button_box.set_margin_bottom(10)
    button_box.set_margin_end(10)

    ok_button = Gtk.Button(label="OK")
    ok_button.get_style_context().add_class("suggested-action")
    ok_button.connect("clicked", lambda _: dialog.destroy())
    button_box.add(ok_button)

    content_area.pack_end(button_box, False, False, 0)
    content_area.show_all()
    dialog.run()

def start_pulse(app, action):
  global PULSE
  PULSE = action
  if PULSE:
    GLib.timeout_add(300, lambda app=app: update_pulse(app))

def update_pulse(app):
  if PULSE:
    app.progress.pulse()
    return True
  return False

def get_info(device):
  result = sgdisk(device, ["-p",])
  if result:
    for line in result.stdout.splitlines():
        line = line.strip()
        if not line:
          continue
        if line.startswith('Disk /dev/'):
          size = line.split(',')[1]
        elif line.startswith('Model:'):
          model = line.split(':')[1]
    return f'{model} {size}'
  return ''

def threading_me(func, *args, **kwargs):
    """Запускает функцию в фоне и возвращает её результат (или None при остановке)."""
    stop_event = threading.Event()
    result_container = [None]
    thread_id = threading.get_ident()

    def wrapper():
        current_thread = threading.current_thread()
        current_thread.stop_event = stop_event
        try:
            if not stop_event.is_set():
                result_container[0] = func(*args, **kwargs)
        finally:
            _threads.pop(thread_id, None)

    thread = threading.Thread(target=wrapper, daemon=True)
    _threads[thread_id] = (thread, stop_event, result_container)
    thread.start()

    while thread.is_alive():
        GLib.MainContext.default().iteration(False)
        if stop_event.is_set():
            thread.join(0.2)
            return None

    return result_container[0]

def stop_all_threads():
     for thread_id, (thread, stop_event, _) in list(_threads.items()):
        stop_event.set()
        thread.join(0.2)
        if thread.is_alive():  # Добить если жив
            try:
                thread._stop()
            except:
                pass
        _threads.pop(thread_id, None)

def write(source, dest, app):
  global RUN
  params = {}
  params['source_file'] = source
  params['destination_device'] = dest
  params['block_size'] = '8M'
  sname = trunk_str(source)
  if not params['source_file'] or not params['destination_device']:
    return
  GLib.idle_add(lambda : app.update_button.set_sensitive(False))
  GLib.idle_add(lambda : app.clear_button.set_sensitive(False))
  GLib.idle_add(lambda : app.write_button.set_sensitive(False))
  start_pulse(app, True)
  if not threading_me(zero_and_verify, dest, custom_print=echo):
    echo(_('Write zero block check - Failed!'), prn=True)
    return
  echo(_('Write zero block check - Complete'), prn=True)
  if not RUN:
    return
  base_check = True
  if app.md5_radio.get_active():
    base_check = threading_me(checkisomd5, source, custom_print=echo)
    if base_check == 'NA':
      echo(_('BuiltIn iso md5 sum - not available'), prn=True)
      iso_size = os.path.getsize(source)
      offset = 0
      iso_md5 = threading_me(calculate_md5, source, offset, iso_size, custom_print=echo)
    elif base_check:
      echo(_('Image builtin md5 sum check - Complete'), prn=True)
    else:
      echo(_('{sname} builtin md5 sum check - Failed').format(sname=sname), prn=True)
      return
  start_pulse(app, False)
  GLib.idle_add(lambda : app.progress.set_fraction(0.0))
  if not threading_me(copy_file_to_device, params, custom_print=echo) and base_check:
    echo(_(f'Copy image data - Failed!'), prn=True)
    return
  echo(_(f'Copy data - Complete!'), prn=True)
  start_pulse(app, True)
  if RUN and app.md5_radio.get_active():
    if base_check == 'NA':
      dev_md5=threading_me(calculate_md5, dest, offset, iso_size, custom_print=echo)
      if dev_md5 != iso_md5:
          echo(_('Md5 sums of {sname} and {dest} not equal!').format(sname=sname, dest=dest), prn=True)
          return
      echo(_('Md5 sum for Image and {dest} are equal!').format(sname=sname, dest=dest), prn=True)
    else:
      if not threading_me(checkisomd5, dest, custom_print=echo):
        echo(_('Checkisomd5 result for {dest} is: Failed!').format(dest=dest), prn=True)
        return
      echo(_('Builtin md5 sum check for {dest} - Complete!').format(dest=dest), prn=True)
  if RUN and app.align_radio.get_active():
    if not sgdisk(dest, "-eg", custom_print=echo):
      echo(_('Set real disk size into GPT partition table - Failed!'), prn=True)
      return
    echo(_('Set real disk size into GPT partition table - Complete!'), prn=True)
  if RUN and app.hybrid_radio.get_active():
    info = get_partition_info(dest)
    if info:
      num_parts = ", ".join(str(i) for i in range(1, len(info) + 1))
      if not sgdisk(dest, {
                    "-h": num_parts,
                    "-U": "random" },
                    custom_print=echo):
        echo(_('Creating of hybrid partiton table - Failed!'), prn=True)
        return
    echo(_('Creating of hybrid partiton table - Complete!'), prn=True)

  threading_me(subprocess.run,
            [ 'partprobe', dest ],
            stdout=subprocess.PIPE,
            text=True )
  return True

def clear(dev):
  if not dev:
    return
  GLib.idle_add(lambda : app.update_button.set_sensitive(False))
  GLib.idle_add(lambda : app.clear_button.set_sensitive(False))
  GLib.idle_add(lambda : app.write_button.set_sensitive(False))
  if RUN and threading_me(zero_and_verify, dev, custom_print=echo):
    echo(_('Write zero block check - Complete!'), prn=True)
    if RUN and threading_me(sgdisk, dev, ['--zap-all', '--new=1:0:0',
                 '--change-name=1:data', '--typecode=1:0700'],
                custom_print=echo, verbose=True ):
      echo(_('Creating partition - Complete!'), prn=True)
      threading_me(subprocess.run,
              ['partprobe', dev ],
              check=True,
              stdout=subprocess.PIPE,
              text=True
          )
      if RUN:
        try:
          threading_me(subprocess.run,
              ['mkfs.exfat', dev + '1'],
              check=True,
              stdout=subprocess.PIPE,
              stderr=subprocess.PIPE,
              text=True
          )
          echo(_('Formatting - Complete'), prn=True)
          return True
        except subprocess.CalledProcessError as e:
          echo(_(f"Formatting error (code {e.returncode}):"), prn=True)
          print(e.stderr)
        except Exception as e:
          echo (_("Forrmatting error"), prn=True)
          print(f"Unknown error: {str(e)}")

def init():
    global RUN
    global digest
    start_pulse(app, False)
    app.update_button.set_sensitive(True)
    app.clear_button.set_sensitive(True)
    app.write_button.set_sensitive(True)
    app.file_button.set_current_folder(config['path'])
    app.window.set_title(_('ddpy-gtk - write ISO into removable devices'))
    app.progress.set_pulse_step(0.1)
    app.progress.set_fraction(0.0)
    n = 0
    if len(sys.argv) > 1:
      file_path = os.path.abspath(sys.argv[1])
      if os.path.exists(file_path):
        app.file_button.set_filename(file_path)
        app.file_button.emit("file-set")
      else:
        print('File not exists: ' + file_path)
    app.dev_combobox.remove_all()
    with open('/proc/mounts', 'r') as mounts:
      procmounts = mounts.read().split()
    for char in 'nmlkjigfedcba':
      model = 'unknown'
      size = 'unknowm'
      dev = '/sys/block/sd'+ char
      if os.path.islink(dev):
        try:
          with open('/dev/sd' + char, 'rb') as device:
            device.seek(0)
            first_bytes = device.read(10)
        except:
          continue
        try:
          with open(dev + '/removable', 'rb') as check_removable:
            code = int(check_removable.read())
            if code != 1:
              print(f'/dev/sd{char} not removable')
              continue
        except:
          continue
        if  any(entry.startswith('/dev/sd' + char) for entry in procmounts):
          app.dev_combobox.insert(n, f'/dev/sd{char}', '/dev/sd{char} - {unmount}'.format(char=char, unmount=_("is in use, unmount first")))
        else:
          info = get_info(f'/dev/sd{char}')
          app.dev_combobox.insert(n, f'/dev/sd{char}', f'/dev/sd{char} - {info}' )
          app.dev_combobox.set_active_id(f'/dev/sd{char}')
        n += 1
      elif os.path.exists('/dev/sd' + char):
        app.dev_combobox.insert(n, f'/dev/sd{char}', '/dev/sd{char} - {unknown}'.format(char=char, unknown=_("Unknown, try to claer first")))
    RUN = False
    digest = []

def echo(text, prn=False):
  global DEBUG
  global digest
  text = text.replace('\n', '').replace('\r', '').strip()
  if text.endswith('%'):
    t = text.split('|')
    percent =t[1]
    text = t[0]
    prg = float(percent[:-1]) / 100
    GLib.idle_add(lambda p=prg: app.progress.set_fraction(p))
  truncated = (text[:25] + '...' + text[len(text)-25:]) if len(text) > 53 else text
  GLib.idle_add(lambda t=truncated: app.status_label.set_text(t))
  if prn or DEBUG:
    digest.append(text)
    print(text)

def exit_(*args):
  Gtk.main_quit()

class WinHandler:
    global app
    def on_file_button_file_set(self, button):
      filename = app.file_button.get_filename()
      config['path'] = os.path.dirname(filename)
      pickle_write()
      if filename:
        app.file_entry.set_text(filename)
        if is_gpt(filename):
          app.align_radio.set_sensitive(True)
          app.hybrid_radio.set_sensitive(True)
        else:
          app.align_radio.set_active(False)
          app.hybrid_radio.set_active(False)
          app.align_radio.set_sensitive(False)
          app.hybrid_radio.set_sensitive(False)
      else:
          app.file_entry.set_text("")
    def on_update_button_clicked(self, button):
      GLib.idle_add(init)

    def on_clear_button_clicked(self, button):
      if not dialog(first=_('Caution!'),
                    second=_("Formatting cannot be undone. All data on the device will be permanently erased.")):
        return
      global RUN
      RUN = True
      device = app.dev_combobox.get_active_id() or None
      if clear(device):
        echo(_('All - Complete!'), prn=True)
      else:
        echo(_('Something went wrong'), prn=True)
      show_digest(app.window, _('digest'), digest)
      GLib.idle_add(init)

    def on_abort_button_clicked(self, button):
      global RUN
      if RUN:
        stop_all_threads()
        RUN = False
        return
      exit_()

    def on_write_button_clicked(self, button):
      if not dialog(first=_('Warning!'),
                    second=_("The process cannot be undone. All existing data on the device will be completely overwritten.")):
        return
      global RUN
      RUN = True
      device = app.dev_combobox.get_active_id() or None
      image = app.file_entry.get_text() or None
      if write(image, device, app):
        echo(_('All - Complete!'), prn=True)
      else:
        echo(_('Something went wrong'), prn=True)
      show_digest(app.window, _('digest'), digest)
      GLib.idle_add(init)

class AppWindow:
  def __init__(self):
    # window
    signal.signal(signal.SIGINT, signal.SIG_DFL)
    builder = Gtk.Builder()
    self.builder = builder
    builder.set_translation_domain(DOMAIN)
    builder.add_from_file(UI)

    self.file_entry = builder.get_object('file_entry')
    self.dev_combobox = builder.get_object('dev_combobox')
    self.progress = builder.get_object('progress')

    self.status_label = builder.get_object('status_label')
    self.align_radio = builder.get_object('align_radio')
    self.hybrid_radio = builder.get_object('hybrid_radio')
    self.md5_radio = builder.get_object('md5_radio')

    self.logo = builder.get_object('logo')

    self.file_button = builder.get_object('file_button')
    self.update_button = builder.get_object('update_button')
    self.clear_button = builder.get_object('clear_button')
    self.write_button = builder.get_object('write_button')
    self.abort_button = builder.get_object('abort_button')
    self.abort_button.set_always_show_image(True)
    self.abort_button.set_label("")
    builder.connect_signals(WinHandler())

    self.window = builder.get_object('window_main')
    self.window.set_icon_from_file(ICON)
    self.window.connect("destroy", Gtk.main_quit)
    self.apply_styles()

  def apply_styles(self):
    css = b"""
    #rosa_label {
        font-size: 30px;
        font-weight: bolder;
    }
    """
    provider = Gtk.CssProvider()
    provider.load_from_data(css)

    screen = Gdk.Screen.get_default()
    style_context = Gtk.StyleContext()
    style_context.add_provider_for_screen(
        screen,
        provider,
        Gtk.STYLE_PROVIDER_PRIORITY_APPLICATION
    )

  def run(self):
    self.window.show_all()
    Gtk.main()

def HLP():
    help_text = """
ddpy-gtk - Graphical tool for writing ISO images to removable devices

Usage:
  ddpy-gtk [OPTIONS]
  ddpy-gtk <iso_path>

Description:
  This is a graphical tool based on GTK3 for writing ISO images to removable
  devices (USB drives, etc.). It provides features like:
  - MD5 checksum verification
  - GPT alignment
  - Hybrid partition table creation
  - Device clearing and formatting

The tool must be run as root. If not run as root, it will automatically
attempt to restart itself with pkexec for authentication.
"""
    print(help_text)


def parse_args():
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument(
        "iso_path",
        nargs="?",
        default=None,
        help="Absolute PATH to ISO image (optional)"
    )
    parser.add_argument(
        "--debug",
        action="store_true",
        help="Enable debug mode"
    )

    parser.add_argument(
        "--help", "-h",
        action="store_true",
        help="Show help message"
    )
    return parser


if __name__ == '__main__':
    # Suppress GTK errors when running as root
    os.environ.update({
        'G_MESSAGES_DEBUG': '',
        'GTK_DEBUG': 'none',
        'DCONF_PROFILE': '',
        'GSETTINGS_BACKEND': 'memory',
        'XDG_RUNTIME_DIR': '/tmp/runtime-root'
    })

    parser = parse_args()
    args = parser.parse_args()

    if args.help:
        HLP()
        parser.print_help()
        sys.exit(0)

    if args.debug:
      DEBUG = True

    if os.getuid() != 0:
        os.execl(
            '/usr/bin/pkexec',
            'pkexec',
            os.path.abspath(sys.argv[0]),
            *([args.iso_path] if args.iso_path else [])
       )
        sys.exit(0)

    config = {}
    config['path'] = os.getenv('XHOME', '/home')

    if os.path.exists(cfgdir + '/dump'):
        try:
          with open(cfgdir + '/dump', 'rb') as f:
            config = pickle.load(f)
        except pickle.UnpicklingError as e:
          print(f"Pickle read data error: {e}")
          try:
            os.remove(cfgdir + '/dump')
          except:
            pass
        except PermissionError as e:
          print(f"Permissions error: {e}")

    app = AppWindow()
    if args.debug:
        os.environ['DEBUG'] = '1'
    GLib.idle_add(init)
    app.run()
