#!/usr/bin/env python3
"""
Reprocess a video file with MediaConvert using production settings.
Takes a hook ID and creates a MediaConvert job for the existing file in S3.
"""

import argparse
import os
import re
import sys
import time

import boto3
import requests
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"
MEDIACONVERT_ENDPOINT = "https://mediaconvert.us-east-1.amazonaws.com"
AWS_REGION = "us-east-1"
CLOUDFRONT_DISTRIBUTION_ID = os.getenv("CLOUDFRONT_DISTRIBUTION_ID")
CDN_BASE_URL = "https://cdn3.suno.ai"

# 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,
    get_metadata_params,
    get_output_path,
)


def get_hook_path_components(hook_id):
    """Get the nested path components for a hook ID"""
    if len(hook_id) < 4:
        hook_id = hook_id.ljust(4, "0")

    first_two = hook_id[:2]
    second_two = hook_id[2:4]
    return first_two, second_two


def find_video_file_in_s3(hook_id, s3_client):
    """Find the video file for a given hook ID in S3 using the nested path structure"""
    # Create nested path from hook ID
    first_two, second_two = get_hook_path_components(hook_id=hook_id)
    prefix = f"{first_two}/{second_two}/"

    # The file is always hook_<hookID>.mp4
    expected_filename = f"hook_{hook_id}.mp4"
    expected_key = f"{prefix}{expected_filename}"

    print(f"🔍 Looking for: s3://{SOURCE_BUCKET}/{expected_key}")

    # Check if the file exists
    try:
        s3_client.head_object(Bucket=SOURCE_BUCKET, Key=expected_key)
        return expected_key
    except s3_client.exceptions.ClientError as e:
        if e.response["Error"]["Code"] == "404":
            print(f"❌ Video file not found: s3://{SOURCE_BUCKET}/{expected_key}")
            return None
        else:
            print(f"❌ Error checking S3: {e}")
            return None


def create_mediaconvert_job(
    source_key,
    mediaconvert_client,
    s3_metadata=None,
):
    """Create a MediaConvert job using the production handler settings and existing metadata"""

    # Use the same output path logic as the handler
    output_path_base = get_output_path(source_key)
    output_path = f"s3://{DEST_BUCKET}/{output_path_base}"

    print("\n📹 Creating MediaConvert job:")
    print(f"  Input: s3://{SOURCE_BUCKET}/{source_key}")
    print(f"  Output: {output_path}")
    print(f"  Landscape: {s3_metadata.is_landscape}")
    print(f"  Environment: {s3_metadata.environment}")
    if s3_metadata.audio_bucket and s3_metadata.audio_key:
        print(f"  Audio: s3://{s3_metadata.audio_bucket}/{s3_metadata.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=s3_metadata.audio_bucket,
        audio_key=s3_metadata.audio_key,
        environment=s3_metadata.environment,
        callback_url=s3_metadata.callback_url,
        skip_callback=s3_metadata.skip_callback,
        is_landscape=s3_metadata.is_landscape,
    )

    # Submit job
    response = mediaconvert_client.create_job(**job_params)
    job_id = response["Job"]["Id"]
    print(f"\n✅ Job created: {job_id}")
    print(f"📍 Output will be at: {output_path}")
    return job_id, output_path


