main.py

import os
import shutil
import asyncio
import subprocess
import numpy as np
from pathlib import Path
from typing import Optional
from contextlib import asynccontextmanager

from fastapi import FastAPI, UploadFile, File, HTTPException, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse

# ── paths ──────────────────────────────────────────────────────────────────────
TMP = Path("/tmp/clipcut")
TMP.mkdir(parents=True, exist_ok=True)
os.environ["INSIGHTFACE_HOME"] = "/tmp"

# ── lazy model ─────────────────────────────────────────────────────────────────
_face_app = None
_swapper = None
_yolo = None
_reid = None

def get_face_app():
    global _face_app
    if _face_app is None:
        from insightface.app import FaceAnalysis
        # Upgraded from buffalo_sc to buffalo_l — larger, more accurate model
        _face_app = FaceAnalysis(name="buffalo_l", root="/tmp/insightface",
                                  providers=["CPUExecutionProvider"])
        _face_app.prepare(ctx_id=-1, det_size=(640, 640))
    return _face_app

def get_yolo():
    """YOLOv8-nano — fast person detector for precise body bounding boxes."""
    global _yolo
    if _yolo is None:
        from ultralytics import YOLO
        model_path = Path("/tmp/insightface/yolov8n.pt")
        model_path.parent.mkdir(parents=True, exist_ok=True)
        _yolo = YOLO("yolov8n.pt")  # auto-downloads on first use (~6MB)
    return _yolo

def get_reid():
    """Lightweight person re-identification model (ONNX) — matches body identity
    across frames using appearance features, more robust than color histograms."""
    global _reid
    if _reid is None:
        import onnxruntime
        model_path = Path("/tmp/insightface/models/reid_osnet.onnx")
        model_path.parent.mkdir(parents=True, exist_ok=True)

        if not model_path.exists():
            import requests
            # NOTE: verified public ONNX mirrors of OSNet ReID model.
            # If these ever go down, code safely falls back to color histogram matching.
            urls = [
                "https://huggingface.co/nomnomnonull/osnet_x0_25_msmt17_onnx/resolve/main/osnet_x0_25_msmt17.onnx",
                "https://huggingface.co/spaces/HaHaBill/LandShapes-Antique/resolve/main/osnet_x0_25.onnx",
            ]
            downloaded = False
            for url in urls:
                try:
                    print(f"[ReID] Downloading model from {url}")
                    r = requests.get(url, stream=True, timeout=60)
                    if r.status_code == 200 and 'text/html' not in r.headers.get('content-type',''):
                        with open(model_path, "wb") as f:
                            for chunk in r.iter_content(chunk_size=65536):
                                if chunk: f.write(chunk)
                        if model_path.stat().st_size > 500_000:  # sanity check — real model, not error page
                            downloaded = True
                            print(f"[ReID] Model ready: {model_path.stat().st_size} bytes")
                            break
                        else:
                            model_path.unlink()
                except Exception as e:
                    print(f"[ReID] Failed from {url}: {e}")
                    continue
            if not downloaded:
                print("[ReID] No working ReID model source — falling back to color histogram matching for body tracking.")
                return None

        try:
            _reid = onnxruntime.InferenceSession(str(model_path), providers=["CPUExecutionProvider"])
        except Exception as e:
            print(f"[ReID] Failed to load model: {e}")
            return None
    return _reid

def reid_embedding(frame, bbox) -> Optional[np.ndarray]:
    """Extract a Re-ID appearance embedding for a person crop."""
    import cv2
    sess = get_reid()
    if sess is None:
        return None
    x1, y1, x2, y2 = [int(v) for v in bbox]
    x1, y1 = max(0, x1), max(0, y1)
    x2, y2 = min(frame.shape[1], x2), min(frame.shape[0], y2)
    crop = frame[y1:y2, x1:x2]
    if crop.size == 0:
        return None
    try:
        img = cv2.resize(crop, (128, 256))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
        mean = np.array([0.485, 0.456, 0.406])
        std = np.array([0.229, 0.224, 0.225])
        img = (img - mean) / std
        img = img.transpose(2, 0, 1)[None].astype(np.float32)
        input_name = sess.get_inputs()[0].name
        out = sess.run(None, {input_name: img})[0]
        emb = out.flatten()
        norm = np.linalg.norm(emb)
        return emb / norm if norm > 0 else emb
    except Exception as e:
        print(f"[ReID] Embedding extraction failed: {e}")
        return None

def get_swapper():
    global _swapper
    if _swapper is None:
        import onnxruntime
        model_path = Path("/tmp/insightface/models/inswapper_128.onnx")
        model_path.parent.mkdir(parents=True, exist_ok=True)

        if not model_path.exists():
            import requests
            urls = [
                "https://huggingface.co/ezioruan/inswapper_128.onnx/resolve/main/inswapper_128.onnx",
                "https://huggingface.co/Patil/inswapper/resolve/main/inswapper_128.onnx",
                "https://huggingface.co/Aitrepreneur/insightface/resolve/fd887cdef0c73f32251198b8160d6771ac413fc0/inswapper_128.onnx",
            ]
            downloaded = False
            for url in urls:
                try:
                    print(f"[FaceSwap] Downloading model from {url}")
                    r = requests.get(url, stream=True, timeout=300)
                    if r.status_code == 200:
                        with open(model_path, "wb") as f:
                            for chunk in r.iter_content(chunk_size=65536):
                                if chunk: f.write(chunk)
                        if model_path.stat().st_size > 100_000_000:  # at least 100MB
                            print(f"[FaceSwap] Model downloaded: {model_path.stat().st_size} bytes")
                            downloaded = True
                            break
                        else:
                            model_path.unlink()  # too small, probably error page
                except Exception as e:
                    print(f"[FaceSwap] Failed from {url}: {e}")
                    continue
            if not downloaded:
                raise Exception("Could not download face swap model. Try again later.")

        _swapper = onnxruntime.InferenceSession(
            str(model_path), providers=["CPUExecutionProvider"]
        )
        # Wrap in insightface compatible interface
        import insightface
        _swapper = insightface.model_zoo.get_model(str(model_path), providers=["CPUExecutionProvider"])

    return _swapper

