core.py

import os
import re
import json
import uuid
import shutil
import subprocess
import threading
from pathlib import Path
from typing import Optional, Dict, Any, List
from datetime import datetime

print("[YTDL] Starting...")
print("[YTDL] Gradio version:", __import__("gradio").__version__)

from gradio import Server
from gradio.data_classes import FileData
from fastapi.responses import HTMLResponse, FileResponse
import spaces

print("[YTDL] All imports OK, Server class available")

# ===== Config =====
DOWNLOAD_DIR = "./downloads"
PROGRESS_DIR = "./.progress"
os.makedirs(DOWNLOAD_DIR, exist_ok=True)
os.makedirs(PROGRESS_DIR, exist_ok=True)

QUALITY_ORDER = ['2160p', '1440p', '1080p', '720p', '480p', '360p', '240p', '144p']

def _yt_dlp_path():
    candidates = [
        shutil.which("yt-dlp"),
        os.path.expanduser("~/.local/bin/yt-dlp"),
        "/usr/local/bin/yt-dlp",
        "/usr/bin/yt-dlp",
    ]
    for c in candidates:
        if c and os.path.isfile(c):
            return c
    raise FileNotFoundError("yt-dlp not found")

YTDLP = _yt_dlp_path()

def _safe_filename(title: str, max_len: int = 80) -> str:
    name = re.sub(r'[^\w\s\-\(\)\[\].]', '', title).strip()
    name = re.sub(r'[-\s]+', '_', name)
    return name[:max_len] if name else "video"

def _run_ytdlp(args: List[str], timeout: int = 600) -> Dict[str, Any]:
    cmd = [YTDLP] + args
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
        if result.returncode != 0:
            stderr = result.stderr.strip()
            if "ERROR:" in stderr:
                error_msg = stderr.split("ERROR:")[-1].strip().split("\n")[0]
            else:
                error_msg = stderr or f"yt-dlp exited with code {result.returncode}"
            return {"success": False, "error": error_msg}
        return {"success": True, "data": result.stdout.strip()}
    except subprocess.TimeoutExpired:
        return {"success": False, "error": f"Download timed out after {timeout}s"}
    except FileNotFoundError:
        return {"success": False, "error": "yt-dlp binary not found"}
    except Exception as e:
        return {"success": False, "error": str(e)}

def _save_progress(task_id: str, data: Dict[str, Any]):
    filepath = os.path.join(PROGRESS_DIR, f"{task_id}.json")
    with open(filepath, "w") as f:
        json.dump(data, f)

def _get_progress_file(task_id: str) -> Optional[Dict[str, Any]]:
    filepath = os.path.join(PROGRESS_DIR, f"{task_id}.json")
    if os.path.exists(filepath):
        with open(filepath, "r") as f:
            return json.load(f)
    return None

def _cleanup_old_downloads(max_age_hours: int = 24):
    now = datetime.now().timestamp()
    for entry in os.scandir(DOWNLOAD_DIR):
        if entry.is_dir():
            try:
                if now - entry.stat().st_mtime > max_age_hours * 3600:
                    shutil.rmtree(entry.path)
            except OSError:
                pass


# ===== Core Functions =====

def get_video_info(url: str) -> Dict[str, Any]:
    result = _run_ytdlp([
        "--dump-json", "--no-playlist", "--no-warnings", "--no-download", url,
    ])
    if not result["success"]:
        return result
    try:
        data = json.loads(result["data"])
        thumbnail = ""
        if data.get("thumbnail"):
            thumbnail = data["thumbnail"]
        elif data.get("thumbnails"):
            thumbs = sorted(
                [t for t in data["thumbnails"] if t.get("url")],
                key=lambda t: t.get("width", 0) or 0, reverse=True,
            )
            thumbnail = thumbs[0]["url"] if thumbs else ""

        duration = data.get("duration", 0) or 0
        h, rem = divmod(duration, 3600)
        m, s = divmod(rem, 60)
        duration_str = f"{h}:{m:02d}:{s:02d}" if h else f"{m}:{s:02d}"

        views = data.get("view_count", 0) or 0
        upload_date = data.get("upload_date")
        if upload_date and len(upload_date) == 8:
            upload_date = f"{upload_date[:4]}-{upload_date[4:6]}-{upload_date[6:8]}"

        captions = list(data["subtitles"].keys()) if data.get("subtitles") else []

        return {
            "success": True,
            "info": {
                "id": data.get("id", ""),
                "title": data.get("title", "Unknown Title"),
                "author": data.get("uploader") or data.get("channel", "Unknown"),
                "duration": duration,
                "duration_string": duration_str,
                "views": views,
                "description": (data.get("description") or "")[:800],
                "upload_date": upload_date,
                "thumbnail": thumbnail,
                "keywords": data.get("tags", []) or [],
                "captions": captions,
                "is_live": data.get("is_live", False),
            },
        }
    except (json.JSONDecodeError, KeyError) as e:
        return {"success": False, "error": f"Failed to parse video info: {str(e)}"}


