#!/usr/bin/python3
import sys, os, locale, signal, re, argparse, stat
import gi, subprocess, time, threading, pickle
from locale import gettext as _
gi.require_version('Gtk', '3.0')
from gi.repository import Gtk, GLib, Gdk

CWD = os.environ.get('DDPYCWD', None)
if CWD:
    os.chdir(CWD)

ui_path = os.path.join(os.getcwd(), 'ddpy.ui')
iconpath = os.path.join(os.getcwd(), 'ddpy.svg')

if os.path.exists(ui_path):
    UI = ui_path
else:
    UI = '/usr/share/ddpy/ddpy.ui'

if os.path.exists(iconpath):
    ICON = iconpath
else:
    ICON = os.path.join('/usr/share/icons', 'ddpy.svg')

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)

CLI_DEVICE = None
CLI_SOURCE = None
CLI_HYBRID = None
CLI_ALIGN = None
CLI_MD5 = None
CLI_CLEAR = None

cfgdir = os.path.join(os.path.expanduser("~/.config/"), os.path.basename(sys.argv[0]))

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 check_rw(device):
    """
    Проверка прав на запись через попытку открытия
    Для надежности )
    """
    try:
        fd = os.open(device, os.O_RDWR | os.O_NONBLOCK)
        os.close(fd)
        return True
    except PermissionError:
        return False
    except BlockingIOError:
        return os.access(device, os.W_OK)
    except Exception as e:
        print(f"Error: {e}")
        return False 

def get_info(device):
  try:
    with open('/sys/block/' + device.split('/')[-1] + '/size', 'r') as f:
        size_blocks = int(f.read().strip())
        size_mb = size_blocks * 512 / 1024 / 1024
        size_str = f"{size_mb:.1f} MB"
    with open ('/sys/block/' + device.split('/')[-1] + '/device/vendor', 'r') as f:
        vendor = f.read().strip()
    return f"{vendor}, {size_str}"
  except:
    return 'Unknown'

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
    RUN = False
    global digest
    global CLI_DEVICE
    global CLI_CLEAR
    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
    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):
        mode = os.stat('/dev/sd' + char).st_mode
        if not stat.S_ISBLK(mode):
            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")))
    if CLI_SOURCE:
        app.file_entry.set_text(CLI_SOURCE)
        app.file_button.set_filename(CLI_SOURCE)
        app.file_button.emit("file-set")
    digest = []
    if CLI_DEVICE:
        app.dev_combobox.set_active_id(CLI_DEVICE)
        app.align_radio.set_active(CLI_ALIGN)
        app.hybrid_radio.set_active(CLI_HYBRID)
        app.md5_radio.set_active(CLI_MD5)
        CLI_DEVICE = None
        RUN = True
        if CLI_CLEAR:
          CLI_CLEAR = None
          GLib.timeout_add(500, lambda x=app.write_button: (app.handler.on_clear_button_clicked(x), False)[1])
        else:
          GLib.timeout_add(500, lambda x=app.write_button: (app.handler.on_write_button_clicked(x), False)[1])
 

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):
      device = app.dev_combobox.get_active_id() or None
      if device and not check_rw(device):
        # Перезапуск с pkexec
        args = [os.path.abspath(sys.argv[0])]
        args.append('--device')
        args.append(device)
        args.append('--clear')
        run_with_pkexec(args)
        return
      
      if not dialog(first=_('Caution!'),
                    second=_("Formatting cannot be undone. All data on the device will be permanently erased.")):
        return
      global RUN
      RUN = True
      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):
      global CLI_DEVICE
      global RUN
      # Проверка прав
      image = app.file_entry.get_text() or None
      device = app.dev_combobox.get_active_id() or None
      if device and not check_rw(device):
          # Перезапуск с pkexec
          if image:
              args = [os.path.abspath(sys.argv[0])]
              args.append(image)
              args.append('--device')
              args.append(device)
              if app.md5_radio.get_active(): 
                  args.append('--md5')
              if app.align_radio.get_active(): 
                  args.append('--align')
              if app.hybrid_radio.get_active(): 
                  args.append('--hybrid')
              run_with_pkexec(args)
              return
          else:
              return
      if not dialog(first=_('Warning!'),
                    second=_("The process cannot be undone. All existing data on the device will be completely overwritten.")):
        return
      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("")
    self.handler = WinHandler()
    builder.connect_signals(self.handler)

    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 run_with_pkexec(args):
    """Перезапуск с pkexec, передавая DISPLAY и XAUTHORITY."""
    try:
        os.execve(
            '/usr/bin/pkexec',
            ['/usr/bin/pkexec', 'env',
            'POLKITSGN=POLKITSGN',  # Фиксированный маркер для polkit policy
             f'DISPLAY={os.environ["DISPLAY"]}', 
             f'XAUTHORITY={os.environ["XAUTHORITY"]}',
             f'DDPYCWD={os.getcwd()}',
             os.path.abspath(sys.argv[0])] + args[1:],
             os.environ
        )
    except Exception as e:
        print(f"Failed to exec pkexec: {e}")
        sys.exit(1)

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(
        "--device", "-d",
        type=str,
        default=None,
        help="Device path (e.g., /dev/sda)"
    )
    parser.add_argument(
        "--debug",
        action="store_true",
        help="Enable debug mode"
    )

    parser.add_argument(
        "--help", "-h",
        action="store_true",
        help="Show help message"
    )
    parser.add_argument(
        "--md5", "-m",
        action="store_true",
        default=None,
        help="Enable MD5 checksum verification"
    )
    parser.add_argument(
        "--align", "-a",
        action="store_true",
        default=None,
        help="Enable GPT alignment"
    )
    parser.add_argument(
        "--hybrid", "-y",
        action="store_true",
        default=None,
        help="Enable hybrid MBR creation"
    )
    parser.add_argument(
        "--clear",
        action="store_true",
        default=None,
        help="Flash drive formatting (gpt / 1 partition / exfat)"
    )

    return parser

if __name__ == '__main__':
    # Suppress GTK errors when running as root
    if os.getuid() == 0:
        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
      
    # Сохраняем аргументы CLI для использования в GUI
    CLI_SOURCE = args.iso_path
    CLI_DEVICE = args.device
    CLI_HYBRID = args.hybrid
    CLI_ALIGN = args.align
    CLI_MD5 = args.md5
    CLI_CLEAR = args.clear

    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()
