app.py

import os
import uuid
import shlex
import time
import threading
import subprocess
from pathlib import Path
from datetime import datetime, timedelta
from concurrent.futures import ThreadPoolExecutor
from typing import Any, Dict, List, Optional

from flask import Flask, render_template, request, redirect, url_for, send_file, jsonify
from werkzeug.utils import secure_filename


BASE_DIR = Path(__file__).resolve().parent
UPLOAD_DIR = BASE_DIR / 'tmp'
UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
ALLOWED_EXTENSIONS = {'.mp4', '.mov', '.mkv', '.avi', '.webm', '.mpg', '.mpeg'}
ENCODER_KEYWORDS = ('nvenc', 'qsv', 'amf', 'vaapi', 'cuvid')
LOG_CHAR_LIMIT = 20000
STRAY_FILE_RETENTION = timedelta(hours=48)

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = str(UPLOAD_DIR)
app.config['MAX_CONTENT_LENGTH'] = 1024 * 1024 * 1024  # 1GB limit

# Jobs store metadata about running/finished jobs
jobs = {}
executor = ThreadPoolExecutor(max_workers=2)


def trim_log(text: str) -> str:
    if len(text) <= LOG_CHAR_LIMIT:
        return text
    return text[-LOG_CHAR_LIMIT:]


def select_encoder(requested: Optional[str]) -> str:
    if requested and requested.lower() != 'auto':
        return requested
    for keyword in ('h264', 'hevc', '265'):
        for encoder in available_encoders:
            if keyword in encoder.lower():
                return encoder
    return 'libx264'


def normalize_resolution(resolution: Optional[str]) -> Optional[str]:
    if not resolution:
        return None
    value = resolution.strip().lower()
    if not value:
        return None
    if 'x' in value and ':' not in value:
        parts = value.split('x', 1)
        value = f"{parts[0]}:{parts[1]}"
    return value


def build_ffmpeg_command(job: Dict[str, Any]) -> List[str]:
    input_path = job['input']
    progress_path = job['progress_file']
    output_path = job['output']

    cmd: List[str] = ['ffmpeg', '-y', '-hide_banner', '-loglevel', 'info', '-nostdin', '-i', input_path]
    raw_extra = job.get('extra_args') or []
    if isinstance(raw_extra, (list, tuple)):
        extra = list(raw_extra)
    else:
        extra = shlex.split(str(raw_extra)) if raw_extra else []
    codec = (job.get('codec') or 'auto').lower()

    if codec == 'copy':
        cmd.extend(['-c:v', 'copy', '-c:a', 'copy'])
        cmd.extend(extra)
    else:
        encoder = select_encoder(job.get('codec'))
        cmd.extend(['-c:v', encoder])
        crf = job.get('crf')
        if crf and encoder != 'copy':
            cmd.extend(['-crf', str(crf)])
        bitrate = job.get('bitrate')
        if bitrate:
            cmd.extend(['-b:v', str(bitrate)])
        resolution = normalize_resolution(job.get('resolution'))
        if resolution:
            cmd.extend(['-vf', f'scale={resolution}'])
        framerate = job.get('framerate')
        if framerate:
            cmd.extend(['-r', str(framerate)])
        cmd.extend(extra)

    cmd.extend(['-progress', progress_path, '-nostats', output_path])
    return cmd


def safe_remove(path: Optional[str]) -> None:
    if not path:
        return
    try:
        p = Path(path)
        if p.exists():
            p.unlink()
    except Exception:
        pass


def cleanup_job_files(job: Dict[str, Any]) -> None:
    safe_remove(job.get('input'))
    safe_remove(job.get('output'))
    safe_remove(job.get('progress_file'))


def run_cmd_check(cmd):
    try:
        out = subprocess.check_output(cmd, stderr=subprocess.STDOUT, shell=False, universal_newlines=True)
        return out
    except Exception:
        return ''


def detect_encoders():
    out = run_cmd_check(['ffmpeg', '-hide_banner', '-encoders'])
    enc: List[str] = []
    for line in out.splitlines():
        parts = line.strip().split()
        if not parts:
            continue
        # basic heuristic: encoder names are last token(s)
        # lines look like " V..... h264_nvenc             NVIDIA NVENC H.264 encoder"
        lower_line = line.lower()
        if any(keyword in lower_line for keyword in ENCODER_KEYWORDS) or 'h264_cuvid' in lower_line or 'hevc_nvenc' in lower_line:
            tokens = line.split()
            for token in tokens:
                token_lower = token.lower()
                if any(keyword in token_lower for keyword in ENCODER_KEYWORDS):
                    enc.append(token.strip())
                    break
    # dedupe while preserving order
    seen = set()
    deduped = []
    for name in enc:
        if name not in seen:
            seen.add(name)
            deduped.append(name)
    return deduped


