#!/usr/bin/env python3
import argparse
import sys
import time
import subprocess
from pathlib import Path

from core import BaseTranscoder

# Add current directory to path to import core
sys.path.append(str(Path(__file__).resolve().parent))

VIDEO_DEFAULTS = {
    "hevc_nvenc": ["-preset", "p7", "-tune", "uhq", "-rc", "vbr_hq", "-cq", "24", "-spatial-aq", "1", "-multipass", "2", "-pix_fmt", "p010le"],
    "av1_nvenc": ["-preset", "p7", "-tune", "uhq", "-rc", "vbr", "-cq", "32", "-multipass", "2", "-pix_fmt", "p010le"],
    "libx265": ["-preset", "slow", "-crf", "24", "-pix_fmt", "yuv420p10le", "-x265-params", "aq-mode=3:aq-strength=0.8:psy-rd=1.0"],
    "libsvtav1": ["-preset", "4", "-crf", "35", "-pix_fmt", "yuv420p10le", "-svtav1-params", "tune=0:enable-qm=1"],
    "libvvenc": ["-preset", "medium", "-qpa", "0", "-qp", "24", "-pix_fmt", "yuv420p10le"],
}

AUDIO_DEFAULTS = {
    "fdkaac": ["-m", "4"],
    "opus": ["-c:a", "libopus", "-b:a", "128k", "-vbr", "on"],
}


class MediaTranscoder(BaseTranscoder):
    def __init__(self, args):
        super().__init__(args, VIDEO_DEFAULTS, AUDIO_DEFAULTS)

    def execute(self):
        print(f"--- Initializing Transcode Task ---")
        self.orig_info = self.get_media_info(self.input_path)
        if not self.orig_info:
            raise ValueError("Failed to read input file metadata.")

        start_time = time.time()
        try:
            if self.aencoder == "fdkaac":
                self._encode_with_external_audio()
            else:
                self._encode_standard()
        except subprocess.CalledProcessError as e:
            self._cleanup()
            self._send_notification(error=f"Transcode failed, exit code: {e.returncode}", extra_title="Transcode ")
            sys.exit(f"\nTranscode failed, exit code: {e.returncode}")

        self._cleanup()
        encode_time = time.time() - start_time

        new_info = self.get_media_info(self.output_path)
        vmaf_score, psnr_score = "N/A", "N/A"

        if self.run_metrics and new_info:
            vmaf_score, psnr_score = self._run_vmaf_psnr(new_info)

        log_content = self._generate_log_content(
            encode_time, new_info, vmaf_score, psnr_score
        )

        with open(self.log_path, "a", encoding="utf-8") as f:
            f.write(log_content)

        self._send_notification(encode_time=encode_time, new_info=new_info,
                                vmaf=vmaf_score, full_text=log_content, extra_title="Transcode ")

        print(f"\nTranscode completed! Time elapsed: {encode_time:.2f} s")
        print(f"Output file: {self.output_path}")

    def _encode_with_external_audio(self):
        print("\n[Audio] Extracting audio stream and encoding with external fdkaac...")
        temp_wav = self.out_dir / f"temp_{self.input_path.stem}.wav"
        temp_m4a = self.out_dir / f"temp_{self.input_path.stem}.m4a"
        self.temp_files.extend([temp_wav, temp_m4a])

        subprocess.run(["ffmpeg", "-y", "-i", str(self.input_path), "-vn", "-c:a", "pcm_s16le", str(temp_wav)], check=True)
        subprocess.run(["fdkaac"] + self.a_params + [str(temp_wav), "-o", str(temp_m4a)], check=True)

        print("\n[Video] Encoding video and muxing...")
        cmd = [
            "ffmpeg", "-y", "-i", str(self.input_path), "-i", str(temp_m4a),
            "-map", "0:v:0", "-map", "1:a:0", "-c:v", self.vencoder,
        ]
        if self.vf_string:
            cmd.extend(["-vf", self.vf_string])
        cmd.extend(self.v_params)
        cmd.extend(["-c:a", "copy", str(self.output_path)])
        subprocess.run(cmd, check=True)

    def _encode_standard(self):
        print("\n[Transcode] Processing video and audio simultaneously with ffmpeg...")
        cmd = ["ffmpeg", "-y", "-i", str(self.input_path), "-c:v", self.vencoder]
        if self.vf_string:
            cmd.extend(["-vf", self.vf_string])
        cmd.extend(self.v_params)
        cmd.extend(self.a_params)
        cmd.extend([str(self.output_path)])
        subprocess.run(cmd, check=True)


def main():
    parser = argparse.ArgumentParser(description="Configurable video transcoding script")
    parser.add_argument("-i", "--input", required=True, help="Input video file")
    parser.add_argument("-d", "--outdir", help="Output directory (default: same as input)")
    parser.add_argument("-cv", "--vencoder", choices=VIDEO_DEFAULTS.keys(), default="libsvtav1", help="Video encoder (default: libsvtav1)")
    parser.add_argument("-ca", "--aencoder", choices=AUDIO_DEFAULTS.keys(), default="opus", help="Audio encoder (default: opus)")
    parser.add_argument("--vargs", type=str, default="", help="Override default video encoding parameters")
    parser.add_argument("--aargs", type=str, default="", help="Override default audio encoding parameters")
    parser.add_argument("--fps", type=str, default="", help="Specify framerate (e.g., 60, 30000/1001)")
    parser.add_argument("--res", type=str, default="", help="Specify resolution, same format as scale filter (e.g., 1920:1080)")
    parser.add_argument("--metrics", action=argparse.BooleanOptionalAction, default=True, help="Run VMAF/PSNR after completion")
    parser.add_argument("--notify", choices=["none", "mail", "notify"], default="mail", help="Notification method after completion (default: mail)")

    try:
        transcoder = MediaTranscoder(parser.parse_args())
        transcoder.execute()
    except Exception as e:
        sys.exit(f"Error occurred: {e}")


if __name__ == "__main__":
    main()