import cv2
import numpy as np

def create_comparison_collage(video_a_path, video_b_path, time_a, offset_b, radius=1, output_name="collage.jpg"):
    caps = {
        "A": cv2.VideoCapture(video_a_path),
        "B": cv2.VideoCapture(video_b_path)
    }

    results = {"A": [], "B": []}
    
    for label, cap in caps.items():
        if not cap.isOpened():
            print(f"Error: Could not open Video {label}")
            return

        fps = cap.get(cv2.CAP_PROP_FPS)
        # Determine base time for this specific video
        base_time = time_a if label == "A" else (time_a + offset_b)
        center_frame = int(round(base_time * fps))

        # Extract range: from (center - radius) to (center + radius)
        for i in range(-radius, radius + 1):
            target_idx = center_frame + i
            cap.set(cv2.CAP_PROP_POS_FRAMES, target_idx)
            success, frame = cap.read()

            if success:
                # Add overlay text (Timestamp and Frame Index)
                current_time = target_idx / fps
                text = f"Vid {label} | F: {target_idx} | {current_time:.4f}s"
                
                # Draw black background for text for readability
                cv2.putText(frame, text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0,0,0), 4, cv2.LINE_AA)
                cv2.putText(frame, text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255,255,255), 2, cv2.LINE_AA)
                
                results[label].append(frame)
            else:
                # Placeholder if frame is missing
                blank = np.zeros((720, 1280, 3), dtype=np.uint8)
                cv2.putText(blank, "FRAME MISSING", (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0,0,255), 2)
                results[label].append(blank)

        cap.release()

    # Ensure all frames are the same size for the collage (using Video A's dimensions)
    h, w, _ = results["A"][0].shape
    
    rows = []
    for i in range(len(results["A"])):
        img_a = cv2.resize(results["A"][i], (w, h))
        img_b = cv2.resize(results["B"][i], (w, h))
        # Combine Left (A) and Right (B)
        combined_row = np.hstack((img_a, img_b))
        rows.append(combined_row)

    # Combine all rows Top to Bottom
    final_collage = np.vstack(rows)

    # Save and show
    cv2.imwrite(output_name, final_collage)
    print(f"Collage saved as {output_name}")

# --- SETTINGS ---
video_a = "/home/honney/projects/mkv/Anohana - The Flower We Saw That Day - S01 (BD 1080p HEVC Opus) [derpie]/[derpie] Anohana - The Flower We Saw That Day - S01E02 (BD 1080p HEVC Opus) [Dual Audio].mkv"
video_b = "/home/honney/projects/mkv/AnoHana The Flower We Saw That Day (2011-2011) [imdbid-tt1913273]/Season 01/AnoHana The Flower We Saw That Day S01E002.mkv"

timestamp_a = 60 * 20#1/(24000/1001) * 250  # Center point for Video A 75
offset_for_b = 1/(24000/1001) * 24
frames_around = 10                  # 2 before, 2 after (5 frames total per video)

create_comparison_collage(video_a, video_b, timestamp_a, offset_for_b, frames_around)