#!/usr/bin/env python3
"""
MediaConvert Test Script - Imports and uses actual media-convert-handler functions
Tests that transcoded videos work correctly with MediaConvert using production settings
"""

import argparse
import os
import sys
import time
from datetime import datetime
from pathlib import Path

import boto3
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# AWS credentials from .env take precedence if set
if os.getenv("AWS_ACCESS_KEY_ID"):
    print("🔐 Using AWS credentials from .env file")


# Mock the datadog_lambda module before importing the handler
class MockDatadogLambda:
    @staticmethod
    def lambda_metric(name, value, tags=None):
        # Just print or ignore for testing
        print(f"[Mock Datadog Metric] {name}: {value} (tags: {tags})")


# Create a fake datadog_lambda module
sys.modules["datadog_lambda"] = MockDatadogLambda()
sys.modules["datadog_lambda.metric"] = MockDatadogLambda

# Constants
SOURCE_BUCKET = "suno-media-sour"
DEST_BUCKET = "suno-media-dest"
INPUTS_FOLDER = "tests/inputs"
OUTPUTS_FOLDER = "tests/outputs"
MEDIACONVERT_ENDPOINT = "https://mediaconvert.us-east-1.amazonaws.com"
AWS_REGION = "us-east-1"

# Set required environment variables BEFORE importing the handler
os.environ["MEDIACONVERT_ENDPOINT"] = MEDIACONVERT_ENDPOINT
os.environ["SOURCE_BUCKET"] = SOURCE_BUCKET
os.environ["DESTINATION_BUCKET"] = DEST_BUCKET
os.environ["AWS_DEFAULT_REGION"] = AWS_REGION

# Add the lambda handler directory to the path
lambda_path = os.path.join(os.path.dirname(__file__), "../lambda/media-convert-handler")
sys.path.insert(0, lambda_path)

# Now import the handler (after env vars are set)
from media_convert_handler import (
    create_hls_job_params,
    get_mediaconvert_client,
)