# ── job store ──────────────────────────────────────────────────────────────────
jobs: dict[str, dict] = {}

@asynccontextmanager
async def lifespan(app):
    yield

app = FastAPI(lifespan=lifespan)
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])

# ── utils ──────────────────────────────────────────────────────────────────────

def cosine_sim(a, b) -> float:
    a, b = a.flatten(), b.flatten()
    d = np.linalg.norm(a) * np.linalg.norm(b)
    return float(np.dot(a, b) / d) if d > 0 else 0.0

def clothing_hist(frame, bbox) -> Optional[np.ndarray]:
    import cv2
    x1, y1, x2, y2 = int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])
    region = frame[y1 + int((y2-y1)*0.3):y2, x1:x2]
    if region.size == 0:
        return None
    hsv = cv2.cvtColor(region, cv2.COLOR_BGR2HSV)
    h = cv2.calcHist([hsv], [0,1], None, [16,16], [0,180,0,256])
    cv2.normalize(h, h)
    return h.flatten()

def hist_sim(h1, h2) -> float:
    import cv2
    if h1 is None or h2 is None:
        return 0.0
    return float(cv2.compareHist(h1.astype(np.float32).reshape(-1,1),
                                  h2.astype(np.float32).reshape(-1,1),
                                  cv2.HISTCMP_CORREL))

def extract_ref_embedding(ref_path: Path) -> Optional[np.ndarray]:
    import cv2
    img = cv2.imread(str(ref_path))
    if img is None:
        return None
    fa = get_face_app()

    # Try original size first
    faces = fa.get(img)
    if not faces:
        # Try resized versions
        for scale in [1.5, 2.0, 0.75]:
            h, w = img.shape[:2]
            resized = cv2.resize(img, (int(w*scale), int(h*scale)))
            faces = fa.get(resized)
            if faces:
                break
    if not faces:
        return None
    largest = max(faces, key=lambda f: (f.bbox[2]-f.bbox[0])*(f.bbox[3]-f.bbox[1]))
    return largest.normed_embedding

def video_info(path: Path) -> dict:
    import json
    r = subprocess.run([
        "ffprobe", "-v", "error", "-select_streams", "v:0",
        "-show_entries", "stream=width,height,r_frame_rate",
        "-show_entries", "format=duration", "-of", "json", str(path)
    ], capture_output=True, text=True)
    d = json.loads(r.stdout)
    s = d.get("streams", [{}])[0]
    num, den = s.get("r_frame_rate", "30/1").split("/")
    return {
        "fps": float(num)/float(den),
        "duration": float(d.get("format",{}).get("duration", 0)),
        "width": int(s.get("width", 1080)),
        "height": int(s.get("height", 1920))
    }

def crop_size(w, h) -> tuple[int, int]:
    if (w/h) > (9/16):
        ch = h - h%2
        cw = int(ch*9/16) - int(ch*9/16)%2
    else:
        cw = w - w%2
        ch = int(cw*16/9) - int(cw*16/9)%2
    return cw, ch