available_encoders = detect_encoders()


def human_readable_size(path: Optional[str]) -> str:
    if not path:
        return ''
    try:
        size = Path(path).stat().st_size
        for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
            if size < 1024:
                return f"{size:.2f}{unit}"
            size /= 1024.0
        return f"{size:.2f}PB"
    except Exception:
        return ''

def get_duration_seconds(path):
    try:
        out = run_cmd_check(['ffprobe', '-v', 'error', '-show_entries', 'format=duration', '-of', 'default=noprint_wrappers=1:nokey=1', str(path)])
        return float(out.strip())
    except Exception:
        return None


def sanitize_extra_args(extra_text: str) -> List[str]:
    # naive sanitization: disallow characters that may be used to chain commands in shell
    if not extra_text:
        return []
    if any(c in extra_text for c in (';', '|', '&', '>', '<')):
        raise ValueError('Invalid characters in extra args')
    # split safely
    try:
        tokens = shlex.split(extra_text)
    except Exception:
        tokens = extra_text.split()
    # disallow tokens that look like paths or contain '/'
    safe_tokens = []
    for t in tokens:
        if t.startswith('/') or t.startswith('\\') or ':' in t and len(t) > 2 and (t[1] == ':' or t[2] == ':'):
            # likely paths or windows drive
            raise ValueError('File paths not allowed in extra args')
        safe_tokens.append(t)
    return safe_tokens


def encode_worker(jobid):
    job = jobs[jobid]
    job['status'] = 'starting'
    job['duration'] = get_duration_seconds(job['input'])
    job['status'] = 'running'

    cmd = build_ffmpeg_command(job)

    job['cmd'] = ' '.join(shlex.quote(p) for p in cmd)
    job['status'] = 'encoding'
    job['start_time'] = datetime.utcnow().isoformat()
    try:
        Path(job['progress_file']).write_text('', encoding='utf-8')
    except Exception:
        pass

    try:
        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)
        job['proc'] = proc
        job['pid'] = proc.pid
        job['log'] = (job.get('cmd', '') or '') + '\n'
        while True:
            line = proc.stderr.readline() if proc.stderr else ''
            if not line:
                if proc.poll() is not None:
                    break
                time.sleep(0.1)
                continue
            if line.strip():
                job['log'] = trim_log((job.get('log') or '') + line)
        returncode = proc.wait()
        if returncode == 0 and os.path.exists(job['output']):
            job['status'] = 'done'
            job['done'] = True
            job['finished_time'] = datetime.utcnow().isoformat()
        else:
            job['status'] = 'error'
            job['done'] = True
            job['err'] = f'Non-zero exit code: {returncode}'
    except Exception as exc:
        job['status'] = 'error'
        job['done'] = True
        job['err'] = str(exc)
    finally:
        try:
            job['proc'] = None
        except Exception:
            pass


def parse_progress(progress_path):
    info = {}
    try:
        with open(progress_path, 'r', encoding='utf-8') as f:
            for line in f.read().splitlines():
                if '=' in line:
                    k, v = line.split('=', 1)
                    info[k.strip()] = v.strip()
    except Exception:
        pass
    return info


def parse_time_to_ms(raw: str) -> Optional[int]:
    try:
        parts = raw.split(':')
        if len(parts) == 3:
            hours = float(parts[0])
            minutes = float(parts[1])
            seconds = float(parts[2])
            total_seconds = hours * 3600 + minutes * 60 + seconds
            return int(total_seconds * 1000)
    except Exception:
        return None
    return None


def safe_int(value: Optional[str]) -> Optional[int]:
    try:
        if value is None or value == '' or value == 'N/A':
            return None
        return int(float(value))
    except Exception:
        return None


def safe_float(value: Optional[str]) -> Optional[float]:
    try:
        if value is None or value == '' or value == 'N/A':
            return None
        return float(value)
    except Exception:
        return None


