import asyncio
import json
import mimetypes
import os
import shutil
import tempfile
import zipfile
from urllib.parse import urlparse

import yt_dlp
from channels.generic.websocket import AsyncWebsocketConsumer
from django.conf import settings
from django.core import signing
from django.utils.text import slugify

YDL_BASE = {
    "quiet": True,
    "no_check_certificates": True,
    "js_runtimes": {"node": {}},
    "remote_components": {"ejs:github"},
}

TOKEN_TTL = 600  # seconds until signed download token expires


class DownloaderConsumer(AsyncWebsocketConsumer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.tmpdir = None

    async def connect(self):
        await self.accept()

    async def disconnect(self, code):
        if self.tmpdir and os.path.exists(self.tmpdir):
            shutil.rmtree(self.tmpdir, ignore_errors=True)
            self.tmpdir = None

    async def receive(self, text_data):
        try:
            params = json.loads(text_data)
        except (json.JSONDecodeError, TypeError):
            await self._send({"type": "error", "message": "Invalid JSON"})
            return

        url = params.get("url")
        if not url:
            await self._send({"type": "error", "message": "URL is required"})
            return

        ext = params.get("ext", "mp4")
        if not ext or not isinstance(ext, str):
            await self._send({"type": "error", "message": "Invalid extension"})
            return

        os.makedirs(settings.DOWNLOADER_TMP_DIR, exist_ok=True)
        self.tmpdir = tempfile.mkdtemp(prefix="dl_ws_", dir=settings.DOWNLOADER_TMP_DIR)

        await self._send({"type": "status", "message": "Starting…"})

        loop = asyncio.get_event_loop()

        def progress_hook(d):
            status = d.get("status")
            if status == "downloading":
                percent = (d.get("_percent_str") or "").strip()
                speed = (d.get("_speed_str") or "").strip()
                eta = (d.get("_eta_str") or "").strip()
                downloaded = d.get("downloaded_bytes", 0)
                total = d.get("total_bytes") or d.get("total_bytes_estimate") or 0
                pct_num = round(downloaded / total * 100, 1) if total else None
                asyncio.run_coroutine_threadsafe(
                    self._send({
                        "type": "progress",
                        "message": f"Downloading: {percent} at {speed}  ETA {eta}",
                        "percent": pct_num,
                        "speed": speed,
                        "eta": eta,
                    }),
                    loop,
                )
            elif status == "finished":
                asyncio.run_coroutine_threadsafe(
                    self._send({"type": "status", "message": "Post-processing…"}),
                    loop,
                )

        try:
            result = await asyncio.to_thread(
                self._run_download,
                url,
                ext,
                params.get("video_quality"),
                params.get("audio_quality"),
                params.get("selected_videos"),
                params.get("subtitles"),
                bool(params.get("embed_subtitles", False)),
                bool(params.get("embed_thumbnail", False)),
                bool(params.get("extract_audio", False)),
                params.get("cookies"),
                progress_hook,
            )
        except Exception as exc:
            if self.tmpdir:
                shutil.rmtree(self.tmpdir, ignore_errors=True)
                self.tmpdir = None
            await self._send({"type": "error", "message": str(exc)})
            return

        # Sign a token containing the file info — no DB needed.
        token = signing.dumps(
            {
                "file_path": result["file_path"],
                "filename": result["filename"],
                "content_type": result["content_type"],
                "file_size": result["file_size"],
                "tmpdir": self.tmpdir,
            },
            salt="downloader-file-token",
        )
        # Consumer no longer owns cleanup — the HTTP serve view will clean up.
        self.tmpdir = None

        await self._send({
            "type": "done",
            "token": token,
            "filename": result["filename"],
            "file_size": result["file_size"],
        })

    # ------------------------------------------------------------------
    def _run_download(self, url, ext, video_quality, audio_quality,
                      selected_videos, subtitles, embed_subtitles,
                      embed_thumbnail, extract_audio, cookies, progress_hook):
        """Runs synchronously inside a thread pool worker."""
        tmpdir = self.tmpdir

        # Format selector
        if video_quality and audio_quality:
            fmt = f"bv[height<={video_quality}]+ba[abr<={audio_quality}]/b"
        elif video_quality:
            fmt = f"bv[height<={video_quality}]+ba/b"
        elif audio_quality:
            fmt = f"bv+ba[abr<={audio_quality}]/b"
        else:
            fmt = "b/bv+ba"

        ydl_opts = {
            **YDL_BASE,
            "format": fmt,
            "merge_output_format": ext,
            "max_filesize": settings.DOWNLOADER_MAX_SIZE_BYTES,
            "postprocessors": [],
            "progress_hooks": [progress_hook],
        }

        if cookies:
            cookie_file = os.path.join(tmpdir, "cookies.txt")
            with open(cookie_file, "w") as f:
                f.write(cookies)
            ydl_opts["cookiefile"] = cookie_file

        if subtitles:
            if subtitles.lower() == "all":
                ydl_opts.update({
                    "writesubtitles": True,
                    "writeautomaticsub": True,
                    "subtitleslangs": ["all"],
                })
            else:
                ydl_opts.update({
                    "writesubtitles": True,
                    "subtitleslangs": [l.strip() for l in subtitles.split(",")],
                })

        if embed_subtitles and subtitles and ext in ("mkv", "mp4"):
            ydl_opts["postprocessors"].append({"key": "FFmpegEmbedSubtitle"})

        if embed_thumbnail:
            ydl_opts["writethumbnail"] = True
            ydl_opts["postprocessors"].append({"key": "EmbedThumbnail"})

        if extract_audio:
            ydl_opts["postprocessors"].append({
                "key": "FFmpegExtractAudio",
                "preferredcodec": ext if ext in ("mp3", "m4a", "opus", "vorbis", "wav") else "mp3",
            })
        else:
            ydl_opts["postprocessors"].append({"key": "FFmpegVideoRemuxer", "preferedformat": ext})

        # Probe to detect playlist
        probe_opts = {**YDL_BASE, "extract_flat": False}
        with yt_dlp.YoutubeDL(probe_opts) as ydl:
            info = ydl.extract_info(url, download=False)

        is_playlist = "entries" in info and info.get("entries") is not None

        if is_playlist:
            if selected_videos:
                ydl_opts["playlist_items"] = ",".join(str(n) for n in selected_videos)
            ydl_opts["outtmpl"] = os.path.join(tmpdir, "%(playlist_index)02d - %(title)s.%(ext)s")
            with yt_dlp.YoutubeDL(ydl_opts) as ydl:
                ydl.download([url])

            zip_path = os.path.join(tmpdir, "playlist.zip")
            files = [
                (fn, os.path.join(tmpdir, fn))
                for fn in os.listdir(tmpdir)
                if fn not in ("playlist.zip", "cookies.txt")
                and not fn.startswith(".")
                and os.path.isfile(os.path.join(tmpdir, fn))
            ]
            if not files:
                raise RuntimeError("No files were downloaded from the playlist")

            with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
                for fn, fp in files:
                    zf.write(fp, fn)

            playlist_title = slugify(info.get("title", "playlist"))
            return {
                "file_path": zip_path,
                "filename": f"{playlist_title}.zip",
                "content_type": "application/zip",
                "file_size": os.path.getsize(zip_path),
            }
        else:
            ydl_opts["outtmpl"] = os.path.join(tmpdir, "download.%(ext)s")
            with yt_dlp.YoutubeDL(ydl_opts) as ydl:
                info = ydl.extract_info(url, download=True)
                base = ydl.prepare_filename(info)
                file_path = base if base.endswith(f".{ext}") else os.path.splitext(base)[0] + f".{ext}"

            safe_title = slugify(info.get("title") or "video")
            filename = f"{safe_title}.{ext}"
            content_type = mimetypes.guess_type(filename)[0] or "application/octet-stream"
            return {
                "file_path": file_path,
                "filename": filename,
                "content_type": content_type,
                "file_size": os.path.getsize(file_path) if os.path.exists(file_path) else 0,
            }

    async def _send(self, data: dict):
        await self.send(text_data=json.dumps(data))