def check_job_status(job_id, mediaconvert_client):
    """Check the status of a MediaConvert job"""
    response = 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']}")

        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 wait_for_job(job_id, mediaconvert_client, 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 = check_job_status(
            job_id=job_id, mediaconvert_client=mediaconvert_client
        )

        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 fetch_m3u8_manifest(hook_id):
    """Fetch the HLS manifest and extract .ts file references"""
    first_two, second_two = get_hook_path_components(hook_id=hook_id)
    manifest_url = f"{CDN_BASE_URL}/{first_two}/{second_two}/hook_{hook_id}.m3u8"

    print(f"\n📥 Fetching original manifest from: {manifest_url}")

    try:
        response = requests.get(url=manifest_url, timeout=10)
        response.raise_for_status()

        # Extract .ts file references from the manifest
        ts_files = re.findall(r"hook_[a-f0-9\-]+_\d+\.ts", response.text)

        print(f"✅ Found {len(ts_files)} .ts files in manifest")
        if ts_files:
            print(f"   Sample files: {ts_files[:3]}")

        return {"manifest_url": manifest_url, "ts_files": set(ts_files), "content": response.text}
    except Exception as e:
        print(f"⚠️  Failed to fetch manifest: {e}")
        return None


def invalidate_cloudfront_cache(output_path):
    """Invalidate CloudFront cache for the processed files"""
    # Extract the path from s3://bucket/path format
    path = output_path.replace(f"s3://{DEST_BUCKET}/", "")
    invalidation_path = f"/{path}*"

    print(f"\n🔄 Invalidating CloudFront cache for: {invalidation_path}")

    cloudfront_client = boto3.client("cloudfront", region_name=AWS_REGION)

    try:
        response = cloudfront_client.create_invalidation(
            DistributionId=CLOUDFRONT_DISTRIBUTION_ID,
            InvalidationBatch={
                "Paths": {"Quantity": 1, "Items": [invalidation_path]},
                "CallerReference": str(time.time()),
            },
        )
        invalidation_id = response["Invalidation"]["Id"]
        print(f"✅ Cache invalidation created: {invalidation_id}")
        print(f"   Path: {invalidation_path}")
        return invalidation_id
    except Exception as e:
        print(f"⚠️  Failed to invalidate cache: {e}")
        return None


def wait_for_cache_update(original_manifest, hook_id, timeout=900):
    """Poll the manifest URL to detect when cache has been updated"""
    manifest_url = original_manifest["manifest_url"]
    original_ts_files = original_manifest["ts_files"]

    print(f"\n⏳ Polling for cache update (timeout: {timeout}s)...")
    print(f"   Watching: {manifest_url}")

    start_time = time.time()
    poll_interval = 20

    while True:
        elapsed = time.time() - start_time

        if elapsed > timeout:
            print(f"⏰ Timeout after {timeout}s - cache may not have updated yet")
            return False

        try:
            response = requests.get(url=manifest_url, timeout=10)
            response.raise_for_status()

            # Extract .ts files from current manifest
            current_ts_files = set(re.findall(r"hook_[a-f0-9\-]+_\d+\.ts", response.text))

            # Check if the manifest has changed
            if current_ts_files != original_ts_files:
                print(f"✅ Cache updated after {elapsed:.1f}s!")
                print(f"   Original files: {len(original_ts_files)}")
                print(f"   New files: {len(current_ts_files)}")

                # Show some examples of new files
                new_files = current_ts_files - original_ts_files
                if new_files:
                    print(f"   New .ts files: {list(new_files)[:3]}")

                return True

            print(
                f"   [{elapsed:.0f}s] Still serving old content, checking again in {poll_interval}s..."
            )

        except Exception as e:
            print(f"   ⚠️  Error fetching manifest: {e}")

        time.sleep(poll_interval)


def main():
    parser = argparse.ArgumentParser(
        description="Reprocess a video file with MediaConvert using production settings"
    )
    parser.add_argument("hook_id", help="Hook ID to reprocess")

    args = parser.parse_args()

    # Initialize clients
    s3_client = boto3.client("s3", region_name=AWS_REGION)
    mediaconvert_client = get_mediaconvert_client()

    print(f"\n🎬 Reprocessing hook ID: {args.hook_id}")

    # Fetch original manifest to track cache updates later
    original_manifest = fetch_m3u8_manifest(hook_id=args.hook_id)

    # Find the video file in S3
    source_key = find_video_file_in_s3(hook_id=args.hook_id, s3_client=s3_client)

    if not source_key:
        sys.exit(1)

    print(f"✅ Found video: s3://{SOURCE_BUCKET}/{source_key}")

    # Get existing metadata from S3
    print("\n📋 Reading existing metadata from S3...")
    s3_metadata = get_metadata_params(bucket=SOURCE_BUCKET, key=source_key)

    # Create MediaConvert job with existing settings
    job_id, output_path = create_mediaconvert_job(
        source_key=source_key,
        mediaconvert_client=mediaconvert_client,
        s3_metadata=s3_metadata,
    )

    # Wait for job to complete
    final_status, final_output_path = wait_for_job(
        job_id=job_id, mediaconvert_client=mediaconvert_client
    )

    print(f"\n📊 Final status: {final_status}")

    if final_status == "COMPLETE":
        print("\n✅ SUCCESS! Video processed without errors.")
        print(f"   Output: {final_output_path}")

        # Invalidate CloudFront cache for the processed files
        invalidate_cloudfront_cache(output_path=final_output_path)

        # Wait for cache to update with new content
        if original_manifest:
            wait_for_cache_update(
                original_manifest=original_manifest, hook_id=args.hook_id, timeout=900
            )
        else:
            print("⚠️  Skipping cache update check (original manifest not available)")
    elif final_status == "ERROR":
        print("\n❌ FAILED! MediaConvert encountered an error.")
        sys.exit(1)
    elif final_status == "TIMEOUT":
        print("\n⏰ Job did not complete within timeout period.")
        print(f"   To check status later: uv run python_tests/test_mediaconvert.py status {job_id}")
        sys.exit(1)

    print("\n✅ Done!")


if __name__ == "__main__":
    main()