# engine.py
import cv2
import torch
import numpy as np
import base64
import re
from collections import deque
from ultralytics import YOLO
from torchvision import transforms
from torchvision.models import mobilenet_v2, MobileNet_V2_Weights
import torch.nn as nn
import time
import os
import json
from datetime import datetime

try:
    DEVICE = "mps" if torch.backends.mps.is_available() else "cpu"
except Exception:
    DEVICE = "cpu"
FRAME_BUFFER_SIZE = 16
VIOLENCE_THRESHOLD = 0.55
YOLO_CONFIDENCE = 0.4
MOTION_THRESHOLD = 0.35
MOTION_SUPPRESS = 0.85
ALERT_COOLDOWN = 10
ALERTS_LOG = "alerts.log"
ALERTS_JSONL = "alerts.jsonl"

# Model
class QuickViolenceNet(nn.Module):
    def __init__(self):
        super().__init__()
        base = mobilenet_v2(weights=MobileNet_V2_Weights.IMAGENET1K_V1)
        self.features = base.features
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.lstm = nn.LSTM(1280, 128, num_layers=2, batch_first=True, dropout=0.3)
        self.dropout = nn.Dropout(0.4)
        self.fc = nn.Linear(128, 2)

    def forward(self, x):
        B, T, C, H, W = x.shape
        x = x.view(B * T, C, H, W)
        x = self.pool(self.features(x)).squeeze(-1).squeeze(-1)
        x = x.view(B, T, -1)
        out, _ = self.lstm(x)
        return self.fc(self.dropout(out[:, -1]))

# Load models once
yolo = YOLO("yolov8n.pt")
classifier = QuickViolenceNet().to(DEVICE)
if os.path.exists("violence_classifier.pt"):
    classifier.load_state_dict(torch.load("violence_classifier.pt", map_location=DEVICE))
else:
    print("WARNING: violence_classifier.pt not found. Predictions may be inaccurate.")
classifier.eval()

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 _decode_data_url(image_data):
    if "," in image_data:
        image_data = image_data.split(",", 1)[1]
    image_data = re.sub(r"\s+", "", image_data)
    raw = base64.b64decode(image_data)
    arr = np.frombuffer(raw, dtype=np.uint8)
    frame = cv2.imdecode(arr, cv2.IMREAD_COLOR)
    if frame is None:
        raise ValueError("Could not decode image frame")
    return frame


def _analyze_frame(frame, state, location):
    curr_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    if state["prev_gray"] is not None:
        flow = cv2.calcOpticalFlowFarneback(
            state["prev_gray"], curr_gray, None, 0.5, 3, 15, 3, 5, 1.2, 0
        )
        mag, _ = cv2.cartToPolar(flow[..., 0], flow[..., 1])
        state["motion_mag"] = float(np.mean(mag))
    state["prev_gray"] = curr_gray.copy()

    results = yolo(frame, classes=[0], conf=YOLO_CONFIDENCE, verbose=False)

    roi = frame
    boxes = results[0].boxes
    person_count = 0
    if boxes is not None and len(boxes) > 0:
        person_count = len(boxes)
        xyxy = boxes.xyxy.cpu().numpy()
        areas = [(b[2] - b[0]) * (b[3] - b[1]) for b in xyxy]
        best = xyxy[np.argmax(areas)].astype(int)
        x1, y1, x2, y2 = best
        x1, y1 = max(0, x1), max(0, y1)
        x2, y2 = min(frame.shape[1], x2), min(frame.shape[0], y2)
        if x2 > x1 and y2 > y1:
            roi = frame[y1:y2, x1:x2]

    rgb = cv2.cvtColor(roi, cv2.COLOR_BGR2RGB)
    tensor = transform(rgb)
    state["buffer"].append(tensor)

    if len(state["buffer"]) == FRAME_BUFFER_SIZE:
        with torch.no_grad():
            clip = torch.stack(list(state["buffer"])).unsqueeze(0).to(DEVICE)
            logits = classifier(clip)
            probs = torch.softmax(logits, dim=1)
            raw_score = probs[0][1].item()

        if state["motion_mag"] < MOTION_THRESHOLD:
            raw_score *= MOTION_SUPPRESS
        state["score"] = raw_score

        if state["score"] > VIOLENCE_THRESHOLD:
            if state["violence_start_t"] is None:
                state["violence_start_t"] = time.time()
        else:
            state["violence_start_t"] = None

    duration = 0.0
    if state["violence_start_t"] is not None:
        duration = time.time() - state["violence_start_t"]

    return {
        "location": location,
        "score": round(float(state["score"]), 4),
        "is_violent": bool(state["score"] > VIOLENCE_THRESHOLD),
        "motion": round(float(state["motion_mag"]), 4),
        "duration_seconds": round(float(duration), 1),
        "threshold": VIOLENCE_THRESHOLD,
        "person_count": int(person_count),
        "frames_collected": len(state["buffer"]),
        "frames_required": FRAME_BUFFER_SIZE,
    }


class BrowserFrameAnalyzer:
    def __init__(self, location="Browser Camera"):
        self.location = location
        self.state = {
            "buffer": deque(maxlen=FRAME_BUFFER_SIZE),
            "score": 0.0,
            "prev_gray": None,
            "motion_mag": 1.0,
            "violence_start_t": None,
        }

    def analyze_data_url(self, image_data):
        frame = _decode_data_url(image_data)
        return _analyze_frame(frame, self.state, self.location)


