import cv2
import torch
import requests
import time
from PIL import Image
from supabase import create_client
from transformers import AutoImageProcessor, AutoModelForImageClassification
import os

# ===============================
# 🗄 Supabase Setup
# ===============================
SUPABASE_URL = os.getenv("SUPABASE_URL")
SUPABASE_KEY = os.getenv("SUPABASE_KEY")

supabase = create_client(SUPABASE_URL, SUPABASE_KEY)

# ===============================
# 🤖 Load NSFW Model (once)
# ===============================
print("🔄 Loading NSFW model...")

model_name = "AdamCodd/vit-base-nsfw-detector"
processor = AutoImageProcessor.from_pretrained(model_name)
model = AutoModelForImageClassification.from_pretrained(model_name)

print("✅ Model loaded")

# ===============================
# 🖼 Image Check
# ===============================
def check_image(url):
    try:
        image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
        inputs = processor(images=image, return_tensors="pt")

        with torch.no_grad():
            outputs = model(**inputs)

        probs = torch.softmax(outputs.logits, dim=1)
        return "explicit" if probs[0][1] > 0.5 else "safe"

    except Exception as e:
        print("❌ Image error:", e)
        return "safe"

# ===============================
# 🎥 Video Check
# ===============================
def check_video(url, frame_sample_rate=30):
    try:
        cap = cv2.VideoCapture(url)
        frame_count = 0

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

            if frame_count % frame_sample_rate == 0:
                img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
                inputs = processor(images=img, return_tensors="pt")

                with torch.no_grad():
                    outputs = model(**inputs)

                probs = torch.softmax(outputs.logits, dim=1)

                if probs[0][1] > 0.5:
                    cap.release()
                    return "explicit"

            frame_count += 1

        cap.release()
        return "safe"

    except Exception as e:
        print("❌ Video error:", e)
        return "safe"

# ===============================
# 🔁 MAIN WORKER LOOP
# ===============================
def moderation_worker():
    while True:
        try:
            print("\n🚀 Running moderation job...")

            # ===============================
            # 🖼 Process Images
            # ===============================
            img_posts = supabase.table("posts") \
                .select("*") \
                .limit(10) \
                .execute().data

            for post in img_posts:
                existing = supabase.table("content_moderation") \
                    .select("id") \
                    .eq("post_id", post["id"]) \
                    .eq("post_type", "image") \
                    .execute()

                if len(existing.data) == 0:
                    result = check_image(post["image_url"])

                    supabase.table("content_moderation").insert({
                        "post_id": post["id"],
                        "post_type": "image",
                        "file_url": post["image_url"],
                        "result": result
                    }).execute()

                    print(f"🖼 IMAGE {post['id']} → {result}")

            # ===============================
            # 🎥 Process Videos
            # ===============================
            vid_posts = supabase.table("trendz") \
                .select("*") \
                .limit(5) \
                .execute().data

            for post in vid_posts:
                existing = supabase.table("content_moderation") \
                    .select("id") \
                    .eq("post_id", post["id"]) \
                    .eq("post_type", "video") \
                    .execute()

                if len(existing.data) == 0:
                    result = check_video(post["video_url"])

                    supabase.table("content_moderation").insert({
                        "post_id": post["id"],
                        "post_type": "video",
                        "file_url": post["video_url"],
                        "result": result
                    }).execute()

                    print(f"🎥 VIDEO {post['id']} → {result}")

        except Exception as e:
            print("❌ Worker error:", e)

        # ⏱ WAIT 5 MINUTES
        print("⏳ Sleeping for 5 minutes...\n")
        time.sleep(300)


# ===============================
# ▶️ AUTO START
# ===============================
if __name__ == "__main__":
    print("🔥 Starting NSFW moderation worker...")
    moderation_worker()