#!/usr/bin/env python3
"""
LTX MSR auto-trim: detects and removes static conditioning frames.
"""

import argparse
import subprocess
import sys
import os
import re
import json


def extract_frame_info(video_path):
    """Extract per-frame mean RGB using ffmpeg showinfo."""
    cmd = [
        "ffmpeg", "-i", video_path,
        "-vf", "select='gt(scene,0)',showinfo",
        "-an", "-f", "null", "-"
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    
    frames = []
    for line in result.stderr.split("\n"):
        if "Parsed_showinfo" not in line or "mean:" not in line:
            continue
        
        # Extract frame number: " n:   123 "
        m_n = re.search(r'\bn:\s*(\d+)', line)
        # Extract mean RGB: "mean:[R G B]"
        m_mean = re.search(r'mean:\[([\d\s]+)\]', line)
        
        if m_n and m_mean:
            frame_n = int(m_n.group(1))
            rgb_str = m_mean.group(1).strip()
            rgb = [float(x) for x in rgb_str.split()]
            frames.append({"frame": frame_n, "rgb": rgb})
    
    # Compute consecutive diffs
    diffs = []
    for i in range(1, len(frames)):
        prev_rgb = frames[i-1]["rgb"]
        curr_rgb = frames[i]["rgb"]
        diff = sum(abs(c - p) for c, p in zip(curr_rgb, prev_rgb)) / 3.0
        diffs.append({
            "from": frames[i-1]["frame"],
            "to": frames[i]["frame"],
            "diff": round(diff, 2)
        })
    
    return frames, diffs


def detect_motion_onset(diffs, threshold=3.0, sustained=3, safety=3):
    """Find first sustained motion frame."""
    motion_streak = 0
    motion_start = None
    
    for i, d in enumerate(diffs):
        if d["diff"] > threshold:
            motion_streak += 1
            if motion_streak >= sustained and motion_start is None:
                # Backtrack to first frame of this streak
                start_idx = max(0, i - sustained + 1)
                motion_start = diffs[start_idx]["from"]
        else:
            motion_streak = 0
    
    if motion_start is None:
        return 0
    
    return motion_start + safety


def trim_video(input_path, output_path, trim_frames, fps=24):
    if trim_frames <= 0:
        subprocess.run(["ffmpeg", "-y", "-i", input_path, "-c", "copy", output_path],
                      capture_output=True)
        return
    
    trim_sec = trim_frames / fps
    cmd = [
        "ffmpeg", "-y",
        "-i", input_path,
        "-vf", f"select=gt(n\\,{trim_frames - 1}),setpts=PTS-STARTPTS",
        "-af", f"atrim=start={trim_sec:.4f},asetpts=PTS-STARTPTS",
        "-c:v", "libx264", "-crf", "18", "-preset", "fast",
        "-c:a", "aac", "-b:a", "128k",
        output_path
    ]
    result = subprocess.run(cmd, capture_output=True, text=True)
    if result.returncode != 0:
        print(f"ffmpeg error: {result.stderr[-500:]}", file=sys.stderr)
        sys.exit(1)


def main():
    parser = argparse.ArgumentParser(description="Auto-trim LTX MSR static conditioning frames")
    parser.add_argument("input", help="Input video path")
    parser.add_argument("output", help="Output video path")
    parser.add_argument("--threshold", type=float, default=3.0,
                       help="Motion detection threshold (default: 3.0)")
    parser.add_argument("--safety", type=int, default=3,
                       help="Extra frames to trim as safety margin (default: 3)")
    parser.add_argument("--fps", type=int, default=24)
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--json", action="store_true")
    parser.add_argument("--verbose", action="store_true")
    args = parser.parse_args()
    
    frames, diffs = extract_frame_info(args.input)
    
    if not diffs:
        print("ERROR: Could not extract frame data", file=sys.stderr)
        sys.exit(1)
    
    if args.verbose:
        print(f"Total frames: {len(frames)}")
        print(f"{'Frame':>6} {'Diff':>8} {'Verdict':>10}")
        print("-" * 28)
        for i, d in enumerate(diffs[:30]):
            v = "STATIC" if d["diff"] < args.threshold else "MOTION"
            print(f"{d['to']:>6} {d['diff']:>8.2f} {v:>10}")
        if len(diffs) > 30:
            print(f"... ({len(diffs) - 30} more)")
    
    motion_onset = detect_motion_onset(diffs, threshold=args.threshold, safety=0)
    trim_count = motion_onset + args.safety
    
    total_frames = max(f["frame"] for f in frames) + 1
    remaining = total_frames - trim_count
    
    if not args.json:
        print(f"\nTotal frames: {total_frames}")
        print(f"Motion onset: frame {motion_onset}")
        print(f"Safety margin: +{args.safety} frames")
        print(f"Trim: {trim_count} frames ({trim_count/args.fps:.2f}s)")
        print(f"Output: {remaining} frames ({remaining/args.fps:.2f}s)")
    
    if args.json:
        print(json.dumps({
            "total_frames": total_frames,
            "motion_onset": motion_onset,
            "trim_frames": trim_count,
            "remaining_frames": remaining,
            "remaining_seconds": round(remaining / args.fps, 2),
        }, indent=2))
    
    if args.dry_run:
        return
    
    trim_video(args.input, args.output, trim_count, args.fps)
    print(f"Output: {args.output}")


if __name__ == "__main__":
    main()