generator_function/video_hashid_function.py

from pathlib import Path
from typing import Optional
import hashlib, base64, binascii, math, subprocess, shutil, os, logging, uuid, tempfile
from dotenv import load_dotenv
from imageio_ffmpeg import get_ffmpeg_exe
import boto3

load_dotenv()
logger = logging.getLogger(__name__)

# ---------- HASH UTILS ----------
def _sha256(path: Path, chunk: int = 1 << 20) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        for b in iter(lambda: f.read(chunk), b""):
            h.update(b)
    return h.hexdigest()

def _hashid_short(sha256_hex: str, length: int = 12) -> str:
    b = binascii.unhexlify(sha256_hex)
    return base64.urlsafe_b64encode(b).decode().rstrip("=")[:length]

# ---------- FFMPEG ----------
def _ffmpeg_bin() -> str:
    try:
        return get_ffmpeg_exe()
    except Exception:
        p = os.environ.get("FFMPEG_BIN") or shutil.which("ffmpeg")
        if not p:
            raise FileNotFoundError("ffmpeg not found. Install it or set FFMPEG_BIN")
        return p

FFMPEG = _ffmpeg_bin()

def _run_ffmpeg(cmd: list[str]) -> None:
    cmd = cmd[:]
    cmd[0] = FFMPEG
    proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    if proc.returncode != 0:
        raise RuntimeError(proc.stderr.decode(errors="ignore") or "ffmpeg failed")

# ---------- THUMBNAIL ----------
def extract_thumbnail(video_path: str, time_position: str = "00:00:01") -> str:
    """
    Extracts a thumbnail (JPEG) from the video and returns it as a base64 string.
    Default frame at 1 second.
    """
    tmp_thumb = Path(tempfile.gettempdir()) / f"{Path(video_path).stem}_thumb.jpg"
    cmd = [
        "ffmpeg", "-nostdin", "-y",
        "-ss", time_position,
        "-i", video_path,
        "-frames:v", "1",
        "-q:v", "2",
        str(tmp_thumb),
    ]
    _run_ffmpeg(cmd)

    with open(tmp_thumb, "rb") as f:
        encoded = base64.b64encode(f.read()).decode("utf-8")
    return encoded

# ---------- AUGMENT VIDEO ----------
def augment_video_random(
    *,
    input_path: str,
    output_path: Optional[str] = None,
    crf: int = 20,
    preset: str = "medium",
) -> str:
    import random
    k_b, k_c, k_h, k_s = [random.randint(-5, 5) for _ in range(4)]
    brightness = max(-1.0, min(1.0, k_b * 0.05))
    contrast = max(0.0, min(2.0, 1.0 + k_c * 0.05))
    hue_rad = math.radians(k_h * 5.0)
    sat_scale = max(0.0, 1.0 + k_s * 0.05)

    vf = f"hue=h={hue_rad:.6f}:s={sat_scale:.4f},eq=contrast={contrast:.4f}:brightness={brightness:.4f}"

    inp = Path(input_path)
    out = Path(output_path) / f"{inp.stem}_augmented.mp4" if output_path else inp.with_name(f"{inp.stem}_augmented.mp4")

    cmd = [
        "ffmpeg", "-nostdin", "-y",
        "-i", str(inp),
        "-map", "0:v:0", "-map", "0:a?",
        "-vf", vf,
        "-c:v", "libx264", "-preset", preset, "-crf", str(crf),
        "-pix_fmt", "yuv420p",
        "-c:a", "copy",
        "-movflags", "+faststart",
        str(out),
    ]
    _run_ffmpeg(cmd)
    return str(out)

# ---------- VIDEO UPLOAD TO R2 ----------
def upload_video_to_r2(video_bytes: bytes, file_name: str) -> str:
    s3 = boto3.client(
        "s3",
        endpoint_url=os.getenv("R2_ENDPOINT"),
        aws_access_key_id=os.getenv("R2_ACCESS_KEY"),
        aws_secret_access_key=os.getenv("R2_SECRET_KEY"),
        region_name="auto",
    )

    bucket = os.getenv("R2_BUCKET_NAME")
    key = f"videos/{uuid.uuid4()}_{file_name}"

    s3.put_object(Bucket=bucket, Key=key, Body=video_bytes, ContentType="video/mp4")

    return f"{os.getenv('NEW_BASE').rstrip('/')}/{key}"

# ---------- PROCESS + UPLOAD ----------
def process_video_with_hash_info(input_path: str, output_path: Optional[str] = None) -> dict:
    in_sha = _sha256(Path(input_path))
    in_id = _hashid_short(in_sha)

    # Augment and get output
    out_path = augment_video_random(input_path=input_path, output_path=output_path)

    # output hashid
    out_sha = _sha256(Path(out_path))
    out_id = _hashid_short(out_sha)

    # copy file to safe tmp path
    safe_tmp = Path(tempfile.gettempdir()) / Path(out_path).name
    shutil.copy(out_path, safe_tmp)

    # thumbnail
    try:
        thumbnail_b64 = extract_thumbnail(str(safe_tmp))
    except Exception as e:
        logger.error(f"Thumbnail extraction failed: {e}")
        thumbnail_b64 = ""


    try:
        with open(safe_tmp, "rb") as f:
            video_bytes = f.read()
        r2_url = upload_video_to_r2(video_bytes, Path(safe_tmp).name)
    except Exception as e:
        logger.error(f"Upload to R2 failed: {e}")
        r2_url = None

    return {
        "input_hashid": in_id,
        "output_name": Path(safe_tmp).name,
        "output_path": str(safe_tmp),
        "output_hashid": out_id,
        "output_r2_url": r2_url,
        "thumbnail": thumbnail_b64,
    }