main.py

"""
╔══════════════════════════════════════════════════════════════╗
║   TITAN SMART DOWNLOADER — API v2.0                         ║
║   Backend: FastAPI + yt-dlp + FFmpeg                        ║
║   Fixes: Audio merge, Proxy download, Reliability           ║
╚══════════════════════════════════════════════════════════════╝
"""
from __future__ import annotations

import asyncio
import hashlib
import logging
import os
import shutil
import subprocess
import tempfile
import time
import urllib.parse
import urllib.request
from contextlib import asynccontextmanager
from typing import Any

import yt_dlp
from fastapi import FastAPI, HTTPException, Query
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pydantic import BaseModel

logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s")
log = logging.getLogger("titan-dl")

# ─── Config ───────────────────────────────────────────────────────────────────
MAX_CONCURRENT  = 10
CACHE_TTL       = 3 * 3600      # 3 hours
MAX_CACHE_ITEMS = 300
FFMPEG_PATH     = shutil.which("ffmpeg") or "/usr/bin/ffmpeg"
CHUNK_SIZE      = 1024 * 256    # 256 KB streaming chunks

UA = (
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
    "AppleWebKit/537.36 (KHTML, like Gecko) "
    "Chrome/124.0.0.0 Safari/537.36"
)

# ─── TTL Cache ─────────────────────────────────────────────────────────────────
class TTLCache:
    def __init__(self, ttl: int, max_items: int):
        self._s: dict[str, tuple[float, Any]] = {}
        self._ttl = ttl
        self._max = max_items
        self._lock = asyncio.Lock()

    async def get(self, key: str) -> Any | None:
        async with self._lock:
            if key not in self._s:
                return None
            ts, v = self._s[key]
            if time.time() - ts > self._ttl:
                del self._s[key]
                return None
            log.info(f"⚡ Cache HIT: {key[:16]}")
            return v

    async def set(self, key: str, value: Any):
        async with self._lock:
            if len(self._s) >= self._max:
                oldest = min(self._s, key=lambda k: self._s[k][0])
                del self._s[oldest]
            self._s[key] = (time.time(), value)

    async def stats(self) -> dict:
        async with self._lock:
            now = time.time()
            alive = sum(1 for ts, _ in self._s.values() if now - ts <= self._ttl)
            return {"total": len(self._s), "alive": alive, "max": self._max}


# ─── Concurrency Queue ────────────────────────────────────────────────────────
class Queue:
    def __init__(self, n: int):
        self._sem = asyncio.Semaphore(n)
        self._active = 0

    @asynccontextmanager
    async def slot(self):
        await self._sem.acquire()
        self._active += 1
        try:
            yield
        finally:
            self._active -= 1
            self._sem.release()

    @property
    def active(self):
        return self._active


cache = TTLCache(CACHE_TTL, MAX_CACHE_ITEMS)
queue = Queue(MAX_CONCURRENT)

# ─── yt-dlp base options ──────────────────────────────────────────────────────
YDL_OPTS_BASE = {
    "quiet":          True,
    "no_warnings":    True,
    "skip_download":  True,
    "extract_flat":   False,
    "noplaylist":     True,
    "socket_timeout": 20,
    "retries":        5,
    "fragment_retries": 5,
    "http_headers": {
        "User-Agent": UA,
        "Accept-Language": "en-US,en;q=0.9",
    },
}

# ─── Helpers ──────────────────────────────────────────────────────────────────
def _size(b) -> str:
    if not b:
        return ""
    for u in ["B", "KB", "MB", "GB"]:
        if b < 1024:
            return f"{b:.0f} {u}"
        b /= 1024
    return f"{b:.1f} GB"


