ingestion/resolver.py

import os
import uuid
import requests
from pathlib import Path
from typing import Optional, Union
from fastapi import UploadFile

# ==============================
# STORAGE CONFIG
# ==============================

BASE_DIR = Path(__file__).resolve().parents[1]
UPLOAD_DIR = str(BASE_DIR / "jobs" / "uploads")
os.makedirs(UPLOAD_DIR, exist_ok=True)

CHUNK_SIZE = 1024 * 1024  # 1MB streaming for large media


# ==============================
# CORE RESOLVER
# ==============================

def resolve_input(
    source: Optional[str] = None,
    upload: Optional[UploadFile] = None,
    raw_bytes: Optional[bytes] = None
) -> str:
    """
    Universal ingestion layer for all pipeline systems.

    Supports:
    - UploadFile (FastAPI / Gradio)
    - URL download (http/https)
    - Local filesystem path
    - Raw bytes input (future automation nodes)
    """

    # ----------------------------------------
    # CASE 1: UploadFile (Gradio / FastAPI)
    # ----------------------------------------
    if upload is not None:
        filename = f"{uuid.uuid4()}_{upload.filename or 'upload.mp4'}"
        path = os.path.join(UPLOAD_DIR, filename)

        with open(path, "wb") as f:
            while True:
                chunk = upload.file.read(CHUNK_SIZE)
                if not chunk:
                    break
                f.write(chunk)

        return path

    # ----------------------------------------
    # CASE 2: Raw bytes (automation / webhook)
    # ----------------------------------------
    if raw_bytes is not None:
        filename = f"{uuid.uuid4()}.mp4"
        path = os.path.join(UPLOAD_DIR, filename)

        with open(path, "wb") as f:
            f.write(raw_bytes)

        return path

    # ----------------------------------------
    # CASE 3: URL input (YouTube, TikTok, direct mp4)
    # ----------------------------------------
    if source and source.startswith(("http://", "https://")):

        filename = f"{uuid.uuid4()}.mp4"
        path = os.path.join(UPLOAD_DIR, filename)

        headers = {
            "User-Agent": "Mozilla/5.0 (compatible; BasyxBot/1.0)"
        }

        with requests.get(source, stream=True, headers=headers, timeout=60) as r:
            r.raise_for_status()

            with open(path, "wb") as f:
                for chunk in r.iter_content(chunk_size=CHUNK_SIZE):
                    if chunk:
                        f.write(chunk)

        return path

    # ----------------------------------------
    # CASE 4: Local file path
    # ----------------------------------------
    if source and os.path.exists(source):
        return source

    # ----------------------------------------
    # INVALID INPUT HANDLING
    # ----------------------------------------
    raise ValueError(
        "resolve_input failed: no valid source, upload, or raw_bytes provided"
    )


# ==============================
# OPTIONAL HELPERS (V11 READY)
# ==============================

def detect_input_type(source: str) -> str:
    """
    Lightweight classifier for routing decisions upstream.
    """

    if source.startswith(("http://", "https://")):
        return "url"

    if os.path.exists(source):
        return "file"

    return "unknown"


def normalize_source(source: str) -> str:
    """
    Cleans input strings for downstream consistency.
    """
    if not source:
        return source

    return source.strip()