import os
import json
import time
import logging
import tempfile
import yt_dlp
from flask import Flask, Response, render_template, request
from werkzeug.utils import secure_filename
from threading import Thread, Lock
import uuid

app = Flask(__name__)

# --- Configuration ---

def get_temp_download_path():
    render_disk_path = os.environ.get('RENDER_DISK_PATH')
    if render_disk_path:
        base_path = render_disk_path
    else:
        base_path = tempfile.gettempdir()
    
    temp_folder = os.path.join(base_path, 'yt_temp_downloads')
    return temp_folder

DOWNLOAD_FOLDER = get_temp_download_path()
COOKIE_FILE = "cookies.txt"

logging.basicConfig(level=logging.INFO)
os.makedirs(DOWNLOAD_FOLDER, exist_ok=True)

# Progress storage with thread safety
progress_data = {}
progress_lock = Lock()

def update_progress(task_id, status, progress=0, **kwargs):
    """Update progress data thread-safely"""
    with progress_lock:
        progress_data[task_id] = {
            'status': status,
            'progress': float(progress),
            'timestamp': time.time(),
            **kwargs
        }
        app.logger.info(f"PROGRESS: {task_id} -> {status} ({progress}%)")

# Global variable to hold current task_id for progress hook
current_download_task = None

def create_progress_hook(task_id):
    """Create a progress hook that captures the task_id in closure"""
    def progress_hook(d):
        try:
            if d['status'] == 'downloading':
                total_bytes = d.get('total_bytes') or d.get('total_bytes_est', 0)
                downloaded_bytes = d.get('downloaded_bytes', 0)
                speed = d.get('speed', 0)
                eta = d.get('eta', 0)
                
                if total_bytes > 0:
                    percent = (downloaded_bytes / total_bytes) * 100
                    # Don't let it reach 100% until actually finished
                    percent = min(99.9, percent)
                else:
                    # If we don't have total size, estimate based on downloaded amount
                    # This is a fallback for cases where total size is unknown
                    percent = min(50, downloaded_bytes / (1024 * 1024))  # Rough estimate
                
                speed_mbps = (speed / (1024 * 1024)) if speed else 0
                
                update_progress(
                    task_id,
                    'downloading',
                    percent,
                    eta=int(eta) if eta else 0,
                    speed=f"{speed_mbps:.2f} MB/s",
                    downloaded_mb=downloaded_bytes / (1024 * 1024),
                    total_mb=total_bytes / (1024 * 1024) if total_bytes else 0
                )
                
            elif d['status'] == 'finished':
                update_progress(task_id, 'processing', 100)
                
        except Exception as e:
            app.logger.error(f"Progress hook error for {task_id}: {e}")
    
    return progress_hook