def _formats(fmts: list) -> list[dict]:
    """
    Return usable formats ranked by quality.
    Priority:
      1. Truly merged (video+audio in single stream)  → direct download
      2. Video-only + best audio paired               → server-side FFmpeg merge
      3. Audio-only (best bitrate)                    → direct download
    """
    merged_formats = []
    video_only     = []
    audio_only     = []

    for f in fmts:
        vcodec = f.get("vcodec", "none") or "none"
        acodec = f.get("acodec", "none") or "none"
        url    = f.get("url", "")
        if not url:
            continue

        is_video = vcodec.lower() not in ("none", "")
        is_audio = acodec.lower() not in ("none", "")

        if is_video and is_audio:
            merged_formats.append(f)
        elif is_video:
            video_only.append(f)
        elif is_audio:
            audio_only.append(f)

    result      = []
    seen_heights= set()

    # ── 1. Real merged ─────────────────────────────────────────────────────
    merged_formats.sort(key=lambda x: x.get("height", 0) or 0, reverse=True)
    for f in merged_formats:
        height = f.get("height", 0) or 0
        if height in seen_heights:
            continue
        seen_heights.add(height)

        fps    = f.get("fps", 0) or 0
        size_b = f.get("filesize") or f.get("filesize_approx")
        lbl    = f"{height}p" if height else "Video"
        if fps > 30:
            lbl += f" {int(fps)}fps"

        result.append({
            "id":            f.get("format_id", ""),
            "label":         lbl,
            "type":          "video_merged",
            "ext":           f.get("ext", "mp4"),
            "height":        height,
            "fps":           int(fps) if fps else 0,
            "has_audio":     True,
            "needs_merge":   False,
            "size":          _size(size_b),
            "size_bytes":    size_b or 0,
            "url":           f["url"],
            "audio_url":     None,
        })

    # ── 2. Video-only + best audio → needs server-side merge ───────────────
    if video_only and audio_only:
        # pick best audio: highest abr, then asr
        best_audio = max(
            audio_only,
            key=lambda x: (x.get("abr", 0) or 0, x.get("asr", 0) or 0)
        )

        video_only.sort(key=lambda x: x.get("height", 0) or 0, reverse=True)
        for f in video_only:
            height = f.get("height", 0) or 0
            if height in seen_heights:
                continue
            seen_heights.add(height)

            fps    = f.get("fps", 0) or 0
            v_size = f.get("filesize") or f.get("filesize_approx") or 0
            a_size = best_audio.get("filesize") or best_audio.get("filesize_approx") or 0
            lbl    = f"{height}p" if height else "Video"
            if fps > 30:
                lbl += f" {int(fps)}fps"

            result.append({
                "id":          f"{f.get('format_id')}+{best_audio.get('format_id')}",
                "label":       lbl,
                "type":        "video_needs_merge",
                "ext":         "mp4",           # ffmpeg will re-mux to mp4
                "height":      height,
                "fps":         int(fps) if fps else 0,
                "has_audio":   True,
                "needs_merge": True,
                "size":        _size(v_size + a_size),
                "size_bytes":  v_size + a_size,
                "url":         f["url"],
                "audio_url":   best_audio["url"],
                "audio_ext":   best_audio.get("ext", "m4a"),
            })

    # ── 3. Audio only ─────────────────────────────────────────────────────
    audio_only.sort(key=lambda x: x.get("abr", 0) or 0, reverse=True)
    for f in audio_only[:1]:        # only best
        url = f.get("url", "")
        if not url:
            continue
        abr   = f.get("abr", 0) or 0
        size_b= f.get("filesize") or f.get("filesize_approx")
        result.append({
            "id":          f.get("format_id", ""),
            "label":       f"{int(abr)}kbps" if abr else "Audio",
            "type":        "audio_only",
            "ext":         f.get("ext", "m4a"),
            "height":      0,
            "fps":         0,
            "has_audio":   True,
            "needs_merge": False,
            "size":        _size(size_b),
            "size_bytes":  size_b or 0,
            "url":         url,
            "audio_url":   None,
        })

    return result