def clamp(cx, cy, cw, ch, fw, fh) -> tuple[int, int]:
    x = max(0, min(cx - cw//2, fw - cw))
    y = max(0, min(cy - ch//2, fh - ch))
    return x, y

# ── scan video ─────────────────────────────────────────────────────────────────

def scan_video(video_path: Path, ref_emb: np.ndarray, fps: float,
               job_id: str, p0=15, p1=82, label="Scanning"):
    import cv2
    import time
    fa = get_face_app()
    yolo = get_yolo()
    FACE_THRESH = 0.15
    REID_THRESH = 0.45   # cosine sim threshold for Re-ID body matching
    BODY_WINDOW = fps * 5  # slightly longer window since Re-ID is more reliable

    cap = cv2.VideoCapture(str(video_path))
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    fw = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    fh = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

    timestamps, positions = [], {}
    ref_reid_emb = None       # Re-ID embedding of the confirmed person (updated when face matches)
    ref_hist = None           # color histogram fallback if Re-ID model unavailable
    last_face_frame = -9999
    idx = 0
    start_time = time.time()
    preview_dir = TMP / job_id
    PREVIEW_EVERY = max(1, total // 20)

    while True:
        if jobs.get(job_id, {}).get("cancelled"):
            cap.release()
            return [], {}

        ret, frame = cap.read()
        if not ret:
            break

        if idx % 2 == 0:
            # ── Step 1: YOLO finds all person bounding boxes (precise, fast) ──
            person_boxes = []
            try:
                results = yolo(frame, classes=[0], verbose=False, conf=0.35)  # class 0 = person
                for r in results:
                    for box in r.boxes:
                        xyxy = box.xyxy[0].cpu().numpy()
                        person_boxes.append(xyxy)
            except Exception:
                person_boxes = []

            # ── Step 2: Face detection + matching (ArcFace buffalo_l) ──
            faces = fa.get(frame)
            if not faces and fw < 1280:
                upscaled = cv2.resize(frame, (fw*2, fh*2))
                faces_up = fa.get(upscaled)
                for f in faces_up:
                    f.bbox = f.bbox / 2
                faces = faces_up if faces_up else faces

            matched = False
            cx_r, cy_r = 0.5, 0.42
            match_bbox = None

            for face in faces:
                if cosine_sim(face.normed_embedding, ref_emb) >= FACE_THRESH:
                    matched = True
                    last_face_frame = idx
                    x1, y1, x2, y2 = face.bbox

                    # Find the YOLO person box that contains this face (precise body crop)
                    face_cx, face_cy = (x1+x2)/2, (y1+y2)/2
                    best_pbox = None
                    for pbox in person_boxes:
                        px1, py1, px2, py2 = pbox
                        if px1 <= face_cx <= px2 and py1 <= face_cy <= py2:
                            best_pbox = pbox
                            break

                    if best_pbox is not None:
                        match_bbox = best_pbox
                        px1, py1, px2, py2 = best_pbox
                        cx_r = ((px1+px2)/2) / fw
                        cy_r = min(0.88, (py1 + (py2-py1)*0.4) / fh)
                    else:
                        fh_box = y2 - y1
                        cx_r = ((x1+x2)/2) / fw
                        cy_r = min(0.88, ((y1+y2)/2 + fh_box*1.2) / fh)

                    # Update Re-ID reference embedding using the matched person's body box
                    if match_bbox is not None:
                        r_emb = reid_embedding(frame, match_bbox)
                        if r_emb is not None:
                            ref_reid_emb = r_emb
                        else:
                            h = clothing_hist(frame, face.bbox)
                            if h is not None:
                                ref_hist = h
                    break

            # ── Step 3: Body fallback — Re-ID first, color histogram as backup ──
            if not matched and (idx - last_face_frame) <= BODY_WINDOW and person_boxes:
                if ref_reid_emb is not None:
                    best_sim, best_pbox = 0.0, None
                    for pbox in person_boxes:
                        p_emb = reid_embedding(frame, pbox)
                        if p_emb is not None:
                            sim = cosine_sim(p_emb, ref_reid_emb)
                            if sim > best_sim:
                                best_sim, best_pbox = sim, pbox
                    if best_sim >= REID_THRESH and best_pbox is not None:
                        matched = True
                        px1, py1, px2, py2 = best_pbox
                        cx_r = ((px1+px2)/2) / fw
                        cy_r = min(0.88, (py1 + (py2-py1)*0.4) / fh)
                elif ref_hist is not None:
                    best_s, best_bbox = 0.0, None
                    for pbox in person_boxes:
                        s = hist_sim(ref_hist, clothing_hist(frame, pbox))
                        if s > best_s:
                            best_s, best_bbox = s, pbox
                    if best_s >= 0.50 and best_bbox is not None:
                        matched = True
                        x1, y1, x2, y2 = best_bbox
                        cx_r = ((x1+x2)/2) / fw
                        cy_r = min(0.88, ((y1+y2)/2) / fh)

            if matched:
                ts = idx / fps
                timestamps.append(ts)
                positions[ts] = (cx_r, cy_r)

            if idx % PREVIEW_EVERY == 0 and frame is not None:
                try:
                    preview_path = preview_dir / "preview.jpg"
                    small = cv2.resize(frame, (360, int(fh * 360 / fw)))
                    cv2.imwrite(str(preview_path), small, [cv2.IMWRITE_JPEG_QUALITY, 70])
                except Exception:
                    pass

        elapsed = time.time() - start_time
        prange = p1 - p0
        prog = p0 + int((idx / max(total, 1)) * prange)
        if idx > 0 and elapsed > 0:
            fps_proc = idx / elapsed
            frames_left = total - idx
            eta_sec = int(frames_left / fps_proc)
            if eta_sec >= 60:
                eta_str = f"{eta_sec//60}m {eta_sec%60}s"
            else:
                eta_str = f"{eta_sec}s"
        else:
            eta_str = "calculating..."

        jobs[job_id]["progress"] = prog
        jobs[job_id]["status_text"] = f"{label}: {idx}/{total} frames"
        jobs[job_id]["eta"] = eta_str
        idx += 1

    cap.release()
    return timestamps, positions

# ── intervals ──────────────────────────────────────────────────────────────────

def to_intervals(timestamps, fps, gap=3.0, min_dur=0.3):
    if not timestamps:
        return []
    fw = 1.0/fps
    intervals, start, end = [], timestamps[0], timestamps[0]+fw
    for ts in timestamps[1:]:
        if ts - end <= gap:
            end = ts + fw
        else:
            if end-start >= min_dur:
                intervals.append((start, end))
            start, end = ts, ts+fw
    if end-start >= min_dur:
        intervals.append((start, end))
    return intervals

# ── cut clips with dynamic crop ────────────────────────────────────────────────

def cut_clip(video_path, start, end, positions, vi, out_path) -> bool:
    w, h = vi["width"], vi["height"]
    fps = vi["fps"]
    cw, ch = crop_size(w, h)
    dur = end - start

    # Gather positions for this interval
    pts = {ts: pos for ts, pos in positions.items() if start-0.1 <= ts <= end+0.1}

    if len(pts) >= 2:
        sorted_ts = sorted(pts.keys())
        rel = [(ts-start, pts[ts]) for ts in sorted_ts]

        def interp(t):
            if t <= rel[0][0]: return rel[0][1]
            if t >= rel[-1][0]: return rel[-1][1]
            for i in range(len(rel)-1):
                t0, p0 = rel[i]; t1, p1 = rel[i+1]
                if t0 <= t <= t1:
                    a = (t-t0)/max(t1-t0, 1e-6)
                    return (p0[0]+a*(p1[0]-p0[0]), p0[1]+a*(p1[1]-p0[1]))
            return rel[-1][1]

        # Build per-frame positions
        step = 1.0/fps
        raw = []
        t = 0.0
        while t <= dur+step:
            cx_r, cy_r = interp(t)
            x, y = clamp(int(cx_r*w), int(cy_r*h), cw, ch, w, h)
            raw.append((t, x, y))
            t += step

        # Smooth with small window + weight recent frames more to reduce lag
        win = max(1, int(fps * 0.12))  # 0.12s — fast response to movement
        smoothed = []
        for i, (t, x, y) in enumerate(raw):
            i0, i1 = max(0, i-win), min(len(raw), i+win+1)
            weights = [1.0 / (abs(j-i)+1) for j in range(i0, i1)]
            xs = [raw[j][1] for j in range(i0, i1)]
            ys = [raw[j][2] for j in range(i0, i1)]
            sw = sum(weights)
            sx = int(sum(wt*xv for wt,xv in zip(weights,xs)) / sw)
            sy = int(sum(wt*yv for wt,yv in zip(weights,ys)) / sw)
            smoothed.append((t, sx, sy))

        # Write sendcmd file
        sc = out_path.parent / f"sc_{out_path.stem}.txt"
        with open(sc, "w") as f:
            for i, (t, x, y) in enumerate(smoothed):
                f.write(f"{t:.4f} crop x {x}; crop y {y};\n")

        cmd = [
            "ffmpeg", "-y", "-ss", str(start), "-i", str(video_path),
            "-t", str(dur),
            "-vf", f"crop={cw}:{ch}:0:0,sendcmd=f={sc},scale=1080:1920",
            "-c:v", "libx264", "-preset", "fast", "-crf", "23",
            "-c:a", "aac", "-b:a", "128k", "-movflags", "+faststart",
            str(out_path)
        ]
        r = subprocess.run(cmd, capture_output=True)
        if sc.exists(): sc.unlink()
        if r.returncode == 0 and out_path.exists():
            return True

    # Static fallback
    vals = list(pts.values()) if pts else [(0.5, 0.42)]
    cx_r = float(np.median([v[0] for v in vals]))
    cy_r = float(np.median([v[1] for v in vals]))
    x, y = clamp(int(cx_r*w), int(cy_r*h), cw, ch, w, h)
    cmd = [
        "ffmpeg", "-y", "-ss", str(start), "-i", str(video_path),
        "-t", str(dur),
        "-vf", f"crop={cw}:{ch}:{x}:{y},scale=1080:1920",
        "-c:v", "libx264", "-preset", "fast", "-crf", "23",
        "-c:a", "aac", "-b:a", "128k", "-movflags", "+faststart",
        str(out_path)
    ]
    r = subprocess.run(cmd, capture_output=True)
    return r.returncode == 0 and out_path.exists()

def merge_clips(clips, out_path) -> bool:
    if not clips: return False
    if len(clips) == 1:
        shutil.copy(clips[0], out_path)
        return True
    lst = out_path.parent / "concat.txt"
    lst.write_text("\n".join(f"file '{c}'" for c in clips))
    r = subprocess.run(["ffmpeg", "-y", "-f", "concat", "-safe", "0",
                         "-i", str(lst), "-c", "copy", str(out_path)],
                        capture_output=True)
    return r.returncode == 0

# ── pipeline ───────────────────────────────────────────────────────────────────

async def process_video(job_id: str, video_path: Path, ref_path: Path):
    job = jobs[job_id]
    try:
        job["status"] = "processing"
        job["status_text"] = "Analyzing reference photo..."
        job["progress"] = 5

        ref_emb = extract_ref_embedding(ref_path)
        if ref_emb is None:
            job["status"] = "failed"
            job["error"] = "No face detected in reference photo. Use a clear front-facing photo."
            return

        job["status_text"] = "Reading video..."
        job["progress"] = 10
        vi = video_info(video_path)
        fps, duration = vi["fps"], vi["duration"]

        loop = asyncio.get_event_loop()
        timestamps, positions = await loop.run_in_executor(
            None, lambda: scan_video(video_path, ref_emb, fps, job_id))

        if not timestamps or jobs[job_id].get("cancelled"):
            if jobs[job_id].get("cancelled"):
                job["status"] = "cancelled"
                job["status_text"] = "Cancelled."
                return
            job["status"] = "failed"
            job["error"] = "Could not find the reference person in the video. Tips: use a clear front-facing photo, make sure the person appears clearly in the video, and avoid blurry or dark photos."
            return

        intervals = to_intervals(timestamps, fps)
        if not intervals:
            job["status"] = "failed"
            job["error"] = "Person found but clips too short to extract."
            return

        out_dir = video_path.parent
        clips = []
        total_intervals = len(intervals)

        for i, (s, e) in enumerate(intervals):
            job["status_text"] = f"Cutting clip {i+1} of {total_intervals}..."
            job["progress"] = 83 + int((i / max(total_intervals, 1)) * 10)
            op = out_dir / f"clip_{i:04d}.mp4"
            if cut_clip(video_path, s, e, positions, vi, op):
                clips.append(op)

        if not clips:
            job["status"] = "failed"
            job["error"] = "Failed to cut clips."
            return

        job["status_text"] = f"Merging {len(clips)} clip(s) into final video..."
        job["progress"] = 94

        final = out_dir / "result.mp4"
        if not merge_clips(clips, final):
            job["status"] = "failed"
            job["error"] = "Failed to merge clips."
            return

        job["status_text"] = "Finalizing video..."
        job["progress"] = 98

        total_kept = sum(e-s for s,e in intervals)
        job.update({
            "status": "done", "progress": 100, "status_text": "Done!",
            "result_path": str(final),
            "stats": {"duration": round(duration,1),
                      "extracted": round(total_kept,1),
                      "clips": len(intervals)}
        })

    except Exception as e:
        import traceback
        job["status"] = "failed"
        job["error"] = f"Error: {str(e)}"
        print(traceback.format_exc())

# ── routes ─────────────────────────────────────────────────────────────────────

@app.get("/health")
def health():
    return {"status": "ok", "face_recognition": True, "device": "cpu"}

@app.post("/upload-chunk")
async def upload_chunk(job_id: str, chunk_index: int, total_chunks: int,
                        file_type: str, chunk: UploadFile = File(...)):
    job_dir = TMP / job_id
    job_dir.mkdir(parents=True, exist_ok=True)
    chunk_dir = job_dir / f"{file_type}_chunks"
    chunk_dir.mkdir(exist_ok=True)
    (chunk_dir / f"{chunk_index:06d}").write_bytes(await chunk.read())

    received = len(list(chunk_dir.glob("*")))
    if received == total_chunks:
        ext = ".mp4" if file_type == "video" else ".jpg"
        final = job_dir / f"{file_type}{ext}"
        with open(final, "wb") as out:
            for i in range(total_chunks):
                out.write((chunk_dir / f"{i:06d}").read_bytes())
        shutil.rmtree(chunk_dir)
        return {"status": "complete"}
    return {"status": "partial", "received": received, "total": total_chunks}

@app.post("/process")
async def process(job_id: str, background_tasks: BackgroundTasks):
    job_dir = TMP / job_id
    video_path = next(
        (job_dir / f"video{ext}" for ext in [".mp4",".mov",".avi",".mkv",".webm"]
         if (job_dir / f"video{ext}").exists()), None)
    ref_path = job_dir / "reference.jpg"
    if not video_path:
        raise HTTPException(400, "Video not found.")
    if not ref_path.exists():
        raise HTTPException(400, "Reference photo not found.")
    jobs[job_id] = {"status": "queued", "progress": 0,
                    "status_text": "Queued...", "result_path": None,
                    "error": None, "stats": None, "cancelled": False, "eta": ""}
    background_tasks.add_task(process_video, job_id, video_path, ref_path)
    return {"job_id": job_id, "status": "queued"}

@app.get("/status/{job_id}")
def get_status(job_id: str):
    if job_id not in jobs:
        raise HTTPException(404, "Job not found")
    j = jobs[job_id]
    return {"status": j["status"], "progress": j["progress"],
            "status_text": j["status_text"], "error": j.get("error"),
            "stats": j.get("stats"), "eta": j.get("eta", "")}

@app.post("/cancel/{job_id}")
def cancel_job(job_id: str):
    if job_id not in jobs:
        raise HTTPException(404, "Job not found")
    jobs[job_id]["cancelled"] = True
    jobs[job_id]["status"] = "cancelled"
    jobs[job_id]["status_text"] = "Cancelling..."
    return {"cancelled": True}

@app.get("/preview/{job_id}")
def get_preview(job_id: str):
    preview_path = TMP / job_id / "preview.jpg"
    if not preview_path.exists():
        raise HTTPException(404, "No preview yet")
    return FileResponse(preview_path, media_type="image/jpeg")

@app.get("/download/{job_id}")
def download(job_id: str):
    if job_id not in jobs:
        raise HTTPException(404, "Job not found")
    j = jobs[job_id]
    if j["status"] != "done":
        raise HTTPException(400, "Not complete")
    p = Path(j["result_path"])
    if not p.exists():
        raise HTTPException(404, "File not found")
    return FileResponse(p, media_type="video/mp4", filename="clipcut_result.mp4")

@app.delete("/job/{job_id}")
def delete_job(job_id: str):
    d = TMP / job_id
    if d.exists(): shutil.rmtree(d)
    jobs.pop(job_id, None)
    return {"deleted": True}

# ── VidFetch ───────────────────────────────────────────────────────────────────

vf_jobs: dict[str, dict] = {}

REPLIT_API_URL = "https://042d7c4e-b637-40b4-b7d2-d68beee78b78-00-3mdbyed99dkeg.pike.replit.dev"

def run_vidfetch(job_id: str, urls: list, mode: str):
    import requests as req_lib
    import time
    job = vf_jobs[job_id]
    out_dir = TMP / job_id
    out_dir.mkdir(parents=True, exist_ok=True)
    downloaded = []

    job["status"] = "processing"
    job["status_text"] = "Starting download..."
    job["progress"] = 5

    try:
        total = len(urls)
        for i, url in enumerate(urls):
            if job.get("cancelled"): break

            job["status_text"] = f"Downloading video {i+1} of {total}..."
            job["progress"] = 10 + int((i / total) * 70)

            out_path = out_dir / f"video_{i:03d}.mp4"
            tmp_path = out_dir / f"tmp_{i:03d}.mp4"

            # Send to Replit API — uses Replit IP + mobile UA, works for Instagram
            try:
                replit_job = f"vf_{job_id}_{i}"
                print(f"[VidFetch] Sending to Replit: {url}")
                r = req_lib.post(
                    f"{REPLIT_API_URL}/download",
                    json={"url": url, "job_id": replit_job},
                    timeout=20
                )
                if r.status_code != 200:
                    raise Exception(f"Replit API error: {r.status_code}")

                for _ in range(100):
                    time.sleep(3)
                    sr = req_lib.get(f"{REPLIT_API_URL}/status/{replit_job}", timeout=10)
                    if sr.status_code != 200: continue
                    d = sr.json()
                    print(f"[VidFetch] Replit status: {d}")
                    if d.get("status") == "done": break
                    if d.get("status") == "failed":
                        raise Exception(d.get("error", "Download failed"))
                else:
                    raise Exception("Download timed out")

                fr = req_lib.get(f"{REPLIT_API_URL}/file/{replit_job}", timeout=300, stream=True)
                if fr.status_code != 200:
                    raise Exception(f"File fetch failed: {fr.status_code}")
                with open(tmp_path, "wb") as f:
                    for chunk in fr.iter_content(chunk_size=65536):
                        if chunk: f.write(chunk)
                try: req_lib.delete(f"{REPLIT_API_URL}/job/{replit_job}", timeout=10)
                except: pass

            except Exception as e:
                # Skip this video and continue with others
                print(f"[VidFetch] Skipping video {i+1}: {str(e)}")
                job["status_text"] = f"Skipped video {i+1} (restricted) — continuing..."
                continue

            if not tmp_path.exists() or tmp_path.stat().st_size == 0:
                job["status"] = "failed"
                job["error"] = f"Video {i+1} file is empty after download."
                return

            # Re-encode to Android-compatible mp4
            job["status_text"] = f"Converting video {i+1}..."
            fix_result = subprocess.run([
                "ffmpeg", "-y", "-i", str(tmp_path),
                "-c:v", "libx264", "-preset", "fast", "-crf", "23",
                "-c:a", "aac", "-b:a", "128k",
                "-movflags", "+faststart", "-pix_fmt", "yuv420p",
                str(out_path)
            ], capture_output=True)
            tmp_path.unlink(missing_ok=True)

            if fix_result.returncode == 0 and out_path.exists() and out_path.stat().st_size > 0:
                downloaded.append(out_path)
            else:
                job["status"] = "failed"
                job["error"] = f"Failed to convert video {i+1}."
                return

        if job.get("cancelled"):
            job["status"] = "cancelled"
            return

        if not downloaded:
            job["status"] = "failed"
            job["error"] = "All videos failed to download — they may be restricted or private. Try different links."
            return

        job["progress"] = 88

        if mode == "merge" and len(downloaded) > 1:
            job["status_text"] = "Merging videos..."
            job["progress"] = 92
            list_file = out_dir / "concat.txt"
            list_file.write_text("\n".join(f"file '{p}'" for p in downloaded))
            merged = out_dir / "merged.mp4"
            r = subprocess.run([
                "ffmpeg", "-y", "-f", "concat", "-safe", "0",
                "-i", str(list_file), "-c", "copy", str(merged)
            ], capture_output=True)
            if r.returncode != 0 or not merged.exists():
                job["status"] = "failed"
                job["error"] = "Failed to merge videos."
                return
            vi = video_info(merged)
            job.update({
                "status": "done", "progress": 100, "status_text": "Done!",
                "mode": "merge", "count": len(downloaded),
                "duration": round(vi["duration"], 1),
                "result_path": str(merged)
            })
        else:
            job.update({
                "status": "done", "progress": 100, "status_text": "Done!",
                "mode": "download", "count": len(downloaded),
                "result_paths": [str(p) for p in downloaded],
                "skipped": total - len(downloaded)
            })

    except Exception as e:
        import traceback
        job["status"] = "failed"
        job["error"] = f"Error: {str(e)}"
        print(traceback.format_exc())






from pydantic import BaseModel

class VidFetchRequest(BaseModel):
    job_id: str
    urls: list
    mode: str = "download"

@app.post("/vidfetch")
async def vidfetch(req: VidFetchRequest, background_tasks: BackgroundTasks):
    vf_jobs[req.job_id] = {
        "status": "queued", "progress": 0, "status_text": "Queued...",
        "cancelled": False, "error": None
    }
    background_tasks.add_task(run_vidfetch, req.job_id, req.urls, req.mode)
    return {"job_id": req.job_id, "status": "queued"}

@app.get("/vidfetch/status/{job_id}")
def vidfetch_status(job_id: str):
    if job_id not in vf_jobs:
        raise HTTPException(404, "Job not found")
    j = vf_jobs[job_id]
    return {
        "status": j["status"], "progress": j["progress"],
        "status_text": j["status_text"], "error": j.get("error"),
        "mode": j.get("mode"), "count": j.get("count"),
        "duration": j.get("duration"),
        "failed_videos": j.get("failed_videos", [])
    }

@app.post("/vidfetch/cancel/{job_id}")
def vidfetch_cancel(job_id: str):
    if job_id not in vf_jobs:
        raise HTTPException(404, "Job not found")
    vf_jobs[job_id]["cancelled"] = True
    vf_jobs[job_id]["status"] = "cancelled"
    return {"cancelled": True}

@app.get("/vidfetch/download/{job_id}")
def vidfetch_download(job_id: str):
    if job_id not in vf_jobs:
        raise HTTPException(404, "Job not found")
    j = vf_jobs[job_id]
    if j["status"] != "done": raise HTTPException(400, "Not complete")
    path = Path(j.get("result_path") or j.get("result_paths", [""])[0])
    if not path.exists(): raise HTTPException(404, "File not found")
    return FileResponse(path, media_type="video/mp4", filename="vidix_download.mp4")

@app.get("/vidfetch/download/{job_id}/{index}")
def vidfetch_download_single(job_id: str, index: int):
    if job_id not in vf_jobs:
        raise HTTPException(404, "Job not found")
    j = vf_jobs[job_id]
    if j["status"] != "done": raise HTTPException(400, "Not complete")
    paths = j.get("result_paths", [])
    if index >= len(paths): raise HTTPException(404, "File not found")
    path = Path(paths[index])
    if not path.exists(): raise HTTPException(404, "File not found")
    return FileResponse(path, media_type="video/mp4", filename=f"vidix_video_{index+1}.mp4")

@app.delete("/vidfetch/job/{job_id}")
def vidfetch_delete(job_id: str):
    d = TMP / job_id
    if d.exists(): shutil.rmtree(d)
    vf_jobs.pop(job_id, None)
    return {"deleted": True}

# ── FaceSwap ───────────────────────────────────────────────────────────────────

fs_jobs: dict[str, dict] = {}

def run_faceswap(job_id: str, source_path: Path, target_path: Path):
    """Swap face from source into target image."""
    import cv2
    job = fs_jobs[job_id]
    job["status"] = "processing"
    job["status_text"] = "Loading models..."
    job["progress"] = 10

    try:
        job["status_text"] = "Loading AI model (first run may take 1-2 min)..."
        job["progress"] = 5
        fa = get_face_app()
        swapper = get_swapper()

        # Load images
        source_img = cv2.imread(str(source_path))
        target_img = cv2.imread(str(target_path))

        if source_img is None:
            job["status"] = "failed"; job["error"] = "Could not read source photo."; return
        if target_img is None:
            job["status"] = "failed"; job["error"] = "Could not read target photo."; return

        job["status_text"] = "Detecting faces..."
        job["progress"] = 30

        # Get source face
        source_faces = fa.get(source_img)
        if not source_faces:
            job["status"] = "failed"; job["error"] = "No face detected in source photo. Use a clear front-facing photo."; return

        # Get target faces
        target_faces = fa.get(target_img)
        if not target_faces:
            job["status"] = "failed"; job["error"] = "No face detected in target photo. Use a clear front-facing photo."; return

        job["status_text"] = "Swapping face..."
        job["progress"] = 60

        # Use largest face from source
        source_face = max(source_faces, key=lambda f: (f.bbox[2]-f.bbox[0]) * (f.bbox[3]-f.bbox[1]))

        # Swap all faces in target
        result = target_img.copy()
        for face in target_faces:
            result = swapper.get(result, face, source_face, paste_back=True)

        job["status_text"] = "Saving result..."
        job["progress"] = 90

        out_path = source_path.parent / "result.jpg"
        cv2.imwrite(str(out_path), result, [cv2.IMWRITE_JPEG_QUALITY, 95])

        job.update({
            "status": "done", "progress": 100,
            "status_text": "Done!",
            "result_path": str(out_path)
        })

    except Exception as e:
        import traceback
        job["status"] = "failed"
        job["error"] = f"Error: {str(e)}"
        print(traceback.format_exc())


@app.post("/faceswap")
async def faceswap(
    background_tasks: BackgroundTasks,
    job_id: str,
    source: UploadFile = File(...),
    target: UploadFile = File(...)
):
    job_dir = TMP / job_id
    job_dir.mkdir(parents=True, exist_ok=True)

    source_path = job_dir / "source.jpg"
    target_path = job_dir / "target.jpg"

    source_path.write_bytes(await source.read())
    target_path.write_bytes(await target.read())

    fs_jobs[job_id] = {
        "status": "queued", "progress": 0,
        "status_text": "Queued...", "error": None,
        "result_path": None
    }
    background_tasks.add_task(run_faceswap, job_id, source_path, target_path)
    return {"job_id": job_id, "status": "queued"}

@app.get("/faceswap/status/{job_id}")
def faceswap_status(job_id: str):
    if job_id not in fs_jobs:
        raise HTTPException(404, "Job not found")
    j = fs_jobs[job_id]
    return {"status": j["status"], "progress": j["progress"],
            "status_text": j["status_text"], "error": j.get("error")}

@app.get("/faceswap/download/{job_id}")
def faceswap_download(job_id: str):
    if job_id not in fs_jobs:
        raise HTTPException(404, "Job not found")
    j = fs_jobs[job_id]
    if j["status"] != "done": raise HTTPException(400, "Not complete")
    path = Path(j["result_path"])
    if not path.exists(): raise HTTPException(404, "File not found")
    return FileResponse(path, media_type="image/jpeg", filename="vidix_faceswap.jpg")

@app.delete("/faceswap/job/{job_id}")
def faceswap_delete(job_id: str):
    d = TMP / job_id
    if d.exists(): shutil.rmtree(d)
    fs_jobs.pop(job_id, None)
    return {"deleted": True}

# ── VidPlay ────────────────────────────────────────────────────────────────────

vp_jobs: dict[str, dict] = {}

def run_vidplay(job_id: str, url: str):
    """Download a single video for VidPlay streaming."""
    import requests as req_lib
    import time
    job = vp_jobs[job_id]
    out_dir = TMP / job_id
    out_dir.mkdir(parents=True, exist_ok=True)

    job["status"] = "processing"
    job["status_text"] = "Downloading video..."
    job["progress"] = 10

    try:
        tmp_path = out_dir / "tmp.mp4"
        out_path = out_dir / "video.mp4"
        replit_job = f"vp_{job_id}"

        # Send to Replit
        r = req_lib.post(
            f"{REPLIT_API_URL}/download",
            json={"url": url, "job_id": replit_job},
            timeout=20
        )
        if r.status_code != 200:
            raise Exception(f"Download service error: {r.status_code}")

        # Poll
        for _ in range(100):
            time.sleep(3)
            sr = req_lib.get(f"{REPLIT_API_URL}/status/{replit_job}", timeout=10)
            if sr.status_code != 200: continue
            d = sr.json()
            if d.get("status") == "done": break
            if d.get("status") == "failed":
                raise Exception(d.get("error", "Download failed"))
        else:
            raise Exception("Download timed out")

        job["status_text"] = "Processing video..."
        job["progress"] = 80

        # Fetch file
        fr = req_lib.get(f"{REPLIT_API_URL}/file/{replit_job}", timeout=300, stream=True)
        if fr.status_code != 200:
            raise Exception(f"File fetch failed: {fr.status_code}")
        with open(tmp_path, "wb") as f:
            for chunk in fr.iter_content(chunk_size=65536):
                if chunk: f.write(chunk)

        # Cleanup Replit
        try: req_lib.delete(f"{REPLIT_API_URL}/job/{replit_job}", timeout=10)
        except: pass

        # Re-encode for browser compatibility
        fix = subprocess.run([
            "ffmpeg", "-y", "-i", str(tmp_path),
            "-c:v", "libx264", "-preset", "fast", "-crf", "23",
            "-c:a", "aac", "-b:a", "128k",
            "-movflags", "+faststart",
            "-pix_fmt", "yuv420p",
            str(out_path)
        ], capture_output=True)
        tmp_path.unlink(missing_ok=True)

        if not out_path.exists() or out_path.stat().st_size == 0:
            raise Exception("Video processing failed")

        job.update({
            "status": "done", "progress": 100,
            "status_text": "Ready!",
            "result_path": str(out_path)
        })

    except Exception as e:
        import traceback
        job["status"] = "failed"
        job["error"] = str(e)
        print(traceback.format_exc())


@app.post("/vidplay/fetch")
async def vidplay_fetch(background_tasks: BackgroundTasks, url: str, job_id: str):
    vp_jobs[job_id] = {
        "status": "queued", "progress": 0,
        "status_text": "Queued...", "error": None,
        "result_path": None
    }
    background_tasks.add_task(run_vidplay, job_id, url)
    return {"job_id": job_id, "status": "queued"}

@app.get("/vidplay/status/{job_id}")
def vidplay_status(job_id: str):
    if job_id not in vp_jobs:
        raise HTTPException(404, "Job not found")
    j = vp_jobs[job_id]
    return {
        "status": j["status"], "progress": j["progress"],
        "status_text": j["status_text"], "error": j.get("error")
    }

@app.get("/vidplay/stream/{job_id}")
def vidplay_stream(job_id: str):
    """Stream the video file — used directly as video src."""
    if job_id not in vp_jobs:
        raise HTTPException(404, "Job not found")
    j = vp_jobs[job_id]
    if j["status"] != "done": raise HTTPException(400, "Not ready")
    path = Path(j["result_path"])
    if not path.exists(): raise HTTPException(404, "File not found")
    return FileResponse(
        path, media_type="video/mp4",
        headers={"Accept-Ranges": "bytes"}
    )

@app.delete("/vidplay/job/{job_id}")
def vidplay_delete(job_id: str):
    d = TMP / job_id
    if d.exists(): shutil.rmtree(d)
    vp_jobs.pop(job_id, None)
    return {"deleted": True}

# ── Profile Scanner ────────────────────────────────────────────────────────────

class ProfileScanRequest(BaseModel):
    url: str

@app.post("/profile/scan")
async def profile_scan(req: ProfileScanRequest):
    """
    Use yt-dlp --flat-playlist to get post list from a profile URL.
    Returns list of {url, title, type, date} — no thumbnails, minimal requests.
    """
    import json as _json

    url = req.url.strip()
    if not url:
        raise HTTPException(400, "No URL provided")

    try:
        result = subprocess.run([
            "yt-dlp",
            "--flat-playlist",
            "-J",                          # dump JSON, don't download
            "--no-warnings",
            "--user-agent", "Mozilla/5.0 (iPhone; CPU iPhone OS 17_0 like Mac OS X) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/17.0 Mobile/15E148 Safari/604.1",
            "--playlist-end", "50",        # cap at 50 posts to stay fast
            url
        ], capture_output=True, text=True, timeout=60)

        if result.returncode != 0:
            err = result.stderr.strip().splitlines()
            msg = err[-1] if err else "Could not scan profile"
            raise HTTPException(400, msg)

        raw = _json.loads(result.stdout)

        # Could be a playlist or a single entry
        entries = raw.get("entries") or []
        if not entries and raw.get("url"):
            # Single video returned
            entries = [raw]

        posts = []
        for entry in entries:
            if entry is None:
                continue
            post_url = (
                entry.get("webpage_url") or
                entry.get("url") or
                entry.get("original_url")
            )
            if not post_url:
                continue

            # Normalise date from YYYYMMDD → YYYY-MM-DD
            raw_date = entry.get("upload_date") or ""
            if len(raw_date) == 8:
                date_str = f"{raw_date[:4]}-{raw_date[4:6]}-{raw_date[6:]}"
            else:
                date_str = raw_date

            # Determine type
            vtype = "video"
            if entry.get("_type") == "url" and not entry.get("duration"):
                vtype = "photo"

            posts.append({
                "url":   post_url,
                "title": entry.get("title") or entry.get("id") or "",
                "type":  vtype,
                "date":  date_str
            })

        if not posts:
            raise HTTPException(404, "No posts found — profile may be private or unsupported")

        return {"count": len(posts), "posts": posts}

    except HTTPException:
        raise
    except subprocess.TimeoutExpired:
        raise HTTPException(504, "Scan timed out — profile may be too large or slow to respond")
    except Exception as e:
        import traceback
        print(traceback.format_exc())
        raise HTTPException(500, f"Scan failed: {str(e)}")