class MediaConvertTester:
    def __init__(self, region="us-east-1"):
        self.region = region

        # Get the MediaConvert client using the handler's function
        self.mediaconvert_client = get_mediaconvert_client()
        self.s3_client = boto3.client("s3", region_name=region)

    def _get_account_id(self):
        """Get AWS account ID"""
        sts = boto3.client("sts")
        return sts.get_caller_identity()["Account"]

    def upload_file(self, local_file_path, metadata=None):
        """Upload a local file to S3 with optional metadata"""
        file_name = Path(local_file_path).name
        s3_key = f"{INPUTS_FOLDER}/{file_name}"
        s3_path = f"s3://{SOURCE_BUCKET}/{s3_key}"

        print(f"📤 Uploading {local_file_path} to {s3_path}")

        # Prepare upload arguments
        upload_args = {
            "Bucket": SOURCE_BUCKET,
            "Key": s3_key,
        }

        # Add metadata if provided
        if metadata:
            upload_args["Metadata"] = metadata
            print(f"   With metadata: {metadata}")

        # Upload file
        with open(local_file_path, "rb") as f:
            self.s3_client.put_object(Body=f, **upload_args)

        print("✅ Upload successful")
        return s3_path, s3_key

    def create_job(
        self,
        source_key,
        test_name="test",
        audio_bucket=None,
        audio_key=None,
        is_landscape=False,
        environment="staging",
        callback_url=None,
        skip_callback=True,
    ):
        """Create a MediaConvert job using the actual handler function"""

        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_path = f"s3://{DEST_BUCKET}/{OUTPUTS_FOLDER}/{test_name}_{timestamp}/"

        print("\n📹 Creating MediaConvert job:")
        print(f"  Name: {test_name}")
        print(f"  Input: s3://{SOURCE_BUCKET}/{source_key}")
        print(f"  Output: {output_path}")
        print(f"  Landscape: {is_landscape}")
        if audio_bucket and audio_key:
            print(f"  Audio: s3://{audio_bucket}/{audio_key}")

        # Use the actual create_hls_job_params function from the handler
        job_params = create_hls_job_params(
            source_bucket=SOURCE_BUCKET,
            source_key=source_key,
            destination_path=output_path,
            audio_bucket=audio_bucket,
            audio_key=audio_key,
            environment=environment,
            callback_url=callback_url,
            skip_callback=skip_callback,
            is_landscape=is_landscape,
        )

        # Submit job
        response = self.mediaconvert_client.create_job(**job_params)
        job_id = response["Job"]["Id"]
        print(f"✅ Job created: {job_id}")
        return job_id

    def check_job_status(self, job_id):
        """Check the status of a MediaConvert job"""
        response = self.mediaconvert_client.get_job(Id=job_id)
        job = response["Job"]

        status = job["Status"]
        print(f"\n📊 Job Status: {status}")

        if status == "COMPLETE":
            print("✅ Job completed successfully!")

            # Get output location
            output_path = job["Settings"]["OutputGroups"][0]["OutputGroupSettings"][
                "HlsGroupSettings"
            ]["Destination"]
            print(f"📍 Output: {output_path}")

            # Check for any warnings
            if "Warnings" in job:
                print(f"⚠️  Warnings: {job['Warnings']}")

            # Check if there were any timing issues reported
            print("\n🎬 MediaConvert handled the video successfully!")
            print("   This means no black frames were injected at segment boundaries.")

            return status, output_path

        elif status == "ERROR":
            print("❌ Job failed!")
            if "ErrorMessage" in job:
                print(f"Error: {job['ErrorMessage']}")
            if "ErrorCode" in job:
                print(f"Error Code: {job['ErrorCode']}")

        elif status == "PROGRESSING":
            percent = job.get("JobPercentComplete", 0)
            print(f"⏳ Processing... {percent}% complete")

        return status, None

    def concatenate_all_qualities(self, local_path):
        """Concatenate HLS segments for 720p only"""
        import subprocess

        # Look for 720p playlist only
        qualities = []
        hls_dir = os.path.join(local_path, "hls")

        # Search in both root and hls subdirectory
        search_dirs = [local_path]
        if os.path.exists(hls_dir):
            search_dirs.append(hls_dir)

        for search_dir in search_dirs:
            for f in os.listdir(search_dir):
                if f.endswith(".m3u8") and f != "index.m3u8":
                    # Only process 720p
                    if "_720p" in f:
                        qualities.append(("720p", os.path.join(search_dir, f)))

        if not qualities:
            print("⚠️  No quality playlists found")
            return []

        print("\n🔨 Concatenating 720p to MP4...")

        output_files = []
        for quality_name, playlist_file in qualities:
            output_mp4 = os.path.join(local_path, f"{quality_name}.mp4")

            print(f"\n   📹 Processing {quality_name}:")
            print(f"      Playlist: {os.path.basename(playlist_file)}")
            print(f"      Output: {output_mp4}")

            # Use ffmpeg to concatenate
            cmd = [
                "ffmpeg",
                "-i",
                playlist_file,
                "-c",
                "copy",
                "-movflags",
                "+faststart",
                "-y",  # Overwrite output
                "-loglevel",
                "error",  # Reduce verbosity
                output_mp4,
            ]

            try:
                subprocess.run(cmd, capture_output=True, text=True, check=True)

                # Get file size
                size_mb = os.path.getsize(output_mp4) / (1024 * 1024)
                print(f"      ✅ Size: {size_mb:.2f} MB")
                output_files.append(output_mp4)

            except subprocess.CalledProcessError as e:
                print(f"      ❌ Failed: {e}")
                if e.stderr:
                    print(f"         Error: {e.stderr}")

        if output_files:
            print(f"\n✅ Created {len(output_files)} MP4 files:")
            for f in output_files:
                print(f"   - {f}")

        return output_files

    def download_hls_output(self, s3_output_path, local_dir="outputs", concat_to_mp4=True):
        """Download HLS output files from S3 to local directory"""
        # Parse S3 path
        if not s3_output_path.startswith("s3://"):
            print(f"❌ Invalid S3 path: {s3_output_path}")
            return None

        path_parts = s3_output_path[5:].split("/", 1)
        bucket = path_parts[0]
        prefix = path_parts[1] if len(path_parts) > 1 else ""

        # Create local output directory
        os.makedirs(local_dir, exist_ok=True)

        # Create subdirectory for this job
        job_name = os.path.basename(prefix.rstrip("/"))
        local_path = os.path.join(local_dir, job_name)
        os.makedirs(local_path, exist_ok=True)

        # Create HLS subdirectory for m3u8 and ts files
        hls_path = os.path.join(local_path, "hls")
        os.makedirs(hls_path, exist_ok=True)

        print("\n📥 Downloading HLS output files (720p only)...")
        print(f"   From: {s3_output_path}")
        print(f"   To: {local_path}/")

        # List and download only 720p files
        paginator = self.s3_client.get_paginator("list_objects_v2")
        pages = paginator.paginate(Bucket=bucket, Prefix=prefix)

        downloaded_files = []
        hls_files = []
        for page in pages:
            if "Contents" not in page:
                continue

            for obj in page["Contents"]:
                key = obj["Key"]
                # Get relative path from prefix
                relative_path = key[len(prefix) :].lstrip("/")
                if not relative_path:
                    continue

                # Skip non-720p files (but always download playlists)
                if relative_path.endswith(".ts") and "_720p" not in relative_path:
                    continue
                elif relative_path.endswith(".m3u8"):
                    # Only download index.m3u8 and 720p playlist
                    if "index" not in relative_path and "_720p" not in relative_path:
                        continue

                # Put HLS files in subfolder
                if relative_path.endswith((".m3u8", ".ts")):
                    local_file = os.path.join(hls_path, relative_path)
                    hls_files.append(local_file)
                else:
                    local_file = os.path.join(local_path, relative_path)

                # Create subdirectories if needed
                os.makedirs(os.path.dirname(local_file), exist_ok=True)

                # Download file (show only first few and last ts file)
                if relative_path.endswith(".ts"):
                    # Extract segment number
                    import re

                    match = re.search(r"(\d+)\.ts$", relative_path)
                    if match:
                        segment_num = int(match.group(1))
                        if segment_num <= 2 or segment_num % 10 == 0:
                            print(f"   📄 {relative_path}")
                else:
                    print(f"   📄 {relative_path}")

                self.s3_client.download_file(bucket, key, local_file)
                downloaded_files.append(local_file)

        if downloaded_files:
            print(f"\n✅ Downloaded {len(downloaded_files)} files")
            print(f"   HLS files: {len(hls_files)} in {hls_path}/")

            # Concatenate to MP4 if requested
            if concat_to_mp4:
                mp4_files = self.concatenate_all_qualities(local_path)
                if mp4_files:
                    print("\n🎬 To play the videos:")
                    for mp4 in mp4_files:
                        basename = os.path.basename(mp4)
                        print(f"   ffplay {mp4}  # {basename}")

                    # Suggest the 720p by default
                    default_mp4 = os.path.join(local_path, "720p.mp4")
                    if os.path.exists(default_mp4):
                        print("\n   Or open with QuickTime:")
                        print(f"   open {default_mp4}")
                    return mp4_files[0] if mp4_files else None

            # Find the master playlist
            master_playlist = os.path.join(hls_path, "index.m3u8")
            if os.path.exists(master_playlist):
                print("\n🎬 To play the HLS stream:")
                print(f"   ffplay {master_playlist}")
                return master_playlist
            else:
                # Look for any .m3u8 file
                for f in hls_files:
                    if f.endswith(".m3u8"):
                        print("\n🎬 To play the HLS stream:")
                        print(f"   ffplay {f}")
                        return f
        else:
            print("⚠️  No files found to download")
            return None

    def wait_for_job(self, job_id, timeout=600):
        """Wait for a job to complete"""
        print(f"\n⏳ Waiting for job {job_id} to complete...")

        start_time = time.time()
        output_path = None
        while True:
            status, output_path = self.check_job_status(job_id)

            if status in ["COMPLETE", "ERROR", "CANCELED"]:
                return status, output_path

            if time.time() - start_time > timeout:
                print(f"⏰ Timeout waiting for job after {timeout} seconds")
                return "TIMEOUT", None

            time.sleep(10)

    def test_file(
        self,
        local_file_path,
        test_name=None,
        is_landscape=False,
        with_audio=False,
        wait=True,
        concat_to_mp4=True,
        refresh_file=False,
    ):
        """Test a local file end-to-end"""
        if not os.path.exists(local_file_path):
            print(f"❌ File not found: {local_file_path}")
            return None

        if test_name is None:
            test_name = Path(local_file_path).stem

        print(f"\n🧪 Testing file: {local_file_path}")
        print(f"   Test name: {test_name}")

        # Check if file already exists in S3
        file_name = Path(local_file_path).name
        s3_key = f"{INPUTS_FOLDER}/{file_name}"

        metadata = {
            "environment": "testing",
            "skip-callback": "true",
            "is-landscape": str(is_landscape).lower(),
        }
        if refresh_file:
            print("📤 Refreshing file in S3, deleting existing file")
            self.s3_client.delete_object(Bucket=SOURCE_BUCKET, Key=s3_key)
            _, s3_key = self.upload_file(local_file_path, metadata)
        else:
            try:
                self.s3_client.head_object(Bucket=SOURCE_BUCKET, Key=s3_key)
                print(f"✅ File already exists in S3: s3://{SOURCE_BUCKET}/{s3_key}")
                print("   Skipping upload, using existing file")
            except self.s3_client.exceptions.ClientError as e:
                if e.response["Error"]["Code"] == "404":
                    print("📤 File not found in S3, uploading...")
                    _, s3_key = self.upload_file(local_file_path, metadata)
                else:
                    print(f"❌ Error checking S3: {e}")
                    return None

        # Add audio if requested (you could point to a real audio file)
        audio_bucket = None
        audio_key = None
        if with_audio:
            # You could upload an audio file here or use an existing one
            # For now, we'll skip audio
            pass

        # Create job using the handler function
        job_id = self.create_job(
            source_key=s3_key,
            test_name=test_name,
            audio_bucket=audio_bucket,
            audio_key=audio_key,
            is_landscape=is_landscape,
            environment="testing",
            skip_callback=True,
        )

        # Wait for completion if requested
        if wait:
            final_status, output_path = self.wait_for_job(job_id)
            print(f"\n📊 Final status: {final_status}")

            if final_status == "COMPLETE":
                print("\n✅ SUCCESS! Video processed without errors.")
                print("   MediaConvert successfully handled the file.")
                print("   No black frames should be present at segment boundaries.")

                # Download the output files
                if output_path:
                    local_playlist = self.download_hls_output(
                        output_path, concat_to_mp4=concat_to_mp4
                    )
                    return job_id, local_playlist

            elif final_status == "ERROR":
                print("\n❌ FAILED! MediaConvert encountered an error.")
                print("   This might indicate timing issues with the transcoded file.")
                print("   Check if the video has start_time=0.000")

        return job_id, None