def _extract_info(url: str) -> dict:
    opts = YDL_OPTS_BASE.copy()

    # Per-platform format selection
    url_lower = url.lower()
    if "youtube" in url_lower or "youtu.be" in url_lower:
        # Get all formats so we can list them; yt-dlp picks best automatically
        opts["format"] = "bestvideo[ext=mp4]+bestaudio[ext=m4a]/bestvideo+bestaudio/best"
    else:
        opts["format"] = "best"

    with yt_dlp.YoutubeDL(opts) as ydl:
        info = ydl.extract_info(url, download=False)
        if not info:
            raise ValueError("لم يُعثر على معلومات")

        all_formats = info.get("formats", [])
        fmts        = _formats(all_formats)

        if not fmts:
            raise ValueError("لا توجد صيغ قابلة للتحميل")

        # Best thumbnail
        thumbs = info.get("thumbnails") or []
        thumb  = (
            sorted(thumbs, key=lambda t: (t.get("width") or 0), reverse=True)[0]["url"]
            if thumbs else info.get("thumbnail", "")
        )

        return {
            "id":           info.get("id", ""),
            "title":        info.get("title", "فيديو"),
            "thumbnail":    thumb,
            "duration":     info.get("duration"),
            "duration_str": info.get("duration_string", ""),
            "uploader":     info.get("uploader") or info.get("channel", ""),
            "platform":     info.get("extractor_key", ""),
            "webpage_url":  info.get("webpage_url", url),
            "view_count":   info.get("view_count"),
            "formats":      fmts,
        }


# ─── FastAPI Setup ─────────────────────────────────────────────────────────────
@asynccontextmanager
async def lifespan(app: FastAPI):
    ffmpeg_ok = os.path.isfile(FFMPEG_PATH)
    log.info(f"🚀 Titan Smart Downloader API v2.0 — ffmpeg: {'✅' if ffmpeg_ok else '❌ NOT FOUND'}")
    yield
    log.info("🛑 Shutdown")


app = FastAPI(
    title="Titan Smart Downloader",
    version="2.0.0",
    lifespan=lifespan,
    redirect_slashes=False,
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
    expose_headers=["Content-Disposition", "Content-Length", "X-Filename"],
)

# ─── Models ───────────────────────────────────────────────────────────────────
class ExtractReq(BaseModel):
    url: str


# ─── Routes ───────────────────────────────────────────────────────────────────
@app.get("/health")
async def health():
    return {
        "status":   "online",
        "version":  "2.0.0",
        "ffmpeg":   os.path.isfile(FFMPEG_PATH),
        "queue":    queue.active,
        "max":      MAX_CONCURRENT,
        "cache":    await cache.stats(),
    }


@app.post("/extract")
@app.post("/extract/")
async def extract(req: ExtractReq):
    url = req.url.strip()
    if not url.startswith(("http://", "https://")):
        raise HTTPException(400, "رابط غير صالح")

    key    = hashlib.sha256(url.encode()).hexdigest()
    cached = await cache.get(key)
    if cached:
        return {"success": True, "cached": True, "data": cached}

    async with queue.slot():
        log.info(f"🔍 [{queue.active}/{MAX_CONCURRENT}] {url[:80]}")
        loop = asyncio.get_event_loop()
        try:
            data = await asyncio.wait_for(
                loop.run_in_executor(None, _extract_info, url),
                timeout=50.0
            )
        except asyncio.TimeoutError:
            raise HTTPException(408, "انتهى وقت الاستخراج — جرّب مرة أخرى")
        except yt_dlp.utils.DownloadError as e:
            msg = str(e)
            if "Private video" in msg:
                raise HTTPException(403, "الفيديو خاص")
            if "rate-limit" in msg.lower():
                raise HTTPException(429, "تم تجاوز الحد المسموح — جرّب بعد شوية")
            if "not available" in msg:
                raise HTTPException(404, "الفيديو غير متاح في منطقتك")
            if "login required" in msg.lower():
                raise HTTPException(401, "يتطلب تسجيل دخول")
            if "copyright" in msg.lower():
                raise HTTPException(451, "محجوب بسبب حقوق النشر")
            if "Unable to extract" in msg:
                raise HTTPException(422, "المنصة حظرت الاستخراج — جرّب رابط تاني")
            raise HTTPException(422, f"لا يمكن استخراج الرابط: {msg[:150]}")
        except ValueError as e:
            raise HTTPException(404, str(e))
        except Exception as e:
            log.error(f"💥 {e}", exc_info=True)
            raise HTTPException(500, "خطأ داخلي في الاستخراج")

    await cache.set(key, data)
    log.info(f"✅ {len(data['formats'])} formats — {data['title'][:50]}")
    return {"success": True, "cached": False, "data": data}


