import math
import os
import shutil
import subprocess
from pathlib import Path
from requests import RequestException
from Core import Bootstrap


class MediaAssembler(Bootstrap):
    """
    Handles downloading, decrypting, and assembling media content such as video, audio, and subtitles.
    It uses aria2c for downloading segment files and Shaka Packager for decryption using Widevine keys.

    Attributes:
        aria2_path (Path): Path to the aria2c binary.
        shaka_packager_path (Path): Path to the Shaka Packager binary.
        video_track (list): List of video segment URLs.
        audio_tracks (dict): Dictionary of audio tracks with keys as key IDs and values as track metadata.
        subtitles_tracks (dict): Dictionary of subtitle tracks by language.
        media_id (str): Unique identifier for the media.
        media_dir (Path): Temporary directory for processing the media.
        aria2c_input_file (Path): File used to store aria2c input configuration.
    """

    def __init__(self, video_track: list, audio_tracks: dict, subtitles_tracks: dict, media_id: str):
        super().__init__()
        self.aria2_path = Path("bin/aria2c.exe").resolve()
        self.shaka_packager_path = Path("bin/packager").resolve()
        self.video_track = video_track
        self.audio_tracks = audio_tracks
        self.subtitles_tracks = subtitles_tracks
        self.media_id = media_id
        self.media_dir = Path(self.temp_dir) / media_id
        self.aria2c_input_file = Path("aria2c_input.txt")

    def call(self, widevine_keys, download_video: bool = False):
        """
        Executes the media assembly pipeline: downloads and decrypts video/audio, and saves subtitles.

        Args:
            widevine_keys (list): List of Widevine decryption keys in the format 'key_id:key'.
            download_video (bool): Whether to download and decrypt the video track.

        Returns:
            tuple: Paths to the decrypted video, decrypted audio, and the media directory.
        """
        try:
            if not self.video_track:
                return self._fail("Invalid video segments list")

            self.media_dir.mkdir(parents=True, exist_ok=True)

            video_decrypted = None

            if download_video:
                video_encrypted = self._download(self.video_track)
                video_decrypted = self._decrypt_files(widevine_keys, video_encrypted) if video_encrypted else None

            audio_decrypted = None
            for key_id, audio in self.audio_tracks.items():
                if key_id and (segments := audio.get("segments")):
                    lang = audio.get("lang", "")
                    audio_encrypted = self._download(segments, lang, "audio")
                    if audio_encrypted:
                        audio_decrypted = self._decrypt_files(widevine_keys, audio_encrypted, "audio")

            for lang, subtitle in self.subtitles_tracks.items():
                if lang and lang.lower() != "none":
                    self._save_subtitles(subtitle)

            return video_decrypted, audio_decrypted, self.media_dir

        except Exception as e:
            return self._fail(f"Error in call method: {e}", exception=True)

    def _decrypt_files(self, widevine_keys, encrypted_file: Path, media: str = "video") -> Path | None:
        """
        Decrypts a media file using Shaka Packager and provided Widevine keys.

        Args:
            widevine_keys (list): List of Widevine decryption keys in the format 'key_id:key'.
            encrypted_file (Path): Path to the encrypted media file.
            media (str): Type of media ('video' or 'audio').

        Returns:
            Path | None: Path to the decrypted file, or None on failure.
        """
        if not encrypted_file or not widevine_keys:
            return self._fail("Missing encrypted file or Widevine keys")

        decrypted_file = encrypted_file.with_name(encrypted_file.name.replace("encrypted", "decrypted"))
        if decrypted_file.exists():
            self.warning(f"Decrypted file {decrypted_file} already exists, continuing")
            encrypted_file.unlink(missing_ok=True)
            return decrypted_file

        keys_arg = ",".join(
            [f"key_id={kid}:key={key}" for kid, key in map(lambda x: x.split(":"), widevine_keys)]
        )

        try:
            self.warning(f"Starting decryption for {encrypted_file.name}")
            #os.environ["PATH"] = str(self.shaka_packager_path.parent) + os.pathsep + os.environ["PATH"]
            subprocess.run([
                str(self.shaka_packager_path),
                f"input={encrypted_file},stream={media},output={decrypted_file}",
                "--quiet",
                "--enable_raw_key_decryption",
                "--keys", keys_arg
            ], check=True)

            encrypted_file.unlink(missing_ok=True)
            self.info(f"Decryption completed: {decrypted_file.name}")
            return decrypted_file

        except subprocess.CalledProcessError as e:
            return self._fail(f"Shaka Packager failed: {e}")
        except Exception as e:
            return self._fail(f"Unexpected error during decryption: {e}")

    def _download(self, segments: list, lang: str = "", media: str = "video") -> Path | None:
        """
        Downloads and joins media segments using aria2c.

        Args:
            segments (list): List of segment URLs.
            lang (str): Language code (used for naming).
            media (str): Type of media ('video' or 'audio').

        Returns:
            Path | None: Path to the encrypted output file, or None on failure.
        """
        try:
            extension = "m4a" if media == "audio" else "mp4"
            suffix = f"_{lang}" if lang else ""
            output_file = self.media_dir / f"{media}{suffix}_encrypted.{extension}"

            if output_file.exists() or output_file.with_name(output_file.name.replace("encrypted", "decrypted")).exists():
                self.success(f"{output_file} already exists, skipping download")
                return output_file

            temp_dir = self.media_dir / media
            temp_dir.mkdir(parents=True, exist_ok=True)

            num_digits = max(1, int(math.log10(len(segments))) + 1)
            lines = []
            for i, url in enumerate(segments):
                lines.extend([
                    url,
                    f"\tout={i:0{num_digits}d}.mp4",
                    f"\tdir={temp_dir}"
                ])

            self.aria2c_input_file.write_text("\n".join(lines), encoding="utf-8")

            self.success(f"Downloading {media} {f'({lang})' if lang else ''} with aria2c...")
            subprocess.run([
                str(self.aria2_path),
                '--allow-overwrite=true',
                '--file-allocation=none',
                '--console-log-level=warn',
                '--download-result=hide',
                "--summary-interval=0",
                '-x16', '-j16', '-s16',
                '-i', str(self.aria2c_input_file)
            ])

            with open(output_file, "wb") as dest:
                for part in sorted(temp_dir.glob("*.mp4")):
                    with open(part, "rb") as src:
                        shutil.copyfileobj(src, dest)

            shutil.rmtree(temp_dir)
            self.success(f"Segments joined into {output_file}")
            return output_file

        except subprocess.CalledProcessError as e:
            self.error(f"Aria2c failed: {e}")
        except Exception as e:
            return self._fail(f"Download error: {e}")

    def _save_subtitles(self, subtitle: dict):
        """
        Downloads and saves subtitle files in ASS format.

        Args:
            subtitle (dict): Dictionary containing subtitle metadata including URL and language.
        """
        try:
            lang = subtitle.get("language")
            base_path = self.media_dir / f"subtitle_{lang}_decrypted.ass"
            path = base_path

            if path.exists():
                if self.get_item("default_language") == lang:
                    path = base_path.with_stem(f"{base_path.stem}_forced")
                    if path.exists():
                        self.success(f"Forced subtitle already exists: {path}")
                        return
                else:
                    self.success(f"Subtitle already exists: {path}")
                    return

            content = self.get(subtitle.get("url")).content
            with open(path, "wb") as f:
                f.write(content)

            self.success(f"Downloaded subtitle [{lang}] to {path}")

        except RequestException as e:
            self.error(self.show_exception(e))

    def _fail(self, message: str, exception: bool = False):
        """
        Logs an error message and returns None values to indicate failure.

        Args:
            message (str): Error message to log.
            exception (bool): If True, returns a tuple of Nones.

        Returns:
            None | tuple: None or (None, None) depending on the exception flag.
        """
        self.error(message)
        if exception:
            return None, None
        return None