def main():
    parser = argparse.ArgumentParser(
        description="Test MediaConvert using actual production handler code"
    )

    subparsers = parser.add_subparsers(dest="command", help="Command to run")

    # Test file command
    test_parser = subparsers.add_parser("test", help="Test a local video file")
    test_parser.add_argument("file", help="Path to video file")
    test_parser.add_argument("--name", help="Test name (default: filename)")
    test_parser.add_argument("--landscape", action="store_true", help="Video is landscape")
    test_parser.add_argument("--with-audio", action="store_true", help="Include audio track")
    test_parser.add_argument("--no-wait", action="store_true", help="Don't wait for completion")
    test_parser.add_argument("--no-concat", action="store_true", help="Don't concatenate to MP4")
    test_parser.add_argument("--refresh-file", action="store_true", help="Refresh file in S3")

    # Check status command
    status_parser = subparsers.add_parser("status", help="Check job status")
    status_parser.add_argument("job_id", help="MediaConvert job ID")

    # Test multiple files
    batch_parser = subparsers.add_parser("test-batch", help="Test multiple files or folder")
    batch_parser.add_argument("path", help="Video files or folder to test")
    batch_parser.add_argument("--refresh-file", action="store_true", help="Refresh file in S3")
    batch_parser.add_argument("--no-concat", action="store_true", help="Don't concatenate to MP4")

    args = parser.parse_args()

    if not args.command:
        parser.print_help()
        sys.exit(1)

    # Create tester
    tester = MediaConvertTester()

    if args.command == "test":
        tester.test_file(
            args.file,
            test_name=args.name,
            is_landscape=args.landscape,
            with_audio=args.with_audio,
            wait=not args.no_wait,
            concat_to_mp4=not args.no_concat if hasattr(args, "no_concat") else True,
            refresh_file=args.refresh_file,
        )

    elif args.command == "test-batch":
        import concurrent.futures
        from threading import Lock

        # Collect files to process
        files_to_test = []
        path = Path(args.path)

        if path.is_file():
            # Single file provided
            files_to_test.append(path)
        elif path.is_dir():
            # Folder provided - find all MP4 files
            files_to_test = list(path.glob("*.mp4"))
        else:
            print(f"❌ Path not found: {args.path}")
            sys.exit(1)

        if not files_to_test:
            print(f"❌ No MP4 files found in: {args.path}")
            sys.exit(1)

        print(f"📦 Submitting {len(files_to_test)} files in parallel...")
        job_ids = []
        lock = Lock()

        def submit_job(file_path):
            test_name = file_path.stem
            with lock:
                print(f"📤 Submitting {test_name}...")
            try:
                job_id, _ = tester.test_file(
                    str(file_path),
                    test_name=test_name,
                    wait=False,
                    refresh_file=args.refresh_file if hasattr(args, "refresh_file") else False,
                )
                with lock:
                    print(f"✅ Submitted {test_name}: Job ID {job_id}")
                return test_name, job_id
            except Exception as e:
                with lock:
                    print(f"❌ Failed to submit {test_name}: {e}")
                return test_name, None

        # Submit all jobs in parallel
        with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
            futures = [executor.submit(submit_job, file_path) for file_path in files_to_test]
            for future in concurrent.futures.as_completed(futures):
                test_name, job_id = future.result()
                if job_id:
                    job_ids.append((test_name, job_id))

        if job_ids:
            print(f"\n⏳ Submitted {len(job_ids)} jobs. Waiting for completion...")
            concat_to_mp4 = not args.no_concat if hasattr(args, "no_concat") else True

            for name, job_id in job_ids:
                print(f"\n📄 {name}:")
                final_status, output_path = tester.wait_for_job(job_id)
                print(f"📊 Status: {final_status}")

                if final_status == "COMPLETE" and output_path:
                    # Download and optionally concatenate the output
                    local_playlist = tester.download_hls_output(
                        output_path, concat_to_mp4=concat_to_mp4
                    )
                    if local_playlist:
                        print(f"✅ Downloaded: {local_playlist}")
                elif final_status == "ERROR":
                    print(f"❌ Failed: {name}")


if __name__ == "__main__":
    main()