# ─── Proxy Download (solves CORS / new-tab problem) ───────────────────────────
@app.get("/download")
async def proxy_download(
    url:      str = Query(..., description="Direct media URL"),
    filename: str = Query("video.mp4", description="Download filename"),
):
    """
    Stream a remote media URL to the client browser with proper
    Content-Disposition so the browser saves the file instead of
    opening a new tab.
    """
    if not url.startswith(("http://", "https://")):
        raise HTTPException(400, "رابط غير صالح")

    # Keep Arabic/Unicode chars; only strip filesystem-dangerous chars
    safe_name = "".join(
        c for c in filename if c not in '\\/:"*?<>|'
    ).strip() or "video.mp4"

    headers_req = {
        "User-Agent": UA,
        "Referer":    url.split("/")[0] + "//" + url.split("/")[2] + "/",
    }

    async def stream_generator():
        req = urllib.request.Request(url, headers=headers_req)
        try:
            with urllib.request.urlopen(req, timeout=30) as resp:
                while True:
                    chunk = resp.read(CHUNK_SIZE)
                    if not chunk:
                        break
                    yield chunk
        except Exception as e:
            log.error(f"Proxy stream error: {e}")
            raise

    try:
        # Peek at headers for content-type
        req      = urllib.request.Request(url, headers=headers_req, method="HEAD")
        try:
            with urllib.request.urlopen(req, timeout=10) as r:
                content_type   = r.headers.get("Content-Type", "video/mp4")
                content_length = r.headers.get("Content-Length")
        except Exception:
            content_type   = "video/mp4"
            content_length = None

        ascii_fn    = safe_name.encode("ascii", "ignore").decode() or "video.mp4"
        utf8_fn     = urllib.parse.quote(safe_name.encode("utf-8"))
        content_disp= f"attachment; filename=\"{ascii_fn}\"; filename*=UTF-8''{utf8_fn}"

        response_headers = {
            "Content-Disposition": content_disp,
            "X-Filename":          urllib.parse.quote(safe_name),
            "Cache-Control":       "no-cache",
        }
        if content_length:
            response_headers["Content-Length"] = content_length

        return StreamingResponse(
            stream_generator(),
            media_type=content_type,
            headers=response_headers,
        )
    except Exception as e:
        log.error(f"Proxy download error: {e}", exc_info=True)
        raise HTTPException(502, "تعذّر جلب الملف — جرّب مرة أخرى")