def get_available_streams(url: str) -> Dict[str, Any]:
    result = _run_ytdlp([
        "--dump-json", "--no-playlist", "--no-warnings", "--no-download", url,
    ])
    if not result["success"]:
        return result
    try:
        data = json.loads(result["data"])
        formats = data.get("formats", [])
        progressive, video_only, audio_only = [], [], []
        seen = set()

        for f in formats:
            fid = str(f.get("format_id", ""))
            if fid in seen:
                continue
            seen.add(fid)

            is_audio = f.get("vcodec") in ("none", None)
            is_video = f.get("acodec") in ("none", None)
            is_prog = not is_audio and not is_video

            height = f.get("height")
            # Normalize resolution to "720p" format
            raw_res = f.get("resolution") or ""
            if is_audio:
                resolution = "audio"
            elif height:
                resolution = f"{height}p"
            elif raw_res and "x" in raw_res:
                try:
                    resolution = f"{int(raw_res.split('x')[1])}p"
                except (ValueError, IndexError):
                    resolution = raw_res
            else:
                resolution = raw_res or None

            ext = f.get("ext", "mp4")
            filesize = f.get("filesize") or f.get("filesize_approx")
            filesize_mb = round(filesize / (1024 * 1024), 2) if filesize else None
            tbr = f.get("tbr")

            entry = {
                "format_id": fid,
                "resolution": resolution,
                "ext": ext,
                "filesize_mb": filesize_mb,
                "tbr": round(tbr) if tbr else None,
                "vcodec": f.get("vcodec"),
                "acodec": f.get("acodec"),
                "fps": f.get("fps"),
                "label": "",
            }

            if is_audio:
                abr = f.get("abr")
                bitrate = f"{round(abr)}kbps" if abr else "Audio"
                entry["abr"] = bitrate
                entry["label"] = f"{bitrate} {ext.upper()}"
                audio_only.append(entry)
            elif is_video:
                codec = (f.get("vcodec") or "").split(".")[0]
                entry["label"] = f"{resolution or 'Video'} {ext.upper()}" + (f" ({codec})" if codec else "")
                video_only.append(entry)
            elif is_prog:
                entry["label"] = f"{resolution or 'Video'} {ext.upper()}"
                progressive.append(entry)

        def res_sort_key(x):
            r = x.get("resolution", "")
            try:
                return QUALITY_ORDER.index(r) if r in QUALITY_ORDER else 999
            except ValueError:
                return 999

        progressive.sort(key=res_sort_key)
        video_only.sort(key=res_sort_key)
        audio_only.sort(key=lambda x: x.get("tbr") or 0, reverse=True)

        all_res = list(set(
            s["resolution"] for s in progressive + video_only
            if s.get("resolution") and s["resolution"] != "audio"
        ))

        def res_str_sort_key(r):
            try:
                return QUALITY_ORDER.index(r) if r in QUALITY_ORDER else 999
            except ValueError:
                return 999

        all_res_sorted = sorted(all_res, key=res_str_sort_key)

        return {
            "success": True,
            "progressive": progressive,
            "video_only": video_only,
            "audio_only": audio_only,
            "available_resolutions": all_res_sorted,
            "best_resolution": all_res_sorted[0] if all_res_sorted else None,
        }
    except (json.JSONDecodeError, KeyError) as e:
        return {"success": False, "error": f"Failed to parse streams: {str(e)}"}


