app.py

import os
import json
import time
import sqlite3
import asyncio
import threading
import subprocess
from datetime import datetime
from flask import Flask, request, jsonify, send_from_directory
from flask_socketio import SocketIO, emit
from flask_cors import CORS
from werkzeug.security import generate_password_hash
import yt_dlp
from providers import MegaProvider, FilenProvider, DrimeProvider
from task_manager import TaskManager

app = Flask(__name__, static_folder='frontend/build', static_url_path='')
CORS(app)
socketio = SocketIO(app, cors_allowed_origins="*", async_mode='eventlet')

# Initialize task manager
task_manager = TaskManager(socketio)

# Database setup
def init_db():
    conn = sqlite3.connect('/tmp/vod-archiver/tasks.db')
    c = conn.cursor()
    c.execute('''
        CREATE TABLE IF NOT EXISTS tasks (
            id TEXT PRIMARY KEY,
            user_id TEXT,
            vod_url TEXT,
            provider TEXT,
            format_id TEXT,
            status TEXT,
            progress_data TEXT,
            created_at TIMESTAMP,
            updated_at TIMESTAMP,
            error TEXT,
            file_info TEXT
        )
    ''')
    conn.commit()
    conn.close()

init_db()

@app.route('/')
def index():
    return send_from_directory(app.static_folder, 'index.html')

@app.route('/api/providers')
def get_providers():
    """Get list of available providers"""
    return jsonify({
        'providers': [
            {'id': 'mega', 'name': 'Mega.nz', 'maxSize': '20GB', 'features': ['streaming', 'encryption']},
            {'id': 'filen', 'name': 'Filen.io', 'maxSize': '20GB', 'features': ['encryption', 'privacy']},
            {'id': 'drime', 'name': 'Drime.cloud', 'maxSize': '30GB', 'features': ['streaming', 'sharing']}
        ]
    })

@app.route('/api/formats/<path:vod_url>')
def get_formats(vod_url):
    """Get available formats for VOD"""
    try:
        ydl_opts = {
            'quiet': True,
            'no_warnings': True,
            'extract_flat': False
        }
        
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            info = ydl.extract_info(vod_url, download=False)
            
            formats = []
            if 'formats' in info:
                # Group formats by quality
                quality_map = {}
                for f in info['formats']:
                    if f.get('height'):
                        quality = f"{f['height']}p"
                        fps = f.get('fps', 30)
                        quality_key = f"{quality}{fps}"
                        
                        if quality_key not in quality_map or (f.get('filesize', 0) > quality_map[quality_key].get('filesize', 0)):
                            quality_map[quality_key] = f
                
                # Convert to list
                for key, f in quality_map.items():
                    formats.append({
                        'format_id': f['format_id'],
                        'quality': f"{f.get('height', '?')}p",
                        'fps': f.get('fps', 30),
                        'ext': f.get('ext', 'mp4'),
                        'filesize': f.get('filesize', 0),
                        'label': f"{f.get('height', '?')}p{f.get('fps', '')}fps ({f.get('filesize', 0) / 1024 / 1024:.1f}MB)" if f.get('filesize') else f"{f.get('height', '?')}p{f.get('fps', '')}fps"
                    })
            
            # Always add best/worst
            formats.insert(0, {'format_id': 'best', 'label': 'Best Quality (Source)'})
            formats.append({'format_id': 'worst', 'label': 'Lowest Quality (Smallest)'})
            
            return jsonify({
                'formats': formats,
                'title': info.get('title', 'Unknown'),
                'duration': info.get('duration', 0),
                'thumbnail': info.get('thumbnail', '')
            })
            
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/api/validate-credentials', methods=['POST'])
def validate_credentials():
    """Validate provider credentials"""
    data = request.json
    provider_id = data.get('provider')
    credentials = data.get('credentials', {})
    
    try:
        if provider_id == 'mega':
            provider = MegaProvider()
            valid = provider.validate_credentials(
                credentials.get('email'),
                credentials.get('password')
            )
        elif provider_id == 'filen':
            provider = FilenProvider()
            valid = provider.validate_credentials(
                credentials.get('email'),
                credentials.get('password')
            )
        elif provider_id == 'drime':
            provider = DrimeProvider()
            valid = provider.validate_credentials(
                credentials.get('email'),
                credentials.get('password')
            )
        else:
            return jsonify({'valid': False, 'message': 'Invalid provider'}), 400
        
        return jsonify({'valid': valid, 'message': 'Valid' if valid else 'Invalid credentials'})
        
    except Exception as e:
        return jsonify({'valid': False, 'message': str(e)}), 500

@app.route('/api/tasks', methods=['GET'])
def get_tasks():
    """Get all tasks for current session"""
    user_id = request.headers.get('X-User-ID', 'default')
    
    conn = sqlite3.connect('/tmp/vod-archiver/tasks.db')
    conn.row_factory = sqlite3.Row
    c = conn.cursor()
    
    tasks = c.execute(
        'SELECT * FROM tasks WHERE user_id = ? ORDER BY created_at DESC',
        (user_id,)
    ).fetchall()
    
    conn.close()
    
    return jsonify({
        'tasks': [dict(task) for task in tasks]
    })

@app.route('/api/tasks', methods=['POST'])
def create_task():
    """Create a new download task"""
    data = request.json
    user_id = request.headers.get('X-User-ID', 'default')
    
    task_id = generate_password_hash(f"{user_id}{time.time()}")[:16]
    
    # Store task in database
    conn = sqlite3.connect('/tmp/vod-archiver/tasks.db')
    c = conn.cursor()
    c.execute('''
        INSERT INTO tasks (id, user_id, vod_url, provider, format_id, status, progress_data, created_at, updated_at)
        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
    ''', (
        task_id,
        user_id,
        data['vod_url'],
        data['provider'],
        data.get('format_id', 'best'),
        'queued',
        json.dumps({'phase': 'queued', 'percent': 0}),
        datetime.now(),
        datetime.now()
    ))
    conn.commit()
    conn.close()
    
    # Start task in background
    task_manager.start_task(task_id, data)
    
    return jsonify({
        'task_id': task_id,
        'status': 'queued'
    })

@app.route('/api/tasks/<task_id>', methods=['DELETE'])
def cancel_task(task_id):
    """Cancel a running task"""
    task_manager.cancel_task(task_id)
    
    conn = sqlite3.connect('/tmp/vod-archiver/tasks.db')
    c = conn.cursor()
    c.execute('UPDATE tasks SET status = ? WHERE id = ?', ('cancelled', task_id))
    conn.commit()
    conn.close()
    
    return jsonify({'status': 'cancelled'})

@socketio.on('connect')
def handle_connect():
    """Handle client connection"""
    emit('connected', {'data': 'Connected to server'})

@socketio.on('subscribe_task')
def handle_subscribe(data):
    """Subscribe to task updates"""
    task_id = data.get('task_id')
    if task_id:
        # Join room for this task
        from flask_socketio import join_room
        join_room(task_id)
        emit('subscribed', {'task_id': task_id})

@socketio.on('unsubscribe_task')
def handle_unsubscribe(data):
    """Unsubscribe from task updates"""
    task_id = data.get('task_id')
    if task_id:
        from flask_socketio import leave_room
        leave_room(task_id)

if __name__ == '__main__':
    socketio.run(app, host='0.0.0.0', port=7860, debug=False)