# ─── Server-side FFmpeg Merge (solves audio problem) ─────────────────────────
@app.get("/merge")
async def merge_and_download(
    video_url: str = Query(..., description="Video-only stream URL"),
    audio_url: str = Query(..., description="Audio-only stream URL"),
    filename:  str = Query("video.mp4", description="Output filename"),
):
    """
    Stream-merge architecture:
    1. FFmpeg reads video + audio URLs directly (no full pre-download)
    2. FFmpeg writes merged output to a temp file
    3. We stream that file to the client chunk by chunk
    This handles films of any size without RAM or timeout issues.
    """
    if not video_url.startswith(("http://", "https://")):
        raise HTTPException(400, "رابط فيديو غير صالح")
    if not audio_url.startswith(("http://", "https://")):
        raise HTTPException(400, "رابط صوت غير صالح")

    if not os.path.isfile(FFMPEG_PATH):
        raise HTTPException(503, "FFmpeg غير متاح على السيرفر")

    # Sanitize filename
    safe_name = "".join(
        c for c in filename if c not in '\\/:"*?<>|'
    ).strip() or "video.mp4"
    if not safe_name.lower().endswith(".mp4"):
        safe_name = safe_name.rsplit(".", 1)[0] + ".mp4"

    # RFC 5987 headers — safe for Arabic/Unicode filenames
    ascii_name   = safe_name.encode("ascii", "ignore").decode() or "video.mp4"
    utf8_encoded = urllib.parse.quote(safe_name.encode("utf-8"))
    content_disp = f"attachment; filename=\"{ascii_name}\"; filename*=UTF-8''{utf8_encoded}"

    # ── Temp dir lives until response is fully sent ────────────────────────
    tmp_dir = tempfile.mkdtemp()
    o_path  = os.path.join(tmp_dir, "output.mp4")

    async def _run_ffmpeg_and_stream():
        """
        Phase 1 — FFmpeg reads directly from CDN URLs (no pre-download).
                   Writes merged MP4 to o_path.
        Phase 2 — Stream o_path to client in chunks, then clean up.
        """
        try:
            # ── Phase 1: FFmpeg merge ──────────────────────────────────────
            # -reconnect flags make ffmpeg retry on network hiccups
            cmd = [
                FFMPEG_PATH,
                "-y",
                "-reconnect", "1",
                "-reconnect_streamed", "1",
                "-reconnect_delay_max", "10",
                "-user_agent", UA,
                "-i", video_url,
                "-reconnect", "1",
                "-reconnect_streamed", "1",
                "-reconnect_delay_max", "10",
                "-user_agent", UA,
                "-i", audio_url,
                "-c:v", "copy",           # no re-encode = fast
                "-c:a", "aac",
                "-movflags", "+faststart",
                "-shortest",
                o_path,
            ]

            log.info(f"🔧 FFmpeg stream-merge → {safe_name}")

            # Run FFmpeg in thread pool — no timeout: FFmpeg handles its own
            # reconnect logic; we trust it to finish or fail on its own.
            loop   = asyncio.get_event_loop()
            result = await loop.run_in_executor(
                None,
                lambda: subprocess.run(cmd, capture_output=True)
            )

            if result.returncode != 0:
                err = result.stderr.decode(errors="ignore")[-600:]
                log.error(f"FFmpeg failed (rc={result.returncode}): {err}")
                raise RuntimeError(f"FFmpeg rc={result.returncode}: {err[:200]}")

            if not os.path.exists(o_path) or os.path.getsize(o_path) == 0:
                raise RuntimeError("FFmpeg produced empty output")

            size = os.path.getsize(o_path)
            log.info(f"✅ Merge complete → {safe_name} ({size/1024/1024:.1f} MB)")

            # ── Phase 2: Stream to client ──────────────────────────────────
            with open(o_path, "rb") as fh:
                while True:
                    chunk = fh.read(CHUNK_SIZE)
                    if not chunk:
                        break
                    yield chunk

        except Exception as e:
            log.error(f"Merge/stream error: {e}", exc_info=True)
            raise
        finally:
            # Clean up temp files regardless of success/failure
            try:
                shutil.rmtree(tmp_dir, ignore_errors=True)
            except Exception:
                pass

    # We don't know Content-Length upfront (streaming), so omit it.
    # Browsers handle chunked transfer encoding fine.
    return StreamingResponse(
        _run_ffmpeg_and_stream(),
        media_type="video/mp4",
        headers={
            "Content-Disposition": content_disp,
            "X-Filename":          urllib.parse.quote(safe_name),
            "Cache-Control":       "no-cache",
            "Transfer-Encoding":   "chunked",
        },
    )


# ─── Cache Management ─────────────────────────────────────────────────────────
@app.get("/cache/stats")
async def cache_stats():
    return await cache.stats()


@app.delete("/cache/clear")
async def cache_clear():
    async with cache._lock:
        n = len(cache._s)
        cache._s.clear()
    return {"cleared": n}