# visualizer.py
import cv2
import os
import uuid
from utils import save_video

OUTPUT_DIR = "output_videos"
os.makedirs(OUTPUT_DIR, exist_ok=True)

def draw_visuals(frames, ball_positions, trajectory, impact_frame_idx, decision):
    font = cv2.FONT_HERSHEY_SIMPLEX
    output_frames = []

    for i, frame in enumerate(frames):
        frame_draw = frame.copy()

        # Draw past ball positions
        for _, x, y in ball_positions:
            cv2.circle(frame_draw, (x, y), 5, (0, 255, 255), -1)

        # Draw trajectory after impact
        if i == impact_frame_idx:
            for x, y in trajectory:
                cv2.circle(frame_draw, (x, y), 4, (0, 0, 255), -1)

        # Highlight impact frame
        if i == impact_frame_idx:
            cv2.putText(frame_draw, "IMPACT!", (50, 50), font, 1.0, (0, 0, 255), 2)
            cv2.rectangle(frame_draw, (10, 10), (630, 470), (0, 0, 255), 4)

        # Final decision on last frame
        if i == len(frames) - 1:
            cv2.putText(frame_draw, f"Decision: {decision}", (100, 420), font, 1.2, (255, 0, 0), 3)

        output_frames.append(frame_draw)

    # Save annotated video
    video_id = str(uuid.uuid4())
    out_path = os.path.join(OUTPUT_DIR, f"{video_id}_annotated.mp4")
    save_video(output_frames, out_path)
    return out_path