pipeline.py

import subprocess
import os
import sys
import zipfile
import time
import re

# Configuration
SOURCE_VIDEO_URL = "https://www.youtube.com/watch?v=FtutLA63Cp8" # Bad Apple
SOURCE_FILENAME = "source_video.mp4"
ARCHIVE_FILENAME = "source_archive.zip"
ENCODED_VIDEO = "encoded_data.mp4"
DECODED_OUTPUT = "decoded_archive.dat"
RESTORED_FILENAME = "decoded_archive.zip"

def run_command(command, capture_output=True):
    print(f"Executing: {' '.join(command)}")
    result = subprocess.run(command, capture_output=capture_output, text=True)
    if result.stdout:
        print(result.stdout)
    if result.returncode != 0:
        print(f"Error executing command: {result.stderr}")
        return None
    return result.stdout

def main():
    # 1. Download source video
    print("--- Step 1: Downloading source video ---")
    if not os.path.exists(SOURCE_FILENAME):
        # Download at lower quality to save time/bandwidth
        run_command(["yt-dlp", "-f", "worst", "-o", SOURCE_FILENAME, SOURCE_VIDEO_URL], capture_output=False)
    else:
        print("Source video already exists.")

    # 2. Archive the video
    print("\n--- Step 2: Archiving the video ---")
    with zipfile.ZipFile(ARCHIVE_FILENAME, 'w', zipfile.ZIP_DEFLATED) as zipf:
        zipf.write(SOURCE_FILENAME)
    print(f"Created {ARCHIVE_FILENAME}")

    # 3. Encode archive to video
    print("\n--- Step 3: Encoding archive to video ---")
    run_command([sys.executable, "main.py", "encode", ARCHIVE_FILENAME, ENCODED_VIDEO, "--block-size", "16"], capture_output=False)

    # 4. Upload to YouTube
    print("\n--- Step 4: Uploading to YouTube ---")
    upload_output = run_command([sys.executable, "main.py", "upload", ENCODED_VIDEO, "--title", "Pipeline Test"])
    if not upload_output:
        print("Upload failed or requires manual interaction.")
        return

    # Extract Video ID from output
    video_id_match = re.search(r"Video ID: ([\w-]+)", upload_output)
    if not video_id_match:
        print("Could not find Video ID in upload output.")
        return
    
    video_id = video_id_match.group(1)
    video_url = f"https://www.youtube.com/watch?v={video_id}"
    print(f"Uploaded successfully! Video ID: {video_id}")

    # 5. Wait for processing (YouTube needs time to make it available for download)
    print("\n--- Step 5: Waiting for YouTube processing (120 seconds) ---")
    time.sleep(120)

    # 6. Download the encoded video back
    print("\n--- Step 6: Downloading encoded video from YouTube ---")
    DOWNLOADED_VIDEO = "downloaded_from_yt.mp4"
    # We try to get the best quality to ensure data integrity
    run_command(["yt-dlp", "-f", "bestvideo", "-o", DOWNLOADED_VIDEO, video_url], capture_output=False)

    # 7. Decode back to archive
    print("\n--- Step 7: Decoding back to archive ---")
    run_command([sys.executable, "main.py", "decode", DOWNLOADED_VIDEO, DECODED_OUTPUT, "--block-size", "16"], capture_output=False)

    # 8. Restore/Verify
    print("\n--- Step 8: Verifying results ---")
    # In this pipeline, we know it's a zip, but let's use the restore command
    run_command([sys.executable, "main.py", "restore", DECODED_OUTPUT], capture_output=False)
    
    # Compare original archive hash with restored one
    # Note: restore_format.py renames the file to restored_output.zip (or similar)
    # Based on our src/restore_format.py, it should be restored_output.zip
    
    restored_file = "restored_output.zip"
    if os.path.exists(restored_file):
        import hashlib
        def get_hash(fname):
            hash_md5 = hashlib.md5()
            with open(fname, "rb") as f:
                for chunk in iter(lambda: f.read(4096), b""):
                    hash_md5.update(chunk)
            return hash_md5.hexdigest()

        original_hash = get_hash(ARCHIVE_FILENAME)
        restored_hash = get_hash(restored_file)

        print(f"Original MD5: {original_hash}")
        print(f"Restored MD5: {restored_hash}")

        if original_hash == restored_hash:
            print("\nSUCCESS: The pipeline completed successfully! Data is intact.")
        else:
            print("\nFAILURE: Hashes do not match. Data corruption occurred.")
    else:
        print(f"Restored file {restored_file} not found.")

if __name__ == "__main__":
    main()