import torch
from torch import nn
from torchvision import models, transforms
import cv2
import face_recognition
import numpy as np
from PIL import Image
import os

# --- 1. MODEL ARCHITECTURE (From views.py) ---
class DeepfakeModel(nn.Module):
    def __init__(self, num_classes=2, latent_dim=2048, lstm_layers=1, hidden_dim=2048, bidirectional=False):
        super(DeepfakeModel, self).__init__()
        resnext = models.resnext50_32x4d(pretrained=True)
        self.model = nn.Sequential(*list(resnext.children())[:-2])
        self.lstm = nn.LSTM(latent_dim, hidden_dim, lstm_layers, bidirectional)
        self.dp = nn.Dropout(0.4)
        self.linear1 = nn.Linear(2048, num_classes)
        self.avgpool = nn.AdaptiveAvgPool2d(1)

    def forward(self, x):
        batch_size, seq_length, c, h, w = x.shape
        x = x.view(batch_size * seq_length, c, h, w)
        fmap = self.model(x)
        x = self.avgpool(fmap)
        x = x.view(batch_size, seq_length, 2048)
        x_lstm, _ = self.lstm(x, None)
        return fmap, self.dp(self.linear1(x_lstm[:, -1, :]))

# --- 2. CUSTOM HANDLER FOR HUGGING FACE ---
class EndpointHandler():
    def __init__(self, path=""):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model = DeepfakeModel().to(self.device)
        
        # Look for the .pt file in the repository
        model_weight_path = os.path.join(path, "pytorch_model.bin")
        if not os.path.exists(model_weight_path):
             # Fallback to any .pt file if pytorch_model.bin isn't found
             for f in os.listdir(path):
                 if f.endswith(".pt"):
                     model_weight_path = os.path.join(path, f)
                     break
        
        self.model.load_state_dict(torch.load(model_weight_path, map_location=self.device))
        self.model.eval()
        
        self.transform = transforms.Compose([
            transforms.ToPILImage(),
            transforms.Resize((112, 112)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])

    def __call__(self, data):
        # Data normally contains "inputs" which is the raw video file (bytes)
        inputs = data.pop("inputs", data)
        
        # Save bytes to temporary file to process with OpenCV
        temp_video = "temp_video.mp4"
        with open(temp_video, "wb") as f:
            f.write(inputs)
            
        # 1. Extract Frames & Crop Faces
        cap = cv2.VideoCapture(temp_video)
        frames = []
        while cap.isOpened() and len(frames) < 20: # Process first 20 frames
            ret, frame = cap.read()
            if not ret: break
            
            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            face_locations = face_recognition.face_locations(rgb_frame)
            
            if len(face_locations) > 0:
                top, right, bottom, left = face_locations[0]
                face_crop = rgb_frame[top:bottom, left:right]
                frames.append(self.transform(face_crop))
        
        cap.release()
        os.remove(temp_video)

        if len(frames) < 10:
            return {"error": "Not enough faces detected in video."}

        # 2. Run Inference
        input_tensor = torch.stack(frames).unsqueeze(0).to(self.device) # [1, Seq, C, H, W]
        with torch.no_grad():
            _, outputs = self.model(input_tensor)
            probabilities = torch.softmax(outputs, dim=1)
            confidence, prediction = torch.max(probabilities, 1)
            
        result = "REAL" if prediction.item() == 1 else "FAKE"
        return {
            "label": result,
            "confidence": round(float(confidence.item()) * 100, 2)
        }