import os
import aiohttp
import asyncio
import json
import gzip
from rlottie_python import LottieAnimation
from urllib.parse import urlparse
from typing import List
from rich.progress import Progress, BarColumn, TextColumn, TimeRemainingColumn
from rich.console import Console

console = Console()


# ======================
#  UNIVERSAL TGS LOADER
# ======================
def load_lottie_auto(path: str):
    with open(path, "rb") as f:
        data = f.read()

    # MODE 1 → GZIP
    if data[:2] == b"\x1f\x8b":
        return LottieAnimation.from_tgs(path)

    # MODE 2 → Direct JSON
    try:
        return LottieAnimation.from_json(data.decode("utf-8"))
    except:
        pass

    # MODE 3 → Try decompress
    try:
        txt = gzip.decompress(data).decode("utf-8")
        return LottieAnimation.from_json(txt)
    except:
        pass

    # MODE 4 → Brutal fallback (SigStickBot)
    try:
        idx = data.find(b"{")
        if idx != -1:
            txt = data[idx:].decode("utf-8")
            return LottieAnimation.from_json(txt)
    except Exception as e:
        raise Exception(f"Brutal mode failed: {e}")

    raise Exception(f"Invalid TGS/JSON file: {path}")


# ======================
#  STICKER DOWNLOADER
# ======================
class TelegramStickerDownloader:
    def __init__(self, token: str, sticker_pack_link: str):
        self.token = token
        self.sticker_pack_link = sticker_pack_link
        self.sticker_set = self._extract_sticker_set()

        self.tgs_dir = f"tgs/{self.sticker_set}_Stickers"
        self.gif_dir = f"{self.sticker_set}_GIFs"

        os.makedirs(self.tgs_dir, exist_ok=True)
        os.makedirs(self.gif_dir, exist_ok=True)

    def _extract_sticker_set(self):
        return urlparse(self.sticker_pack_link).path.split('/')[-1]

    @staticmethod
    async def retry_request(func, *args, retries=3, delay=2, **kwargs):
        for attempt in range(retries):
            try:
                return await func(*args, **kwargs)
            except:
                if attempt == retries - 1:
                    raise
                await asyncio.sleep(delay)

    async def fetch_sticker_set_info(self, session: aiohttp.ClientSession):
        async def inner():
            url = f"https://api.telegram.org/bot{self.token}/getStickerSet?name={self.sticker_set}"
            async with session.get(url, timeout=10) as resp:
                res = await resp.json()
                if not res.get("ok"):
                    raise Exception(f"Failed to get sticker set info: {res}")
                return res["result"]["stickers"]
        return await TelegramStickerDownloader.retry_request(inner)

    async def download_file(self, session: aiohttp.ClientSession, file_id: str, save_path: str):
        async def inner():
            # get file info
            url = f"https://api.telegram.org/bot{self.token}/getFile?file_id={file_id}"
            async with session.get(url, timeout=10) as resp:
                res = await resp.json()
                if not res.get("ok"):
                    raise Exception(f"Failed to get file info: {file_id}")
                file_path = res["result"]["file_path"]

            # download file
            download_url = f"https://api.telegram.org/file/bot{self.token}/{file_path}"
            async with session.get(download_url, timeout=10) as r:
                data = await r.read()
                with open(save_path, "wb") as f:
                    f.write(data)

        return await TelegramStickerDownloader.retry_request(inner)

    async def convert_tgs_to_gif(self, tgs_path: str, gif_path: str):
        try:
            anim = load_lottie_auto(tgs_path)
            anim.save_animation(gif_path)
        except Exception as e:
            console.print(f"[red]❌ Failed to convert {tgs_path}: {e}[/red]")

    async def process_sticker(self, session, sticker, progress, task):
        safe_emoji = sticker["emoji"]
        unique_id = sticker["file_unique_id"]

        file_name = f"{safe_emoji}_{unique_id}.tgs"
        tgs_path = os.path.join(self.tgs_dir, file_name)
        gif_path = os.path.join(self.gif_dir, file_name.replace(".tgs", ".gif"))

        # Already created
        if os.path.exists(tgs_path) and os.path.exists(gif_path):
            progress.update(task, advance=1)
            return

        # Download (if needed)
        if not os.path.exists(tgs_path):
            try:
                await self.download_file(session, sticker["file_id"], tgs_path)
            except Exception as e:
                console.print(f"[red]❌ Download failed: {file_name} — {e}[/red]")
                progress.update(task, advance=1)
                return

        # Convert to GIF
        if not os.path.exists(gif_path):
            await self.convert_tgs_to_gif(tgs_path, gif_path)

        progress.update(task, advance=1)

    async def run(self, selected_indexes: List[int] = None):
        async with aiohttp.ClientSession() as session:
            try:
                stickers = await self.fetch_sticker_set_info(session)
            except Exception as e:
                console.print(f"[red]{e}[/red]")
                return

            if selected_indexes:
                stickers = [stickers[i] for i in selected_indexes]

            console.print(f"[bold blue]Sticker pack:[/bold blue] {self.sticker_set}")
            console.print(f"[bold green]Total stickers:[/bold green] {len(stickers)}\n")

            with Progress(
                TextColumn("[progress.description]{task.description}"),
                BarColumn(),
                "[progress.percentage]{task.percentage:>3.0f}%",
                TimeRemainingColumn(),
                console=console
            ) as progress:
                task = progress.add_task("[cyan]Processing...[/cyan]", total=len(stickers))

                await asyncio.gather(*[
                    self.process_sticker(session, s, progress, task)
                    for s in stickers
                ])
