import cv2
import tempfile
import requests
import os
from PIL import Image
from transformers import pipeline

classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection")


def download_file(url):
    response = requests.get(url, stream=True)
    tmp = tempfile.NamedTemporaryFile(delete=False)
    for chunk in response.iter_content(1024):
        tmp.write(chunk)
    tmp.close()
    return tmp.name


def get_video_duration(video_path):
    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS)
    frame_count = cap.get(cv2.CAP_PROP_FRAME_COUNT)
    cap.release()
    return frame_count / fps if fps > 0 else 0


def extract_frame(video_path, second):
    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS)
    frame_number = int(fps * second)

    cap.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
    success, frame = cap.read()
    cap.release()

    if not success:
        return None

    tmp_file = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
    cv2.imwrite(tmp_file.name, frame)
    return tmp_file.name


def get_frame_times(duration, file_size_mb):
    if duration <= 2:
        return [1, 2]

    elif duration <= 10:
        return [2]

    elif duration <= 15:
        return [4, 9, 13]

    if file_size_mb > 14:
        return [4, 9, 13]

    return [2]


def check_image_nsfw(image_path):
    image = Image.open(image_path).convert("RGB")
    result = classifier(image)

    for r in result:
        if r["label"] == "nsfw" and r["score"] > 0.5:
            return True

    return False


def check_video_nsfw(video_path):
    size_mb = os.path.getsize(video_path) / (1024 * 1024)
    duration = get_video_duration(video_path)

    times = get_frame_times(duration, size_mb)

    for t in times:
        frame = extract_frame(video_path, t)
        if frame:
            if check_image_nsfw(frame):
                return True  # 🚨 return immediately if ANY frame is NSFW

    return False