def download_worker(url, format_choice, task_id):
    """Download worker function"""
    global current_download_task
    current_download_task = task_id
    
    try:
        # Step 1: Initialize
        update_progress(task_id, 'initializing', 5)
        time.sleep(0.5)
        
        # Step 2: Get video info
        update_progress(task_id, 'fetching_info', 10)
        info_opts = {
            'quiet': True,
            'no_warnings': True,
            'cookiefile': COOKIE_FILE if os.path.exists(COOKIE_FILE) else None
        }
        
        with yt_dlp.YoutubeDL(info_opts) as ydl:
            info = ydl.extract_info(url, download=False)
            clean_title = secure_filename(info.get('title', 'video'))
            app.logger.info(f"Video title: {clean_title}")
        
        # Step 3: Prepare download
        update_progress(task_id, 'preparing', 15)
        time.sleep(0.5)
        
        # Configure download options
        timestamp = int(time.time())
        audio_formats = {"mp3", "m4a", "webm", "aac", "flac", "opus", "ogg", "wav"}
        
        if format_choice in audio_formats:
            unique_name = f"{clean_title}_{timestamp}"
            final_filename = f"{unique_name}.{format_choice}"
            ydl_opts = {
                'format': 'bestaudio/best',
                'postprocessors': [{
                    'key': 'FFmpegExtractAudio',
                    'preferredcodec': format_choice,
                    'preferredquality': '192',
                }]
            }
        elif format_choice in ["1080", "720", "480", "360", "1440"]:
            res = int(format_choice)
            unique_name = f"{clean_title}_{res}p_{timestamp}"
            final_filename = f"{unique_name}.mp4"
            ydl_opts = {
                'format': f'bestvideo[height<={res}]+bestaudio/best[height<={res}]',
                'merge_output_format': 'mp4'
            }
        else:
            raise ValueError(f"Invalid format: {format_choice}")
        
        # Create progress hook with task_id captured in closure
        progress_hook = create_progress_hook(task_id)
        
        # Set common options
        ydl_opts.update({
            'outtmpl': os.path.join(DOWNLOAD_FOLDER, f"{unique_name}.%(ext)s"),
            'progress_hooks': [progress_hook],
            'quiet': True,
            'no_warnings': True,
            'cookiefile': COOKIE_FILE if os.path.exists(COOKIE_FILE) else None
        })
        
        expected_path = os.path.join(DOWNLOAD_FOLDER, final_filename)
        
        # Step 4: Start download
        update_progress(task_id, 'starting_download', 20)
        time.sleep(0.5)
        
        app.logger.info(f"Starting yt-dlp download for {task_id}")
        
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            ydl.download([url])
        
        app.logger.info(f"yt-dlp download completed for {task_id}")
        
        # Step 5: Find the downloaded file
        actual_filename = final_filename
        if not os.path.exists(expected_path):
            # Look for any file with the base name
            base_name = unique_name
            found_files = []
            
            for filename in os.listdir(DOWNLOAD_FOLDER):
                if filename.startswith(base_name) and not filename.endswith('.part'):
                    found_files.append(filename)
            
            if found_files:
                # Use the most recent file
                found_files.sort(key=lambda x: os.path.getctime(os.path.join(DOWNLOAD_FOLDER, x)), reverse=True)
                actual_filename = found_files[0]
                expected_path = os.path.join(DOWNLOAD_FOLDER, actual_filename)
                app.logger.info(f"Found downloaded file: {actual_filename}")
        
        # Verify file exists and isn't empty
        if not os.path.exists(expected_path):
            raise FileNotFoundError(f"Download completed but file not found: {actual_filename}")
        
        file_size = os.path.getsize(expected_path)
        if file_size == 0:
            raise ValueError("Downloaded file is empty")
            
        app.logger.info(f"File verified: {actual_filename} ({file_size} bytes)")
        
        # Step 6: Complete
        update_progress(task_id, 'complete', 100, filename=actual_filename)
        
    except Exception as e:
        error_msg = str(e)
        app.logger.error(f"Download error for {task_id}: {error_msg}")
        update_progress(task_id, 'error', 0, message=error_msg)
        
        # Cleanup on error
        try:
            if 'expected_path' in locals() and os.path.exists(expected_path):
                os.remove(expected_path)
                app.logger.info(f"Cleaned up failed download: {expected_path}")
        except Exception as cleanup_error:
            app.logger.error(f"Cleanup error: {cleanup_error}")
    
    finally:
        # Clear global task reference
        if current_download_task == task_id:
            current_download_task = None

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/stream-download', methods=['GET'])
def stream_download():
    url = request.args.get('url')
    format_choice = request.args.get('format')
    
    if not url or not format_choice:
        return Response(
            json.dumps({"error": "Missing parameters"}),
            status=400,
            mimetype='application/json'
        )
    
    task_id = str(uuid.uuid4())
    app.logger.info(f"New download request: {task_id} - {url} - {format_choice}")
    
    # Initialize progress
    update_progress(task_id, 'waiting', 0)
    
    # Start download thread
    thread = Thread(target=download_worker, args=(url, format_choice, task_id))
    thread.daemon = True
    thread.start()
    
    def generate():
        try:
            start_time = time.time()
            timeout = 1800  # 30 minutes
            
            while True:
                # Check timeout
                if time.time() - start_time > timeout:
                    update_progress(task_id, 'error', 0, message='Download timeout')
                    break
                
                # Get current progress
                with progress_lock:
                    data = progress_data.get(task_id, {
                        'status': 'waiting',
                        'progress': 0,
                        'timestamp': time.time()
                    })
                
                # Send data
                json_data = json.dumps(data)
                yield f"data: {json_data}\n\n"
                
                # Check if finished
                if data.get('status') in ['complete', 'error']:
                    break
                
                time.sleep(0.5)
                
        except GeneratorExit:
            app.logger.info(f"Client disconnected: {task_id}")
        except Exception as e:
            app.logger.error(f"Stream error for {task_id}: {e}")
        finally:
            # Cleanup after delay
            def cleanup():
                time.sleep(30)
                with progress_lock:
                    if task_id in progress_data:
                        del progress_data[task_id]
                        app.logger.info(f"Cleaned up progress data for {task_id}")
            
            Thread(target=cleanup, daemon=True).start()
    
    response = Response(generate(), mimetype='text/event-stream')
    response.headers.update({
        'Cache-Control': 'no-cache',
        'Connection': 'keep-alive',
        'Access-Control-Allow-Origin': '*',
        'X-Accel-Buffering': 'no'
    })
    
    return response

@app.route('/download-file/<filename>')
def download_file(filename):
    safe_folder = os.path.abspath(DOWNLOAD_FOLDER)
    filepath = os.path.join(safe_folder, filename)
    
    # Security check
    if not os.path.abspath(filepath).startswith(safe_folder):
        return "Forbidden", 403
    
    if not os.path.exists(filepath):
        return "File not found", 404
    
    def generate_and_cleanup():
        try:
            app.logger.info(f"Serving file: {filename}")
            with open(filepath, 'rb') as f:
                while True:
                    chunk = f.read(8192)
                    if not chunk:
                        break
                    yield chunk
        finally:
            # Remove file after download
            def remove_file():
                time.sleep(2)
                try:
                    if os.path.exists(filepath):
                        os.remove(filepath)
                        app.logger.info(f"Cleaned up file: {filename}")
                except Exception as e:
                    app.logger.error(f"Error removing file {filename}: {e}")
            
            Thread(target=remove_file, daemon=True).start()
    
    return Response(
        generate_and_cleanup(),
        headers={
            'Content-Disposition': f'attachment; filename="{filename}"',
            'Content-Type': 'application/octet-stream'
        }
    )

@app.route('/health')
def health_check():
    return {
        "status": "healthy", 
        "active_downloads": len(progress_data),
        "current_task": current_download_task
    }

if __name__ == "__main__":
    app.run(debug=True, threaded=True, host='0.0.0.0', port=5000)