def download_file(url: str, format_id: str, is_audio: bool = False,
                  audio_quality: str = "128") -> Dict[str, Any]:
    task_id = uuid.uuid4().hex[:12]
    info_result = get_video_info(url)
    if not info_result["success"]:
        return info_result

    video_info = info_result["info"]
    video_id = video_info["id"]
    safe_title = _safe_filename(video_info["title"])
    out_dir = Path(DOWNLOAD_DIR) / video_id
    out_dir.mkdir(parents=True, exist_ok=True)

    _save_progress(task_id, {
        "status": "starting", "percent": 0,
        "speed": None, "eta": None, "filename": None,
    })

    if is_audio:
        output_template = str(out_dir / f"{safe_title}_audio.%(ext)s")
        args = ["-x", "--audio-format", "mp3", "--audio-quality", audio_quality,
                "-o", output_template, "--no-playlist", "--newline", "--progress", url]
    else:
        output_template = str(out_dir / f"{safe_title}_{format_id}.%(ext)s")
        args = ["-f", format_id, "--merge-output-format", "mp4",
                "-o", output_template, "--no-playlist", "--newline", "--progress", url]

    def _download():
        try:
            _save_progress(task_id, {"status": "downloading", "percent": 0, "speed": None, "eta": None})
            process = subprocess.Popen([YTDLP] + args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
            last_percent = 0
            while True:
                line = process.stdout.readline()
                if not line and process.poll() is not None:
                    break
                line = line.strip()
                if not line:
                    continue
                match = re.search(r'(\d+(?:\.\d+)?)%', line)
                if match:
                    percent = float(match.group(1))
                    last_percent = percent
                    speed = None
                    eta = None
                    sm = re.search(r'at\s+([\d.]+\w+/s)', line)
                    em = re.search(r'ETA\s+([\d:]+)', line)
                    if sm:
                        speed = sm.group(1)
                    if em:
                        eta = em.group(1)
                    _save_progress(task_id, {"status": "downloading", "percent": percent, "speed": speed, "eta": eta})
            process.wait()
            if process.returncode != 0:
                stderr = process.stderr.read()
                error_msg = "Download failed"
                if "ERROR:" in stderr:
                    error_msg = stderr.split("ERROR:")[-1].strip().split("\n")[0]
                _save_progress(task_id, {"status": "error", "percent": last_percent, "error": error_msg})
                return
            files = list(out_dir.glob(f"{safe_title}*"))
            files = [f for f in files if f.suffix not in (".part", ".ytdl", ".temp")]
            if not files:
                files = [f for f in out_dir.iterdir() if f.is_file() and f.stat().st_size > 1024]
            if files:
                downloaded = max(files, key=lambda f: f.stat().st_mtime)
                _save_progress(task_id, {"status": "completed", "percent": 100, "filename": downloaded.name, "filepath": str(downloaded)})
            else:
                _save_progress(task_id, {"status": "error", "percent": last_percent, "error": "File not found after download"})
        except Exception as e:
            _save_progress(task_id, {"status": "error", "percent": 0, "error": str(e)})

    thread = threading.Thread(target=_download, daemon=True)
    thread.start()
    return {"success": True, "task_id": task_id, "message": "Download started"}


def get_progress(task_id: str) -> Dict[str, Any]:
    progress = _get_progress_file(task_id)
    if progress is None:
        return {"success": False, "error": "Task not found"}
    result = {"success": True}
    result.update(progress)
    return result


def download_subtitles(url: str, lang: str = "en") -> Dict[str, Any]:
    info_result = get_video_info(url)
    if not info_result["success"]:
        return info_result
    video_id = info_result["info"]["id"]
    safe_title = _safe_filename(info_result["info"]["title"])
    out_dir = Path(DOWNLOAD_DIR) / video_id
    out_dir.mkdir(parents=True, exist_ok=True)
    output_template = str(out_dir / f"{safe_title}_subtitles")
    result = _run_ytdlp([
        "--write-sub", "--write-auto-sub", "--sub-langs", lang,
        "--sub-format", "srt", "--skip-download", "-o", output_template, "--no-playlist", url,
    ], timeout=60)
    if not result["success"]:
        return result
    sub_files = list(out_dir.glob("*.srt"))
    if sub_files:
        return {"success": True, "message": f"Subtitles downloaded in {len(sub_files)} file(s)", "files": [f.name for f in sub_files]}
    return {"success": False, "error": "No subtitles found for this video"}


# ===== Gradio Server =====
print("[YTDL] Creating Server instance...")
demo = Server()
print("[YTDL] Server instance created")

@demo.api(name="get_video_info")
def video_info_api(url: str) -> Dict[str, Any]:
    return get_video_info(url)

@demo.api(name="get_streams")
def streams_info_api(url: str) -> Dict[str, Any]:
    return get_available_streams(url)

@demo.api(name="download")
def download_api(url: str, format_id: str, is_audio: bool = False,
                 audio_quality: str = "128") -> Dict[str, Any]:
    return download_file(url, format_id, is_audio, audio_quality)

@demo.api(name="get_progress")
def progress_api(task_id: str) -> Dict[str, Any]:
    return get_progress(task_id)

@demo.api(name="download_subtitles")
def subtitles_api(url: str, lang: str = "en") -> Dict[str, Any]:
    return download_subtitles(url, lang)

@demo.get("/")
async def homepage():
    html_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html")
    if os.path.exists(html_path):
        with open(html_path, "r", encoding="utf-8") as f:
            return HTMLResponse(f.read())
    return HTMLResponse("<h1>YTDL</h1><p>index.html not found.</p>")

@demo.get("/api-guide")
async def api_guide_page():
    guide_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "api_guide.html")
    if os.path.exists(guide_path):
        with open(guide_path, "r", encoding="utf-8") as f:
            return HTMLResponse(f.read())
    return HTMLResponse("<h1>API Guide</h1>")

@demo.get("/serve/{video_id}/{filename:path}")
async def serve_download(video_id: str, filename: str):
    filepath = os.path.join(DOWNLOAD_DIR, video_id, filename)
    real_path = os.path.realpath(filepath)
    real_dir = os.path.realpath(DOWNLOAD_DIR)
    if not real_path.startswith(real_dir):
        return HTMLResponse("Forbidden", status_code=403)
    if not os.path.exists(real_path):
        return HTMLResponse("File not found", status_code=404)
    return FileResponse(real_path, filename=os.path.basename(real_path), media_type="application/octet-stream")

print("[YTDL] Cleaning old downloads...")
_cleanup_old_downloads()
print("[YTDL] All ready. HF Space launcher will call demo.launch()")