// ==UserScript==
// @name         Hugging Face Repo Size Display (Recursive with TB and Copy URLs)
// @namespace    http://tampermonkey.net/
// @version      0.5
// @description  Show repo and directory file sizes with recursive directory size calculation, TB support, and copy download URLs
// @author       You
// @match        https://huggingface.co/*/blob/*
// @match        https://huggingface.co/*/tree/*
// @match        https://huggingface.co/*/src/*
// @grant        none
// ==/UserScript==

(function () {
    'use strict';

    // Function to extract repo name and path from URL
    function getRepoInfo() {
        const url = window.location.pathname;
        const match = url.match(/\/([^\/]+)\/([^\/]+)(?:\/(?:blob|tree|src)\/([^\/]+)(\/.*)?)?/);
        if (!match) return null;

        const [, username, reponame] = match;
        const repoFullName = `${username}/${reponame}`;

        let path = '';
        if (match[4]) {
            path = match[4].startsWith('/') ? match[4].substring(1) : match[4];
        }

        return { repoFullName, path };
    }

    // Fetch repository files using Hugging Face API
    async function fetchRepoFiles(repoName) {
        try {
            const response = await fetch(`https://huggingface.co/api/models/${repoName}/tree/main?recursive=true`);
            if (!response.ok) throw new Error(`HTTP ${response.status}`);
            return await response.json();
        } catch (error) {
            console.log('Failed to fetch repo files:', error);
            return [];
        }
    }

    // Recursively fetch and calculate size of a directory
    async function getDirectorySize(repoName, directoryPath) {
        try {
            const url = directoryPath ?
                `https://huggingface.co/api/models/${repoName}/tree/main/${directoryPath}?recursive=true` :
                `https://huggingface.co/api/models/${repoName}/tree/main?recursive=true`;

            const response = await fetch(url);
            if (!response.ok) return 0;

            const files = await response.json();
            let totalSize = 0;

            files.forEach(file => {
                if (file.size && (file.type === 'file' || file.type === 'file')) {
                    totalSize += file.size;
                }
            });

            return totalSize;
        } catch (error) {
            console.log(`Failed to fetch directory size for ${directoryPath}:`, error);
            return 0;
        }
    }

    // Calculate the total repository size by recursively processing all directories
    async function getRepoSize(repoName) {
        try {
            // Get the complete file tree
            const response = await fetch(`https://huggingface.co/api/models/${repoName}/tree/main?recursive=true`);
            if (!response.ok) return 0;

            const files = await response.json();
            let totalSize = 0;

            // Process each file/directory
            for (const item of files) {
                if (item.type === 'file' && item.size) {
                    totalSize += item.size;
                } else if (item.type === 'dir') {
                    // Recursively calculate size of this directory
                    const dirSize = await getDirectorySize(repoName, item.path);
                    totalSize += dirSize;
                }
            }

            return totalSize;
        } catch (error) {
            console.log('Failed to calculate repository size:', error);
            return 0;
        }
    }

    // Calculate the size of files in the current directory only
    async function getDirectorySizeCurrent(repoName, currentPath) {
        try {
            // If currentPath is empty, we're at root
            if (!currentPath) {
                // Get root-level files
                const response = await fetch(`https://huggingface.co/api/models/${repoName}/tree/main?recursive=false`);
                if (!response.ok) return 0;

                const items = await response.json();
                let totalSize = 0;

                items.forEach(item => {
                    if (item.type === 'file' && item.size) {
                        totalSize += item.size;
                    }
                });

                return totalSize;
            }

            // For subdirectories, get files directly in this directory
            const response = await fetch(`https://huggingface.co/api/models/${repoName}/tree/main/${currentPath}?recursive=false`);
            if (!response.ok) return 0;

            const items = await response.json();
            let totalSize = 0;

            items.forEach(item => {
                if (item.type === 'file' && item.size) {
                    totalSize += item.size;
                }
            });

            return totalSize;
        } catch (error) {
            console.log('Failed to calculate directory size:', error);
            return 0;
        }
    }

    // Format bytes to human readable format with TB support
    function formatBytes(bytes) {
        if (bytes === 0) return '0 B';

        // Define units including TB
        const sizes = ['B', 'KB', 'MB', 'GB', 'TB'];
        const k = 1024;

        // Calculate the power
        const i = Math.floor(Math.log(bytes) / Math.log(k));

        // Cap at TB (index 4)
        const displayIndex = Math.min(i, 4);
        const divisor = Math.pow(k, displayIndex);
        const formatted = (bytes / divisor).toFixed(2);

        return `${formatted} ${sizes[displayIndex]}`;
    }

    // Create and style the size display element
    function createSizeDisplay() {
        const display = document.createElement('div');
        display.id = 'hf-size-display';
        display.style.cssText = `
            position: fixed;
            bottom: 10px;
            right: 10px;
            background: rgba(0, 0, 0, 0.7);
            color: white;
            padding: 8px 12px;
            border-radius: 6px;
            font-size: 12px;
            font-family: monospace;
            z-index: 10000;
            pointer-events: auto;
            opacity: 0.8;
            transition: opacity 0.3s;
            box-shadow: 0 2px 4px rgba(0, 0, 0, 0.2);
            cursor: pointer;
        `;
        display.title = 'Click to copy download URLs of files in current directory';
        document.body.appendChild(display);
        return display;
    }

    // Generate download URLs for files in the current directory and copy to clipboard
    async function copyDirectoryFileUrls(repoName, directoryPath) {
        try {
            // Get the files in the current directory
            const url = directoryPath ?
                `https://huggingface.co/api/models/${repoName}/tree/main/${directoryPath}?recursive=false` :
                `https://huggingface.co/api/models/${repoName}/tree/main?recursive=false`;

            const response = await fetch(url);
            if (!response.ok) {
                throw new Error(`Failed to fetch directory contents: ${response.status}`);
            }

            const items = await response.json();

            // Filter only files (not directories)
            const filesOnly = items.filter(item => item.type === 'file');

            if (filesOnly.length === 0) {
                console.log('No files found in this directory.');
                return;
            }

            // Generate download URLs for each file
            const fileUrls = [];
            filesOnly.forEach(file => {
                const filePath = file.path;
                // Create direct download URL
                const downloadUrl = `https://huggingface.co/${repoName}/resolve/main/${filePath}`;
                fileUrls.push(downloadUrl);
            });

            // Copy to clipboard
            const urlsText = fileUrls.join('\n');
            navigator.clipboard.writeText(urlsText).then(() => {
                console.log('Download URLs copied to clipboard:', urlsText);

                // Create a temporary notification element
                const notification = document.createElement('div');
                notification.style.cssText = `
                    position: fixed;
                    bottom: 40px;
                    right: 10px;
                    background: rgba(0, 0, 0, 0.8);
                    color: #00ff00;
                    padding: 8px 12px;
                    border-radius: 6px;
                    font-size: 12px;
                    font-family: monospace;
                    z-index: 10001;
                    opacity: 1;
                    transition: opacity 0.5s;
                    pointer-events: none;
                    box-shadow: 0 2px 4px rgba(0, 0, 0, 0.3);
                `;
                notification.textContent = '✓ URLs copied to clipboard!';
                document.body.appendChild(notification);

                // Fade out and remove after 2 seconds
                setTimeout(() => {
                    notification.style.opacity = '0';
                    setTimeout(() => {
                        if (notification.parentNode) {
                            document.body.removeChild(notification);
                        }
                    }, 500);
                }, 2000);
            }).catch(err => {
                console.log('Failed to copy URLs to clipboard:', err);
            });

        } catch (error) {
            console.log('Error generating download URLs:', error);
        }
    }

    // Update the display with sizes
    function updateSizeDisplay(display, repoSize, dirSize) {
        const repoFormatted = formatBytes(repoSize);
        const dirFormatted = formatBytes(dirSize);

        display.textContent = `📁 ${dirFormatted} | 💾 ${repoFormatted}`;

        // Add hover effect for visibility
        display.style.opacity = '1';
        display.onmouseenter = () => display.style.opacity = '1';
        display.onmouseleave = () => display.style.opacity = '0.8';
    }

    // Main function to calculate and display sizes
    async function calculateAndDisplaySizes() {
        const repoInfo = getRepoInfo();
        if (!repoInfo) return;

        const display = document.getElementById('hf-size-display') || createSizeDisplay();

        // Add click event listener to copy download URLs
        display.onclick = async () => {
            await copyDirectoryFileUrls(repoInfo.repoFullName, repoInfo.path);
        };

        display.textContent = 'Loading...';

        try {
            // Calculate repository size (includes recursive directory sizes)
            const repoSize = await getRepoSize(repoInfo.repoFullName);

            // Calculate current directory size (only files in current directory)
            const dirSize = await getDirectorySizeCurrent(repoInfo.repoFullName, repoInfo.path);

            // Update display
            updateSizeDisplay(display, repoSize, dirSize);
        } catch (error) {
            console.log('Error calculating sizes:', error);
            display.textContent = 'Error loading sizes';
        }
    }

    // Run the calculation after a short delay
    setTimeout(() => {
        calculateAndDisplaySizes();
    }, 500);

})();