from flask import Flask, Response, request
import subprocess
import os
import argparse
import time
import threading

app = Flask(__name__)

BASE_STORAGE_PATH = "/vPath"  # Root directory for MKV files
DEBUG_MODE = False  # Default: Debugging on

class bcolors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKCYAN = '\033[96m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'

@app.route('/convert/vpath/<path:filepath>.webm')
def convert(filepath):
    # Log the incoming request
    if DEBUG_MODE:
        print(f"{bcolors.OKGREEN}[DEBUG] Received request: {request.url}{bcolors.ENDC}")
        print(f"{bcolors.OKGREEN}[DEBUG] Mapped file path: {os.path.join(BASE_STORAGE_PATH, filepath + '.mkv')}{bcolors.ENDC}")

    mkv_file = os.path.join(BASE_STORAGE_PATH, filepath + ".mkv")

    # Check if the file exists before proceeding
    if not os.path.isfile(mkv_file):
        if DEBUG_MODE:
            print(f"{bcolors.FAIL}[ERROR] File not found: {mkv_file}{bcolors.ENDC}")
        return "File not found", 404

    def generate():
        # Use ffprobe to get the audio codec info
        probe_cmd = ["ffprobe", "-v", "error", "-select_streams", "a:0", "-show_entries", "stream=codec_name", "-of", "default=nw=1:nk=1", mkv_file]
        try:
            audio_codec_info = subprocess.check_output(probe_cmd, stderr=subprocess.PIPE).decode('utf-8').strip()
            if DEBUG_MODE:
                print(f"{bcolors.OKGREEN}[DEBUG] Audio codec: {audio_codec_info}{bcolors.ENDC}")
        except subprocess.CalledProcessError as e:
            if DEBUG_MODE:
                print(f"{bcolors.FAIL}[ERROR] Audio probe failed: {e}{bcolors.ENDC}")
            audio_codec_info = "unknown"

        # Use ffprobe to get the video information
        probe_cmd = ["ffprobe", "-v", "error", "-select_streams", "v:0", "-show_entries", "stream=codec_name", "-of", "default=nw=1:nk=1", mkv_file]
        try:
            video_codec_info = subprocess.check_output(probe_cmd, stderr=subprocess.PIPE).decode('utf-8').strip()
            if DEBUG_MODE:
                print(f"{bcolors.WARNING}[DEBUG] Video codec: {video_codec_info}{bcolors.ENDC}")
        except subprocess.CalledProcessError as e:
            if DEBUG_MODE:
                print(f"{bcolors.FAIL}[ERROR] Video probe failed: {e}{bcolors.ENDC}")
            video_codec_info = "unknown"

        # Get available encoders in this FFmpeg installation
        available_encoders = []
        try:
            encoders_output = subprocess.check_output(["ffmpeg", "-encoders"], stderr=subprocess.PIPE).decode('utf-8')
            if DEBUG_MODE:
                print("{bcolors.WARNING}[DEBUG] Checking available encoders{bcolors.ENDC}")
            for line in encoders_output.split('\n'):
                if ' V' in line or ' A' in line:  # Video or Audio encoder
                    parts = line.split()
                    if len(parts) > 2:
                        available_encoders.append(parts[1])
            if DEBUG_MODE:
                print(f"{bcolors.WARNING}[DEBUG] Found encoders: {', '.join(available_encoders)}{bcolors.ENDC}")
        except Exception as e:
            if DEBUG_MODE:
                print(f"{bcolors.FAIL}[ERROR] Failed to get encoder list: {str(e)}{bcolors.ENDC}")

        # Build FFmpeg command with fallback options
        ffmpeg_cmd = ["ffmpeg", "-i", mkv_file]
        
        # Add video options based on codec and available encoders
        if video_codec_info in ["vp8", "vp9"]:
            ffmpeg_cmd.extend(["-c:v", "copy"])
        elif "libvpx" in available_encoders:
            ffmpeg_cmd.extend(["-c:v", "libvpx", "-b:v", "1M"])
        elif "vp8" in available_encoders:
            ffmpeg_cmd.extend(["-c:v", "vp8", "-b:v", "1M"])
        else:
            # Last resort fallback
            ffmpeg_cmd.extend(["-c:v", "libx264", "-b:v", "1M"])
        
        # Add audio options based on codec and available encoders
        # We dont need to transcode Audio if its already in opus or vorbis format.
        if audio_codec_info in ["opus", "vorbis"]:
            ffmpeg_cmd.extend(["-c:a", "copy"])
        elif "libvorbis" in available_encoders:
            ffmpeg_cmd.extend(["-c:a", "libvorbis", "-b:a", "128k"])
        elif "vorbis" in available_encoders:
            ffmpeg_cmd.extend(["-c:a", "vorbis", "-b:a", "128k"])
        elif "aac" in available_encoders:
            ffmpeg_cmd.extend(["-c:a", "aac", "-b:a", "128k"])
        else:
            # Try to copy audio as a last resort
            ffmpeg_cmd.extend(["-c:a", "copy"])
        
        # Add output format and destination - simplified with just essential options
        ffmpeg_cmd.extend([
            "-f", "webm",
            "-y",               # Overwrite output without asking
            "pipe:1"
        ])
        
        if DEBUG_MODE:
            print(f"{bcolors.WARNING}[DEBUG] Executing command: {' '.join(ffmpeg_cmd)}{bcolors.ENDC}")

        # Start FFmpeg process
        print(f"{bcolors.OKGREEN}Dispatching ffmpeg Task with: {' '.join(ffmpeg_cmd)}{bcolors.ENDC}")
        process = subprocess.Popen(
            ffmpeg_cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            bufsize=1024*1024    # 1MB buffer
        )
        
        # Create a flag to track if we're actually getting video data
        received_data = False
        error_output = []
        
        # Create a thread to capture and log stderr
        def log_stderr():
            for line in iter(process.stderr.readline, b''):
                line_text = line.decode('utf-8', errors='replace').strip()
                error_output.append(line_text)
                if DEBUG_MODE:
                    print(f"{bcolors.OKCYAN}[FFMPEG] {line_text}{bcolors.ENDC}")
                    
        stderr_thread = threading.Thread(target=log_stderr)
        stderr_thread.daemon = True
        stderr_thread.start()
        
        try:
            # Stream output to client
            start_time = time.time()
            byte_count = 0
            
            while True:
                chunk = process.stdout.read(65536)  # 64KB chunk size
                
                # Check if we got any data
                if chunk:
                    byte_count += len(chunk)
                    received_data = True
                    if byte_count == len(chunk):  # First chunk
                        if DEBUG_MODE:
                            print(f"{bcolors.WARNING}[DEBUG] First chunk received after {time.time() - start_time:.2f} seconds, size: {len(chunk)} bytes{bcolors.ENDC}")
                    yield chunk
                else:
                    # Check if process has ended
                    if process.poll() is not None:
                        if DEBUG_MODE:
                            print(f"{bcolors.WARNING}[DEBUG] FFmpeg process completed with code: {process.returncode}{bcolors.ENDC}")
                            print(f"{bcolors.WARNING}[DEBUG] Total bytes streamed: {byte_count}{bcolors.ENDC}")
                        
                        # If we never got any data and process exited with error
                        if not received_data and process.returncode != 0:
                            if DEBUG_MODE:
                                print("{bcolors.FAIL}[ERROR] FFmpeg failed to produce any output{bcolors.ENDC}")
                                if error_output:
                                    print(f"{bcolors.FAIL}[ERROR] Last errors: {error_output[-10:]}{bcolors.ENDC}")
                        break
                    
                    # Short wait before trying again
                    time.sleep(0.1)
                
                # Safety timeout - if no data after 30 seconds, abort
                if not received_data and (time.time() - start_time) > 30:
                    if DEBUG_MODE:
                        print("{bcolors.FAIL}[ERROR] Timed out waiting for FFmpeg output{bcolors.ENDC}")
                    break
        
        except Exception as e:
            if DEBUG_MODE:
                print(f"{bcolors.FAIL}[ERROR] Streaming exception: {str(e)}{bcolors.ENDC}")
                
        finally:
            try:
                if process.poll() is None:
                    process.terminate()
                    time.sleep(0.5)
                    if process.poll() is None:
                        process.kill()
                
                # Wait for stderr thread to finish
                stderr_thread.join(timeout=1.0)
                
            except Exception as e:
                if DEBUG_MODE:
                    print(f"{bcolors.FAIL}[ERROR] Process cleanup error: {str(e)}{bcolors.ENDC}")
                    
            if DEBUG_MODE:
                print(f"{bcolors.WARNING}[DEBUG] Streaming ended for: {mkv_file}{bcolors.ENDC}")

    # Set proper headers for streaming
    headers = {
        'Content-Type': 'video/webm',
        'Cache-Control': 'no-cache, no-store, must-revalidate',
        'Pragma': 'no-cache',
        'X-Content-Type-Options': 'nosniff'
    }
    
    return Response(
        generate(),
        headers=headers,
        direct_passthrough=True
    )


@app.route('/test')
def test():
    """Simple test endpoint to verify the server is running"""
    return "Server is running"


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run Flask video streaming server")
    parser.add_argument("--debug", action="store_true", help="Enable debug mode")
    args = parser.parse_args()

    if args.debug:
        DEBUG_MODE = True
        print(f"{bcolors.WARNING}[DEBUG] Debug mode enabled{bcolors.ENDC}")

    app.run(host="0.0.0.0", port=8080, debug=DEBUG_MODE, threaded=True)