#!/usr/bin/env python3
"""
pyannote speaker-diarization-3.1 Downloader
Downloads the repo files for pyannote/speaker-diarization-3.1
"""

import subprocess
import sys
import os

# Auto-install required packages
for package in ["huggingface_hub", "PySimpleGUI", "requests"]:
    try:
        __import__(package.replace("-", "_"))
    except ImportError:
        subprocess.check_call([sys.executable, "-m", "pip", "install", package])

import PySimpleGUI as sg
from huggingface_hub import list_repo_files, hf_hub_download
import shutil
import threading
import time
import traceback
from datetime import datetime

sg.theme('Dark Blue 3')
sg.set_options(font=('Segoe UI', 10))

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
REPO_ID = "pyannote/speaker-diarization-3.1"
REPO_TYPE = "model"
DOWNLOAD_DIR = os.path.join(SCRIPT_DIR, REPO_ID.split("/")[-1])

# Empty set => download everything. Add specific relative paths to whitelist if needed.
KEEP_ONLY = set()


def should_download(filename):
    return True if not KEEP_ONLY else filename in KEEP_ONLY


HF_TOKEN = os.environ.get("HF_TOKEN", "")


class DownloaderUI:
    def __init__(self):
        self.is_downloading = False
        self.is_paused = False
        self.start_time = None
        self.window = None
        self.download_thread = None
        self.force_download = False
        self.hf_token = HF_TOKEN

    def format_bytes(self, bytes_val):
        for unit in ['B', 'KB', 'MB', 'GB']:
            if bytes_val < 1024:
                return f"{bytes_val:.2f} {unit}"
            bytes_val /= 1024
        return f"{bytes_val:.2f} TB"

    def log_message(self, message):
        if self.window:
            current = self.window['-LOG-'].get()
            self.window['-LOG-'].update(current + message + '\n')

    def _apply_token(self):
        pass  # token passed directly to hf_hub_download

    def download_file(self, filename, filepath):
        try:
            if not self.force_download and os.path.exists(filepath):
                size = os.path.getsize(filepath)
                self.log_message(f"   ⏭️ Already exists ({self.format_bytes(size)}), skipping")
                return True, size

            self.log_message(f"   ⬇️ Downloading... (token={'set' if self.hf_token else 'MISSING'})")
            self.window['-PROGRESS-'].update(0)
            self.window['-PERCENT-'].update("0%")
            self.window['-SPEED-'].update("-- MB/s")
            self.window['-ETA-'].update("--:--")

            os.makedirs(os.path.dirname(filepath) if os.path.dirname(filepath) else DOWNLOAD_DIR, exist_ok=True)

            cached = hf_hub_download(
                repo_id=REPO_ID,
                repo_type=REPO_TYPE,
                filename=filename,
                token=self.hf_token or None,
                local_dir=DOWNLOAD_DIR,
                local_dir_use_symlinks=False,
            )

            if not self.is_downloading:
                return False, 0

            actual_path = os.path.join(DOWNLOAD_DIR, filename)
            if os.path.exists(actual_path):
                size = os.path.getsize(actual_path)
                self.window['-PROGRESS-'].update(100)
                self.window['-PERCENT-'].update("100%")
                self.window['-DOWNLOADED-'].update(self.format_bytes(size))
                self.log_message(f"   ✅ Saved: {self.format_bytes(size)}")
                return True, size
            else:
                self.log_message("   ❌ File not found after download!")
                return False, 0

        except Exception as e:
            traceback.print_exc()
            error_msg = str(e).encode('ascii', errors='replace').decode('ascii')[:120]
            self.log_message(f"   ❌ Error: {error_msg}")
            return False, 0
    def download_worker(self):
        retry_count = 0
        max_retries = 3

        while retry_count < max_retries:
            try:
                self.log_message(f"[{datetime.now().strftime('%H:%M:%S')}] 🔄 Fetching file list...")
                self.window['-STATUS-'].update("🔄 Fetching file list...")

                all_files = list_repo_files(repo_id=REPO_ID, repo_type=REPO_TYPE,
                                            token=self.hf_token or None)
                files_to_download = [f for f in all_files if should_download(f)]

                self.log_message(
                    f"[{datetime.now().strftime('%H:%M:%S')}] 📋 WHITELISTED: {len(files_to_download)} files"
                )

                total_files = len(files_to_download)
                for idx, filename in enumerate(files_to_download):
                    if not self.is_downloading:
                        self.log_message("⏸ Download cancelled")
                        return

                    filepath = os.path.join(DOWNLOAD_DIR, filename)
                    self.log_message(f"📥 [{idx + 1}/{total_files}] {filename}")
                    self.window['-STATUS-'].update(f"[{idx + 1}/{total_files}] {filename}")
                    self.window['-PROGRESS-'].update(0)
                    self.window['-PERCENT-'].update("0%")

                    success, size = self.download_file(filename, filepath)

                    if success:
                        self.log_message(f"   ✅ {self.format_bytes(size)}")
                    else:
                        self.log_message("   ⚠️ Skipped or failed")

                if self.is_downloading:
                    file_count = len([
                        f for f in os.listdir(DOWNLOAD_DIR)
                        if os.path.isfile(os.path.join(DOWNLOAD_DIR, f))
                    ])
                    total_size = sum(
                        os.path.getsize(os.path.join(DOWNLOAD_DIR, f))
                        for f in os.listdir(DOWNLOAD_DIR)
                        if os.path.isfile(os.path.join(DOWNLOAD_DIR, f))
                    )
                    self.window['-STATUS-'].update(
                        f"✅ Complete! {file_count} files, {self.format_bytes(total_size)}"
                    )
                    self.log_message(f"[{datetime.now().strftime('%H:%M:%S')}] ✅ All done!")
                    self.log_message(f"📁 Total: {file_count} files, {self.format_bytes(total_size)}")
                    self.window['-PROGRESS-'].update(100)
                    self.window['-PERCENT-'].update("100%")

                self.is_downloading = False
                break

            except Exception as e:
                retry_count += 1
                traceback.print_exc()
                error_msg = str(e).encode('ascii', errors='replace').decode('ascii')[:120]
                self.log_message(f"[ERROR] {error_msg}")
                self.log_message(f"[RETRY] {retry_count}/{max_retries} ...")
                self.window['-STATUS-'].update(f"Error, retrying... ({retry_count}/{max_retries})")
                time.sleep(3)

        if retry_count >= max_retries:
            self.window['-STATUS-'].update("❌ Failed - too many errors")
            self.log_message("❌ Download failed after multiple retries")

        self.is_downloading = False
        if self.window:
            self.window['-BUTTON-'].update("Start Download")
            self.window['-FORCE-'].update(disabled=False)
            self.window['-PAUSE-'].update(visible=False)

    def create_layout(self):
        layout = [
            [sg.Text("⚡ pyannote/speaker-diarization-3.1 Downloader", font=('Segoe UI', 14, 'bold'))],
            [sg.Text(
                "Downloads repository files for pyannote/speaker-diarization-3.1",
                text_color='#aac8e8',
                font=('Segoe UI', 9)
            )],

            [sg.Text("HF Token:", font=('Segoe UI', 9, 'bold')),
             sg.Input(default_text=self.hf_token, key='-TOKEN-', size=(40, 1),
                      password_char='*', font=('Segoe UI', 9))],

            [sg.Text("Status:", font=('Segoe UI', 10, 'bold')),
             sg.Text("Ready", key='-STATUS-', text_color='#4fc3f7', font=('Segoe UI', 10))],

            [sg.ProgressBar(100, size=(45, 20), key='-PROGRESS-',
                            bar_color=('#4fc3f7', '#1a3a5c'))],

            [sg.Column([
                [sg.Text("Progress:", font=('Segoe UI', 9, 'bold')),
                 sg.Text("0%", key='-PERCENT-', font=('Segoe UI', 9))],
                [sg.Text("Speed:", font=('Segoe UI', 9, 'bold')),
                 sg.Text("0 MB/s", key='-SPEED-', font=('Segoe UI', 9))],
            ]), sg.Column([
                [sg.Text("Downloaded:", font=('Segoe UI', 9, 'bold')),
                 sg.Text("0 B", key='-DOWNLOADED-', font=('Segoe UI', 9))],
                [sg.Text("ETA:", font=('Segoe UI', 9, 'bold')),
                 sg.Text("--:--", key='-ETA-', font=('Segoe UI', 9))],
            ])],

            [sg.Button('Start Download', key='-BUTTON-', size=(15, 2),
                       button_color=('#fff', '#0d6eaf')),
             sg.Button('Force Fresh', key='-FORCE-', size=(15, 2),
                       button_color=('#fff', '#b03030')),
             sg.Button('Pause', key='-PAUSE-', size=(15, 2), visible=False),
             sg.Button('Exit', size=(10, 2))],

            [sg.Multiline(size=(50, 12), key='-LOG-', disabled=True, autoscroll=True)],
        ]
        return layout

    def run(self):
        layout = self.create_layout()
        self.window = sg.Window(
            'pyannote/speaker-diarization-3.1 Downloader', layout, finalize=True, size=(560, 620)
        )

        self.log_message(f"📁 Download directory: {DOWNLOAD_DIR}")
        self.log_message("")
        self.log_message("⚠️  WARNING: This may download multiple large files.")
        self.log_message("💡 Click 'Start Download' to resume incomplete downloads")
        self.log_message("💡 Click 'Force Fresh' to re-download everything from scratch")
        self.log_message("")

        while True:
            event, values = self.window.read(timeout=500)

            if event == sg.WINDOW_CLOSED or event == 'Exit':
                if self.is_downloading:
                    if sg.popup_yes_no("Download in progress. Exit anyway?") != "Yes":
                        continue
                break

            if event == '-BUTTON-':
                if not self.is_downloading:
                    self.hf_token = values['-TOKEN-'].strip()
                    self._apply_token()
                    self.is_downloading = True
                    self.force_download = False
                    self.start_time = time.time()
                    self.log_message(f"[{datetime.now().strftime('%H:%M:%S')}] ⬇️ Download started (resume mode)")
                    self.window['-BUTTON-'].update('Cancel')
                    self.window['-PAUSE-'].update(visible=True)
                    self.window['-STATUS-'].update("🔄 Initializing...")
                    self.window['-PERCENT-'].update("0%")
                    self.window['-SPEED-'].update("0 MB/s")
                    self.window['-DOWNLOADED-'].update("0 B")
                    self.window['-PROGRESS-'].update(0)
                    self.download_thread = threading.Thread(target=self.download_worker, daemon=True)
                    self.download_thread.start()
                else:
                    self.is_downloading = False
                    self.window['-STATUS-'].update("⏹ Cancelled by user")
                    self.window['-BUTTON-'].update('Start Download')
                    self.window['-PAUSE-'].update(visible=False)
                    self.log_message(f"[{datetime.now().strftime('%H:%M:%S')}] ⏹ Download cancelled")

            if event == '-FORCE-':
                if not self.is_downloading:
                    self.hf_token = values['-TOKEN-'].strip()
                    self._apply_token()
                    self.is_downloading = True
                    self.force_download = True
                    self.start_time = time.time()
                    self.log_message(f"[{datetime.now().strftime('%H:%M:%S')}] 🔄 FORCE download started (will re-download all files)")
                    self.window['-BUTTON-'].update('Cancel')
                    self.window['-FORCE-'].update(disabled=True)
                    self.window['-PAUSE-'].update(visible=True)
                    self.window['-STATUS-'].update("🔄 Initializing...")
                    self.window['-PERCENT-'].update("0%")
                    self.window['-SPEED-'].update("0 MB/s")
                    self.window['-DOWNLOADED-'].update("0 B")
                    self.window['-PROGRESS-'].update(0)
                    self.download_thread = threading.Thread(target=self.download_worker, daemon=True)
                    self.download_thread.start()

            if event == '-PAUSE-' and self.is_downloading:
                self.is_paused = not self.is_paused
                if self.is_paused:
                    self.window['-PAUSE-'].update('Resume')
                    self.window['-STATUS-'].update("⏸ Paused")
                    self.log_message("⏸ Paused - click Resume to continue")
                else:
                    self.window['-PAUSE-'].update('Pause')
                    self.window['-STATUS-'].update("🔄 Resuming...")
                    self.log_message("▶️ Resumed")

        self.window.close()


if __name__ == "__main__":
    downloader = DownloaderUI()
    downloader.run()
