import os
import sys
import torch
import logging
from pathlib import Path
from tqdm import tqdm
import cv2
import numpy as np

# Add src to path to import DeepfakeDetectorAPI
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
from src.deepfake.deepfake_detector import DeepfakeDetectorAPI

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(message)s')
logger = logging.getLogger("DeepfakeAnalyzer")

def analyze_folder(detector, folder_path, label):
    folder = Path(folder_path)
    if not folder.exists():
        logger.error(f"Folder not found: {folder}")
        return None

    images = list(folder.glob("*.jpg")) + list(folder.glob("*.png")) + list(folder.glob("*.jpeg"))
    if not images:
        logger.error(f"No images found in {folder}")
        return None
    
    logger.info(f"\nAnalyzing '{label}' ({folder})... Found {len(images)} images.")
    
    total_prob = 0
    attack_count = 0
    real_count = 0
    
    # Increase sample size for robustness
    sample_size = min(len(images), 50)
    logger.info(f"Sampling {sample_size} images using Deepfake Detector (ViT)...")
    
    for img_path in tqdm(images[:sample_size]):
        try:
            # Detect returns dict with 'deepfake_score'
            result = detector.detect(str(img_path))
            score = result['deepfake_score']
            
            # Threshold 0.5 
            threshold = 0.5
            is_attack = score > threshold
            
            total_prob += score
            if is_attack:
                attack_count += 1
            else:
                real_count += 1
                
        except Exception as e:
            logger.warning(f"Failed to analyze {img_path.name}: {e}")
            
    avg_prob = total_prob / sample_size if sample_size > 0 else 0
    attack_percent = (attack_count / sample_size) * 100 if sample_size > 0 else 0
    
    return {
        'label': label,
        'path': str(folder),
        'sample_size': sample_size,
        'avg_score': avg_prob,
        'attack_percentage': attack_percent,
        'real_count': real_count,
        'attack_count': attack_count
    }

def main():
    folder_a = "/media/juanquy/Dev/Users photos/"
    folder_b = "/media/juanquy/Dev/a small sample/"
    
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    logger.info(f"Using device: {device}")
    
    logger.info("Initializing DeepfakeDetectorAPI (dima806/ViT)...")
    try:
        detector = DeepfakeDetectorAPI(device=device)
        logger.info("Detector initialized successfully.")
    except Exception as e:
        logger.error(f"Failed to init detector: {e}")
        return

    results = []
    
    res_a = analyze_folder(detector, folder_a, "Real Photos Folder")
    if res_a: results.append(res_a)
    
    res_b = analyze_folder(detector, folder_b, "AI Generated Folder")
    if res_b: results.append(res_b)
    
    print("\n" + "="*80)
    print(f"{'Folder Analysis Report (Deepfake/GenAI Detector)':^80}")
    print("="*80)
    
    for res in results:
        print(f"\n--- {res['label']} ---")
        print(f"Path: {res['path']}")
        print(f"Avg AI Probability: {res['avg_score']:.4f} (0=Real, 1=AI)")
        print(f"AI Detection Rate: {res['attack_percentage']:.1f}% ({res['attack_count']}/{res['sample_size']})")
        
        verdict = "UNKNOWN"
        if res['attack_percentage'] > 70:
            verdict = ">> AI GENERATED <<"
        elif res['attack_percentage'] < 30: 
            verdict = ">> REAL PHOTOS <<"
        else:
            verdict = ">> MIXED / UNCERTAIN <<"
            
        print(f"Verdict: {verdict}")

if __name__ == "__main__":
    main()