class CameraEngine:
    def __init__(self, cam_id, location="Camera"):
        self.cam_id = cam_id
        self.cap = cv2.VideoCapture(cam_id)
        self.buffer = deque(maxlen=FRAME_BUFFER_SIZE)
        self.score = 0.0
        self.location = location
        self.prev_gray = None
        self.motion_mag = 1.0
        self.violence_start_t = None
        self.last_alert_time = 0.0
        self.last_alert_record = None
        self.is_open = self.cap.isOpened()

    def _ensure_camera_open(self):
        if self.cap is not None and self.cap.isOpened():
            self.is_open = True
            return True
        # Retry opening camera if it was disconnected or busy before.
        self.cap = cv2.VideoCapture(self.cam_id)
        self.is_open = self.cap.isOpened()
        return self.is_open

    def _trigger_alert(self, confidence, duration_seconds):
        now = time.time()
        if now - self.last_alert_time < ALERT_COOLDOWN:
            return
        self.last_alert_time = now

        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        msg = (
            f"[ALERT] {timestamp} | Location: {self.location} | "
            f"Confidence: {confidence:.1%} | Duration: {duration_seconds:.1f}s"
        )

        with open(ALERTS_LOG, "a") as f:
            f.write(msg + "\n")

        record = {
            "timestamp": timestamp,
            "location": self.location,
            "confidence": round(confidence, 4),
            "duration_seconds": round(duration_seconds, 1),
            "camera": str(self.location),
        }
        with open(ALERTS_JSONL, "a") as f:
            f.write(json.dumps(record) + "\n")
        self.last_alert_record = record

    def get_frame(self):
        if not self._ensure_camera_open():
            return None
        ret, frame = self.cap.read()
        if not ret:
            self.is_open = False
            return None

        curr_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        if self.prev_gray is not None:
            flow = cv2.calcOpticalFlowFarneback(
                self.prev_gray, curr_gray, None, 0.5, 3, 15, 3, 5, 1.2, 0
            )
            mag, _ = cv2.cartToPolar(flow[..., 0], flow[..., 1])
            self.motion_mag = float(np.mean(mag))
        self.prev_gray = curr_gray.copy()

        results = yolo(frame, classes=[0], conf=YOLO_CONFIDENCE, verbose=False)

        roi = frame
        boxes = results[0].boxes
        if boxes is not None and len(boxes) > 0:
            xyxy = boxes.xyxy.cpu().numpy()
            areas = [(b[2] - b[0]) * (b[3] - b[1]) for b in xyxy]
            best = xyxy[np.argmax(areas)].astype(int)
            x1, y1, x2, y2 = best
            x1, y1 = max(0, x1), max(0, y1)
            x2, y2 = min(frame.shape[1], x2), min(frame.shape[0], y2)
            if x2 > x1 and y2 > y1:
                roi = frame[y1:y2, x1:x2]

        try:
            rgb = cv2.cvtColor(roi, cv2.COLOR_BGR2RGB)
            tensor = transform(rgb)
            self.buffer.append(tensor)
        except Exception:
            return frame

        if len(self.buffer) == FRAME_BUFFER_SIZE:
            with torch.no_grad():
                clip = torch.stack(list(self.buffer)).unsqueeze(0).to(DEVICE)
                logits = classifier(clip)
                probs = torch.softmax(logits, dim=1)
                raw_score = probs[0][1].item()

            if self.motion_mag < MOTION_THRESHOLD:
                raw_score *= MOTION_SUPPRESS
            self.score = raw_score

            if self.score > VIOLENCE_THRESHOLD:
                if self.violence_start_t is None:
                    self.violence_start_t = time.time()
                duration = time.time() - self.violence_start_t
                self._trigger_alert(self.score, duration)
            else:
                self.violence_start_t = None

        annotated = results[0].plot()
        is_violent = self.score > VIOLENCE_THRESHOLD
        status_color = (0, 60, 220) if is_violent else (20, 180, 60)
        status_text = "!! VIOLENCE DETECTED !!" if is_violent else "NORMAL"
        suppressed = self.motion_mag < MOTION_THRESHOLD

        h_w = annotated.shape[1]
        overlay = annotated.copy()
        cv2.rectangle(overlay, (0, 0), (h_w, 75), (10, 10, 10), -1)
        cv2.addWeighted(overlay, 0.75, annotated, 0.25, 0, annotated)

        cv2.putText(
            annotated,
            f"Status: {status_text}",
            (10, 26),
            cv2.FONT_HERSHEY_DUPLEX,
            0.75,
            status_color,
            2,
        )
        cv2.putText(
            annotated,
            f"Score: {self.score:.2f}{' [suppressed]' if suppressed else ''} | {self.location}",
            (10, 52),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.50,
            (255, 255, 255),
            1,
        )
        cv2.putText(
            annotated,
            f"Motion: {'HIGH' if self.motion_mag >= MOTION_THRESHOLD else 'LOW'} ({self.motion_mag:.2f})",
            (10, 72),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.42,
            (100, 200, 100) if self.motion_mag >= MOTION_THRESHOLD else (100, 100, 200),
            1,
        )

        if is_violent and int(time.time() * 2) % 2 == 0:
            cv2.rectangle(annotated, (2, 2), (h_w - 2, annotated.shape[0] - 2), (0, 0, 255), 3)

        return annotated

    def status(self):
        duration = 0.0
        if self.violence_start_t is not None:
            duration = time.time() - self.violence_start_t
        return {
            "location": self.location,
            "camera_id": self.cam_id,
            "camera_open": bool(self._ensure_camera_open()),
            "score": round(float(self.score), 4),
            "is_violent": bool(self.score > VIOLENCE_THRESHOLD),
            "motion": round(float(self.motion_mag), 4),
            "duration_seconds": round(float(duration), 1),
            "threshold": VIOLENCE_THRESHOLD,
        }