import asyncio
import os
from dotenv import load_dotenv
import aiohttp
from telegram import Bot
from telegram.error import TelegramError
from telegram.ext import Application

async def download_single_pack(bot: Bot, session: aiohttp.ClientSession, sticker_url: str, allowed_types: list):
    """
    Asynchronously downloads specific file types from a sticker pack into its own folder.
    Includes improved logging for skipped files.
    """
    master_download_dir = "downloads"
    try:
        if 't.me/addstickers/' not in sticker_url:
            print(f"--> Invalid URL format: {sticker_url}. Skipping.")
            return

        pack_name = sticker_url.split('/')[-1]
        print(f"\n--- Processing pack: {pack_name} ---")

        pack_path = os.path.join(master_download_dir, pack_name)
        os.makedirs(pack_path, exist_ok=True)
        print(f"Saving to directory: ./{pack_path}/")

        sticker_set = await bot.get_sticker_set(pack_name)
        total_stickers = len(sticker_set.stickers)
        print(f"Found {total_stickers} stickers. Filtering for types: {', '.join(allowed_types)}")

        download_count = 0
        skipped_count = 0
        for i, sticker in enumerate(sticker_set.stickers):
            was_skipped = True # Ga ervan uit dat de sticker wordt overgeslagen, tenzij we iets downloaden

            # 1. Check voor primaire sticker (webp, tgs, webm)
            primary_extension = ".webp"
            if sticker.is_animated:
                primary_extension = ".tgs"
            elif sticker.is_video:
                primary_extension = ".webm"
            
            primary_type = primary_extension.strip('.')
            if primary_type in allowed_types:
                was_skipped = False
                file_to_download = await bot.get_file(sticker.file_id)
                file_name = f"{i+1:03d}{primary_extension}"
                file_path = os.path.join(pack_path, file_name)
                
                async with session.get(file_to_download.file_path) as response:
                    if response.status == 200:
                        content = await response.read()
                        with open(file_path, 'wb') as f:
                            f.write(content)
                        download_count += 1

            # 2. Check voor PNG thumbnail
            if 'png' in allowed_types and sticker.thumbnail:
                was_skipped = False
                thumb_to_download = await bot.get_file(sticker.thumbnail.file_id)
                thumb_name = f"{i+1:03d}_thumb.png"
                thumb_path = os.path.join(pack_path, thumb_name)

                async with session.get(thumb_to_download.file_path) as response:
                    if response.status == 200:
                        content = await response.read()
                        with open(thumb_path, 'wb') as f:
                            f.write(content)
                        download_count += 1
            
            # NIEUW: Verbeterde logging voor overgeslagen bestanden
            if was_skipped:
                skipped_count += 1
            
            # Update de status op de console
            print(f"\rProcessing... (sticker {i+1}/{total_stickers}, downloaded: {download_count}, skipped: {skipped_count})", end="", flush=True)

        print(f"\n✅ Download complete for pack: {pack_name}. Total files downloaded: {download_count}")

    except TelegramError as e:
        print(f"\nError for pack '{pack_name}': {e}")
    except aiohttp.ClientError as e:
        print(f"\nNetwork error while downloading for '{pack_name}': {e}")
    except Exception as e:
        print(f"\nAn unexpected error occurred with pack '{pack_name}': {e}")

async def main():
    """
    Main asynchronous function to run the sticker downloader script.
    """
    print("--- Telegram Sticker Pack Downloader ---")
    
    load_dotenv()
    bot_token = os.getenv("TELEGRAM_BOT_TOKEN")

    if not bot_token:
        print("\n[ERROR] Telegram Bot Token niet gevonden!")
        print("Maak een '.env' bestand aan en voeg 'TELEGRAM_BOT_TOKEN=jouw_token' toe.")
        return

    application = Application.builder().token(bot_token).build()
    bot = application.bot
    
    async with aiohttp.ClientSession() as session:
        try:
            bot_user = await bot.get_me()
            print(f"Bot '{bot_user.first_name}' initialized successfully.")
        except Exception as e:
            print(f"Error initializing bot. Is je token in .env correct? Details: {e}")
            return

        print("\nAvailable file types: webp (static), tgs (animated), webm (video), png (thumbnail)")
        type_input = input("Enter desired file types (comma-separated), or type 'all': ").strip().lower()

        # GEWIJZIGD: Correcte logica voor 'all' of lege input
        if not type_input or type_input == 'all':
            allowed_types = ['webp', 'tgs', 'webm', 'png']
            print("--> Downloading all available types.")
        else:
            allowed_types = [t.strip() for t in type_input.split(',')]
            print(f"--> Selected types: {', '.join(allowed_types)}")

        print("\nChoose an option:")
        print("1: Download a single sticker pack from a URL")
        print("2: Download multiple packs from a 'urls.txt' file")
        
        choice = input("Enter your choice (1 or 2): ").strip()

        if choice == '1':
            sticker_url = input("Enter the Sticker Pack URL: ").strip()
            if not sticker_url:
                print("Error: Sticker URL cannot be empty.")
            else:
                await download_single_pack(bot, session, sticker_url, allowed_types)

        elif choice == '2':
            try:
                with open('urls.txt', 'r') as f:
                    urls = [line.strip() for line in f if line.strip()]
                
                if not urls:
                    print("'urls.txt' is empty.")
                    return

                print(f"\nFound {len(urls)} URLs. Starting batch download...")
                for url in urls:
                    await download_single_pack(bot, session, url, allowed_types)

            except FileNotFoundError:
                print("\nError: 'urls.txt' not found.")
            except Exception as e:
                print(f"An error occurred while reading the file: {e}")

        else:
            print("Invalid choice. Please run the script again and enter 1 or 2.")

    print("\n--- Script finished ---")

if __name__ == '__main__':
    asyncio.run(main())