def parse_speed(value: Optional[str]) -> Optional[float]:
    if not value:
        return None
    speed_str = value.strip()
    if speed_str.endswith('x'):
        speed_str = speed_str[:-1]
    return safe_float(speed_str)


def format_hms(seconds: Optional[float]) -> Optional[str]:
    if seconds is None:
        return None
    try:
        total = max(0, int(seconds))
        h = total // 3600
        m = (total % 3600) // 60
        s = total % 60
        return f"{h:02d}:{m:02d}:{s:02d}"
    except Exception:
        return None


def build_progress_payload(job: Dict[str, Any], progress_fields: Dict[str, str]) -> Dict[str, Any]:
    out_ms = None
    if progress_fields.get('out_time_ms'):
        out_ms = safe_int(progress_fields.get('out_time_ms'))
    elif progress_fields.get('out_time'):
        out_ms = parse_time_to_ms(progress_fields.get('out_time', ''))

    frame = safe_int(progress_fields.get('frame'))
    fps = safe_float(progress_fields.get('fps'))
    speed = parse_speed(progress_fields.get('speed'))
    bitrate = progress_fields.get('bitrate') if progress_fields.get('bitrate') else None
    total_size = safe_int(progress_fields.get('total_size'))

    current_time_seconds = out_ms / 1000.0 if out_ms is not None else None
    if current_time_seconds is None and frame is not None and fps and fps > 0:
        current_time_seconds = frame / fps
        out_ms = int(current_time_seconds * 1000)

    percentage = float(job.get('last_progress', 0.0))
    eta_seconds = None
    duration = job.get('duration')
    try:
        if out_ms is not None and duration:
            duration_sec = float(duration)
            if duration_sec > 0:
                percentage = max(0.0, min(1.0, out_ms / (duration_sec * 1000.0)))
                if speed and speed > 0:
                    remaining = max(0.0, duration_sec - (out_ms / 1000.0))
                    eta_seconds = remaining / speed if speed else None
    except Exception:
        percentage = float(job.get('last_progress', 0.0))

    if progress_fields.get('progress') == 'end' or job.get('done'):
        percentage = 1.0
        eta_seconds = 0

    job['last_progress'] = percentage

    return {
        'percentage': percentage,
        'out_ms': out_ms,
        'current_time_seconds': current_time_seconds,
        'eta_seconds': eta_seconds,
        'frame': frame,
        'fps': fps,
        'speed': speed,
        'bitrate': bitrate,
        'total_size': total_size,
    }


@app.route('/')
def main():
    return render_template('main.html', encoders=available_encoders)


@app.route('/encode', methods=['POST'])
def encode():
    file = request.files.get('file')
    if not file:
        return 'Missing file', 400
    if not file.filename:
        return 'Invalid filename', 400
    filename = secure_filename(file.filename)
    _, ext = os.path.splitext(filename)
    if ext.lower() not in ALLOWED_EXTENSIONS:
        return 'Unsupported file type', 400

    jobid = uuid.uuid4().hex
    upload_dir = Path(app.config['UPLOAD_FOLDER'])
    in_path = upload_dir / f'{jobid}_in{ext}'
    # sanitize container
    container = request.form.get('container') or 'mp4'
    if not container.isalnum():
        container = 'mp4'
    out_filename = f'{jobid}_out.{container}'
    out_path = upload_dir / out_filename
    progress_file = upload_dir / f'{jobid}.progress'

    # ensure the parent directory exists (race conditions may remove it)
    try:
        in_path.parent.mkdir(parents=True, exist_ok=True)
    except Exception:
        pass
    try:
        file.save(str(in_path))
    except Exception as e:
        # If saving fails, return a 400 with a helpful error to the user
        return f'Failed to save uploaded file: {e}', 400

    # parse form
    form_codec = request.form.get('codec')
    bitrate = request.form.get('bitrate')
    crf = request.form.get('crf')
    resolution = request.form.get('resolution')
    framerate = request.form.get('framerate')
    lifetime_min = int(request.form.get('lifetime') or 30)
    extra = request.form.get('extra') or ''
    try:
        extra_args = sanitize_extra_args(extra)
    except Exception as e:
        return f'Invalid extra args: {e}', 400

    # build job record
    job = {
        'id': jobid,
        'input': str(in_path),
        'output': str(out_path),
        'progress_file': str(progress_file),
        'status': 'queued',
        'done': False,
        'created_time': datetime.utcnow().isoformat(),
        'lifetime_min': lifetime_min,
        'cmd': None,
        'log': '',
        'last_progress': 0.0,
        'codec': form_codec,
        'container': container,
        'bitrate': bitrate,
        'crf': crf,
        'resolution': resolution,
        'framerate': framerate,
        'extra_args': extra_args,
    }
    jobs[jobid] = job

    # Submit encoding job
    jobs[jobid]['future'] = executor.submit(encode_worker, jobid)

    return redirect(url_for('job_status', jobid=jobid))


