# ffmpeg_service.py

from flask import Flask, request, Response, stream_with_context
import requests
import subprocess
import tempfile
import os
import shutil
import time

ffmpeg_app = Flask(__name__)

@ffmpeg_app.route('/convert', methods=['GET'])
def convert_media():
    """
    Downloads the raw media file from the provided URL and converts it using FFmpeg.
    """
    media_url = request.args.get('url')
    target_format = request.args.get('format') # e.g., 'mp3', 'm4a', 'wav'

    if not media_url or not target_format:
        return "Missing 'url' or 'format' parameter", 400

    temp_dir = tempfile.mkdtemp()
    
    # Use the original format's extension for the input file to help FFmpeg
    input_ext = media_url.split('.')[-1].split('?')[0]
    input_file_path = os.path.join(temp_dir, f"input.{input_ext}")
    output_file_path = os.path.join(temp_dir, f"output.{target_format}")

    try:
        start_time = time.time()
        print(f"[{time.strftime('%H:%M:%S')}] Starting conversion for {target_format} from URL: {media_url[:100]}...")

        # 1. Download the raw media file to a temporary file
        with requests.get(media_url, stream=True, timeout=120) as r:
            r.raise_for_status()
            content_length = int(r.headers.get('content-length', 0))
            print(f"[{time.strftime('%H:%M:%S')}] Raw media size: {content_length/1024/1024:.2f} MB")
            with open(input_file_path, 'wb') as f:
                for chunk in r.iter_content(chunk_size=8192):
                    f.write(chunk)
        
        download_time = time.time()
        print(f"[{time.strftime('%H:%M:%S')}] Download complete. Time taken: {download_time - start_time:.2f} seconds.")

        # 2. Run FFmpeg conversion
        command = [
            'ffmpeg',
            '-y', # Overwrite output files without asking
            '-i', input_file_path,
            # Audio-specific options for quality and codec
            '-q:a', '0', # Variable bitrate, highest quality
            '-map', 'a',  # Only include audio streams
            # Codec selection based on format
            '-c:a', 'libmp3lame' if target_format == 'mp3' else 
                    ('pcm_s16le' if target_format == 'wav' else 'aac'),
            '-f', target_format,
            output_file_path
        ]
        
        subprocess.run(command, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, timeout=180) # 3-minute timeout
        
        conversion_time = time.time()
        print(f"[{time.strftime('%H:%M:%S')}] FFmpeg conversion complete. Time taken: {conversion_time - download_time:.2f} seconds.")

        # 3. Stream the converted file back
        def stream_output_file():
            with open(output_file_path, 'rb') as f:
                chunk = True
                while chunk:
                    chunk = f.read(8192)
                    yield chunk

        mime_type = f'audio/{target_format}' if target_format != 'm4a' else 'audio/mp4'
        
        return Response(stream_with_context(stream_output_file()), 
                        mimetype=mime_type)

    except subprocess.CalledProcessError as e:
        print(f"[{time.strftime('%H:%M:%S')}] FFmpeg command failed with error: {e.stderr.decode()}")
        return f"Conversion failed: {e.stderr.decode()}", 500
    except requests.exceptions.Timeout:
        return "Media download or conversion timed out.", 504
    except Exception as e:
        print(f"[{time.strftime('%H:%M:%S')}] An error occurred in the FFmpeg service: {e}")
        return f"Internal Server Error: {e}", 500
    
    finally:
        # 4. Cleanup temporary files
        if os.path.exists(temp_dir):
            shutil.rmtree(temp_dir)
            print(f"[{time.strftime('%H:%M:%S')}] Cleaned up temp directory: {temp_dir}")

@ffmpeg_app.route('/', methods=['GET'])
def index():
    return "FFmpeg Conversion Service is running."

if __name__ == '__main__':
    # For local development/testing of the FFmpeg service
    # NOTE: When deployed via Docker/Gunicorn on Hugging Face, this block is ignored.
    ffmpeg_app.run(host='0.0.0.0', port=7860, debug=True)