import cv2
import torch
import os
import warnings
from typing import Dict
from dataclasses import dataclass

warnings.filterwarnings('ignore')

try:
    from ultralytics import YOLO
    from transformers import pipeline
    from PIL import Image
except ImportError as e:
    print(f"Missing dependency: {e}")


@dataclass
class DetectionResult:
    """Simple detection result"""
    nude_count: int = 0
    gun_count: int = 0
    knife_count: int = 0
    fight_count: int = 0
    is_safe: bool = True

    def to_dict(self):
        return {
            'nude': self.nude_count,
            'gun': self.gun_count,
            'knife': self.knife_count,
            'fight': self.fight_count,
            'is_safe': self.is_safe
        }


class SmartSequentialModerator:
    """
    Smart Sequential Pipeline with balanced thresholds:
    1. NSFW Check with BALANCED threshold
    2. Only if NSFW is clean → Check Weapons/Fights
    """

    def __init__(self):
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'

        # Models
        self.nsfw_classifier = None
        self.weapon_model = None

        # BALANCED Thresholds
        self.nsfw_threshold = 0.75  # Balanced: not too high, not too low
        self.nsfw_safe_threshold = 0.25  # If below this, definitely safe
        self.gun_threshold = 0.7
        self.knife_threshold = 0.65
        self.fight_threshold = 0.75

        print(f"🚀 Smart Sequential Moderator initialized on {self.device}")
        print(f"📋 Pipeline: NSFW (0.75) → Weapons/Fights")

        self._setup_models()

    def _setup_models(self):
        """Initialize models"""
        try:
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

            # 1. NSFW Classifier (PRIORITY)
            self._setup_nsfw()

            # 2. Weapon/Fight Model
            self._setup_weapons()

            print("✅ All models ready!")

        except Exception as e:
            print(f"❌ Setup error: {e}")

    def _setup_nsfw(self):
        """Setup NSFW classifier"""
        try:
            print("🔞 Loading NSFW classifier...")

            device_id = 0 if self.device == 'cuda' else -1

            # Use the NSFW detection model
            self.nsfw_classifier = pipeline(
                "image-classification",
                model="Falconsai/nsfw_image_detection",
                device=device_id
            )
            print("✅ NSFW classifier loaded")

        except Exception as e:
            print(f"⚠️ NSFW failed: {e}")
            self.nsfw_classifier = None

    def _setup_weapons(self):
        """Setup weapon/fight model"""
        try:
            print("🔫 Loading weapon/fight model...")

            # Custom model path
            custom_path = "models/best_ft4.pt"
            if os.path.exists(custom_path):
                self.weapon_model = YOLO(custom_path)
                print(f"✅ Custom model loaded")

                # Show available classes
                if hasattr(self.weapon_model, 'names'):
                    classes = list(self.weapon_model.names.values())
                    print(f"   Classes: {classes}")
            else:
                # Fallback
                self.weapon_model = YOLO('yolo11n.pt')
                print("✅ General model loaded")

        except Exception as e:
            print(f"⚠️ Weapon model failed: {e}")
            self.weapon_model = None

    def process_image(self, image) -> DetectionResult:
        """
        STRICT SEQUENTIAL:
        1. NSFW first (balanced threshold)
        2. If NSFW detected → STOP
        3. If clean → check weapons/fights
        """

        result = DetectionResult()

        try:
            # Load image
            if isinstance(image, str):
                image = cv2.imread(image)
                if image is None:
                    return result

            print(f"\n{'=' * 40}")
            print(f"📸 Processing: {image.shape}")

            # ========== STAGE 1: NSFW ==========
            print("\n🔞 Stage 1: NSFW Check")

            nsfw_score = self._check_nsfw(image)

            if nsfw_score > self.nsfw_threshold:
                print(f"   🚨 NSFW DETECTED: {nsfw_score:.3f}")
                print(f"   ⛔ STOPPING - Returning NSFW only")

                result.nude_count = 1
                result.is_safe = False
                return result  # STOP HERE

            elif nsfw_score < self.nsfw_safe_threshold:
                print(f"   ✅ Definitely safe: {nsfw_score:.3f}")
            else:
                print(f"   ⚠️ Borderline safe: {nsfw_score:.3f} - Continuing checks")

            # ========== STAGE 2: WEAPONS/FIGHTS ==========
            print("\n🔫 Stage 2: Weapons & Fights")

            if self.weapon_model:
                detections = self._detect_threats(image)
                result.gun_count = detections['guns']
                result.knife_count = detections['knives']
                result.fight_count = detections['fights']

                if detections['total'] > 0:
                    print(f"   Found: G:{detections['guns']} K:{detections['knives']} F:{detections['fights']}")

            # Final safety
            total = result.nude_count + result.gun_count + result.knife_count + result.fight_count
            result.is_safe = (total == 0)

            print(
                f"\n📊 Result: N:{result.nude_count} G:{result.gun_count} K:{result.knife_count} F:{result.fight_count} Safe:{result.is_safe}")
            print(f"{'=' * 40}\n")

            return result

        except Exception as e:
            print(f"❌ Error: {e}")
            return result

    def _check_nsfw(self, image) -> float:
        """
        Check NSFW with proper scoring
        Returns confidence score (0-1)
        """
        try:
            if not self.nsfw_classifier:
                return 0.0

            # Convert to RGB
            rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            pil_image = Image.fromarray(rgb_image)

            # Run classifier
            results = self.nsfw_classifier(pil_image)

            # Get NSFW score
            nsfw_score = 0.0
            for result in results:
                label = result['label'].lower()
                score = result['score']

                # Check for NSFW label
                if 'nsfw' in label or 'unsafe' in label or 'explicit' in label:
                    nsfw_score = max(nsfw_score, score)
                    print(f"   {label}: {score:.3f}")

            return nsfw_score

        except Exception as e:
            print(f"   ⚠️ NSFW error: {e}")
            return 0.0

    def _detect_threats(self, image) -> Dict[str, int]:
        """Detect weapons and fights"""
        counts = {
            'guns': 0,
            'knives': 0,
            'fights': 0,
            'total': 0
        }

        try:
            # Run detection with low base threshold
            results = self.weapon_model(
                image,
                conf=0.4,  # Low base threshold
                device=self.device,
                verbose=False
            )

            for result in results:
                if result.boxes is None:
                    continue

                for box in result.boxes:
                    class_id = int(box.cls[0])
                    confidence = float(box.conf[0])

                    if hasattr(result, 'names'):
                        class_name = result.names[class_id].lower()
                    else:
                        continue

                    # Check each category with proper threshold
                    if self._is_gun(class_name) and confidence > self.gun_threshold:
                        counts['guns'] += 1

                    elif self._is_knife(class_name) and confidence > self.knife_threshold:
                        counts['knives'] += 1

                    elif self._is_fight(class_name) and confidence > self.fight_threshold:
                        counts['fights'] += 1

            counts['total'] = counts['guns'] + counts['knives'] + counts['fights']
            return counts

        except Exception as e:
            print(f"   ⚠️ Detection error: {e}")
            return counts

    def _is_gun(self, name: str) -> bool:
        gun_words = ['gun', 'pistol', 'rifle', 'firearm', 'súng']
        return any(w in name for w in gun_words)

    def _is_knife(self, name: str) -> bool:
        knife_words = ['knife', 'dao', 'blade', 'sword']
        return any(w in name for w in knife_words)

    def _is_fight(self, name: str) -> bool:
        fight_words = ['fight', 'fighting', 'combat', 'violence']
        return any(w in name for w in fight_words)

    def process_video(self, video_path: str) -> Dict:
        """
        Process video with SMART frame skipping
        Auto-adjusts based on video duration
        """

        total = DetectionResult()

        try:
            cap = cv2.VideoCapture(video_path)
            if not cap.isOpened():
                return total.to_dict()

            # Get video info
            fps = cap.get(cv2.CAP_PROP_FPS)
            total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            duration = total_frames / fps if fps > 0 else 0

            # SMART frame skip based on duration
            if duration <= 10:  # Short video
                frame_skip = 5  # Check every 5th frame
                max_frames = 100
            elif duration <= 30:
                frame_skip = 10  # Check every 10th frame
                max_frames = 150
            elif duration <= 60:
                frame_skip = 15
                max_frames = 200
            else:  # Long video
                frame_skip = 30
                max_frames = 300

            print(f"\n📹 Video: {duration:.1f}s, {total_frames} frames")
            print(f"   Auto settings: skip={frame_skip}, max={max_frames}")

            frame_count = 0
            processed = 0
            nsfw_strikes = 0  # Count NSFW detections

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

                frame_count += 1

                # Skip frames
                if frame_count % frame_skip != 0:
                    continue

                # Max frame limit
                if processed >= max_frames:
                    break

                processed += 1

                # Process frame
                result = self.process_image(frame)

                # Accumulate
                total.nude_count += result.nude_count
                total.gun_count += result.gun_count
                total.knife_count += result.knife_count
                total.fight_count += result.fight_count

                # Early stop on multiple NSFW
                if result.nude_count > 0:
                    nsfw_strikes += 1
                    if nsfw_strikes >= 3:  # Stop after 3 NSFW frames
                        print(f"⛔ Early stop: {nsfw_strikes} NSFW frames")
                        break

                # Progress
                if processed % 50 == 0:
                    print(f"   Processed {processed} frames...")

            cap.release()

            # Final safety
            total_threats = total.nude_count + total.gun_count + total.knife_count + total.fight_count
            total.is_safe = (total_threats == 0)

            print(f"\n📊 Video complete: {processed} frames analyzed")
            print(f"   Total: N:{total.nude_count} G:{total.gun_count} K:{total.knife_count} F:{total.fight_count}")

            return total.to_dict()

        except Exception as e:
            print(f"❌ Video error: {e}")
            return total.to_dict()


def main():
    """Test the moderator"""

    moderator = SmartSequentialModerator()

    print("\n" + "=" * 50)
    print("🎯 SMART SEQUENTIAL MODERATOR")
    print("=" * 50)
    print("• Balanced NSFW threshold: 0.75")
    print("• Auto frame skipping for videos")
    print("• Simple output: counts + boolean")
    print("=" * 50)

    # Test
    test_image = "test.jpg"
    if os.path.exists(test_image):
        result = moderator.process_image(test_image)
        print(f"\nResult: {result.to_dict()}")


if __name__ == "__main__":
    main()