@app.route('/job/<jobid>')
def job_status(jobid):
    j = jobs.get(jobid)
    if not j:
        return 'Unknown job', 404
    return render_template('job.html', job=j)


@app.route('/progress/<jobid>')
def progress(jobid):
    job = jobs.get(jobid)
    if not job:
        return jsonify({'error': 'Unknown job'}), 404
    progress_fields = parse_progress(job['progress_file'])
    metrics = build_progress_payload(job, progress_fields)
    done = job.get('done', False)
    current_time_seconds = metrics.get('current_time_seconds')
    eta_seconds = metrics.get('eta_seconds') if (current_time_seconds is not None and job.get('duration')) else None

    return jsonify({
        'id': jobid,
        'status': job.get('status'),
        'progress': metrics.get('percentage'),
        'log': (job.get('log') or '')[-8000:],
        'frame': metrics.get('frame'),
        'fps': metrics.get('fps'),
        'speed': metrics.get('speed'),
        'bitrate': metrics.get('bitrate'),
        'total_size': human_readable_size(job.get('output')) if done else None,
        'out_time_ms': metrics.get('out_ms'),
        'current_time': format_hms(current_time_seconds),
        'eta_seconds': eta_seconds,
        'eta': format_hms(eta_seconds),
        'done': done,
        'size': human_readable_size(job.get('output')) if done else None,
    })


@app.route('/download/<jobid>')
def download(jobid):
    j = jobs.get(jobid)
    if not j:
        return 'Unknown job', 404
    output_path = Path(j.get('output')) if j.get('output') else None
    if not j.get('done') or not output_path or not output_path.exists():
        return 'Not ready', 404
    return send_file(str(output_path), as_attachment=True)


@app.route('/remove/<jobid>', methods=['POST'])
def remove_job(jobid):
    j = jobs.get(jobid)
    if not j:
        return jsonify({'error': 'Unknown job'}), 404
    # Try to terminate a running process if present
    try:
        proc = j.get('proc')
        if proc and isinstance(proc, subprocess.Popen):
            try:
                proc.terminate()
            except Exception:
                try:
                    proc.kill()
                except Exception:
                    pass
    except Exception:
        pass
    # Remove files
    try:
        cleanup_job_files(j)
        jobs.pop(jobid, None)
    except Exception as e:
        return jsonify({'error': f'Failed to remove job: {e}'}), 500
    return jsonify({'id': jobid, 'removed': True})


def cleanup_worker():
    while True:
        try:
            now = datetime.utcnow()
            for jid, j in list(jobs.items()):
                # if done and older than lifetime, remove files and job
                created_raw = j.get('created_time')
                if not created_raw:
                    continue
                try:
                    created = datetime.fromisoformat(created_raw)
                except Exception:
                    continue
                lifetime = timedelta(minutes=int(j.get('lifetime_min', 30)))
                if now - created > lifetime:
                    cleanup_job_files(j)
                    jobs.pop(jid, None)
            # also clean up any stray tmp files older than 48h
            for path in Path(app.config['UPLOAD_FOLDER']).iterdir():
                try:
                    mtime = datetime.utcfromtimestamp(path.stat().st_mtime)
                    if now - mtime > STRAY_FILE_RETENTION and path.is_file():
                        path.unlink()
                except Exception:
                    pass
        except Exception:
            pass
        time.sleep(60)


# Start cleanup thread
cleanup_thread = threading.Thread(target=cleanup_worker, daemon=True)
cleanup_thread.start()


if __name__ == '__main__':
    # quick check that ffmpeg/ffprobe are available
    if not run_cmd_check(['ffmpeg', '-version']):
        print('FFmpeg is not available in PATH. Please install ffmpeg and ensure it is in PATH.')
    app.run(host='0.0.0.0', port=5000, debug=True)