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__)
CANONICAL_STATIC_BASE = "https://adgen.static.lookfinity.com"

# ---------- COLOR FILTERS ----------
SAFE_EFFECTS = {
    "warm": "eq=brightness=0.01:contrast=1.02:saturation=1.08,colorbalance=rs=0.04:bs=-0.03",
    "fresh": "eq=brightness=0.015:contrast=1.03:saturation=1.12,colorbalance=gs=0.03:bs=0.02",
    "bright_clean": "eq=brightness=0.02:contrast=1.03:saturation=1.04",
    "natural": "eq=brightness=0.006:contrast=1.015:saturation=1.025",
    "product_clean": "eq=brightness=0.018:contrast=1.04:saturation=1.07,unsharp=3:3:0.2",
    "light_sharpen": "eq=brightness=0.005:contrast=1.02:saturation=1.04,unsharp=3:3:0.18",
    "soft": "eq=brightness=0.01:contrast=0.98:saturation=1.04",
    "muted": "eq=brightness=0.005:contrast=0.97:saturation=0.92",
    "cinematic": "eq=brightness=-0.005:contrast=1.06:saturation=0.98,colorbalance=rs=0.015:bs=0.025",
    "subtle": "eq=brightness=0.004:contrast=1.01:saturation=1.015",
}


def _public_base_url() -> str:
    base = (os.getenv("NEW_BASE") or "").strip().rstrip("/")
    return base or CANONICAL_STATIC_BASE

# ---------- 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",
    variation_idx: int = 0,
    border_width: int = 0,
    border_color: str = "black",
    effect_name: Optional[str] = None,
) -> str:
    if effect_name and effect_name in SAFE_EFFECTS:
        vf = SAFE_EFFECTS[effect_name]
    else:
        import random
        k_b, k_c, k_h, k_s = [random.randint(-3, 3) for _ in range(4)]
        brightness = max(-1.0, min(1.0, k_b * 0.2))
        contrast = max(0.0, min(2.0, 1.0 + k_c * 0.03))
        hue_rad = math.radians(k_h * 3.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}"

    if border_width > 0:
        ffmpeg_color = border_color.replace("#", "0x") if border_color.startswith("#") else border_color
        bw = int(border_width)
        vf += f",pad=iw+{bw*2}:ih+{bw*2}:{bw}:{bw}:{ffmpeg_color}"

    inp = Path(input_path)
    suffix = f"_augmented_{variation_idx}" if variation_idx > 0 else "_augmented"
    out = Path(output_path) / f"{inp.stem}{suffix}.mp4" if output_path else inp.with_name(f"{inp.stem}{suffix}.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"{_public_base_url()}/{key}"

# ---------- PROCESS + UPLOAD ----------
def process_video_with_hash_info(
    input_path: str,
    output_path: Optional[str] = None,
    variation_idx: int = 0,
    border_width: int = 0,
    border_color: str = "black",
    effect_name: Optional[str] = None,
) -> dict:
    logger.info(f"[process] Starting process_video_with_hash_info for {Path(input_path).name}")
    logger.info(f"[process] Step 1/5 — computing input hash")
    in_sha = _sha256(Path(input_path))
    in_id = _hashid_short(in_sha)

    logger.info(f"[process] Step 2/5 — augmenting video (effect={effect_name or 'random'})")
    out_path = augment_video_random(
        input_path=input_path,
        output_path=output_path,
        variation_idx=variation_idx,
        border_width=border_width,
        border_color=border_color,
        effect_name=effect_name,
    )

    logger.info(f"[process] Step 3/5 — computing output hash")
    out_sha = _sha256(Path(out_path))
    out_id = _hashid_short(out_sha)

    logger.info(f"[process] Step 4/5 — extracting thumbnail")
    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 = ""

    logger.info(f"[process] Step 5/5 — uploading to R2")
    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

    logger.info(f"[process] Done. output_hashid={out_id}, r2_url={r2_url}")
    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,
        "effect": effect_name or "random",
    }