import cv2
import time
import os
import threading
import warnings
from datetime import datetime
from feat import Detector
import torch
import mediapipe as mp
import pandas as pd
import db_helper as db 

from pathlib import Path
from eyeFeature_new import compute_features_from_parquet

# Suppress specific Protobuf deprecation warnings from the library
warnings.filterwarnings(
    "ignore",
    message=r".*SymbolDatabase\.GetPrototype\(\) is deprecated.*",
    category=UserWarning,
    module=r"google\.protobuf\.symbol_database"
)

# --- Configuration & Hyperparameters ---
DB_PATH = Path("~/MSY_FS/databases/database.sqlite").expanduser()
CAMERA_INDEX = 0
OUTPUT_DIR = Path("recordings")
VIDEO_DURATION = 50  # Seconds per recording segment
START_INTERVAL = 5   # Delay between starting overlapping recordings
FPS = 25.0           # Target Frames Per Second

# Global feature storage - Updated to be thread-safe in production environments
eye_tracking_features = {}

if not OUTPUT_DIR.exists():
    OUTPUT_DIR.mkdir(parents=True)

# Initialize the AU-Detector globally to optimize VRAM/RAM usage
print("[INFO] Initializing Facial Action Unit Detector (XGB)...")
detector = Detector(au_model="xgb") 

# --- MediaPipe FaceMesh Configuration ---
mp_face_mesh = mp.solutions.face_mesh
face_mesh = mp_face_mesh.FaceMesh(
    static_image_mode=False,
    max_num_faces=1,
    refine_landmarks=True,  # Mandatory for Iris tracking
    min_detection_confidence=0.5,
    min_tracking_confidence=0.5
)

# Landmark Indices for Oculometrics
LEFT_IRIS = [474, 475, 476, 477]
RIGHT_IRIS = [469, 470, 471, 472]
LEFT_EYE_LIDS = (159, 145)
RIGHT_EYE_LIDS = (386, 374)
EYE_OPEN_THRESHOLD = 6

# Bounding box indices for eye regions
LEFT_EYE_ALL = [33, 7, 163, 144, 145, 153, 154, 155, 133, 173, 157, 158, 159, 160, 161, 246]
RIGHT_EYE_ALL = [263, 249, 390, 373, 374, 380, 381, 382, 362, 398, 384, 385, 386, 387, 388, 466]

def eye_openness(landmarks, top_idx, bottom_idx, img_height):
    """Calculates the vertical distance between eyelids normalized by image height."""
    top = landmarks[top_idx]
    bottom = landmarks[bottom_idx]
    return abs(top.y - bottom.y) * img_height

def compute_gaze(landmarks, iris_center, eye_indices, w, h):
    """
    Computes normalized gaze coordinates (0.0 to 1.0) relative to the eye's 
    internal bounding box.
    """
    iris_x, iris_y = iris_center

    eye_points = []
    for idx in eye_indices:
        lm = landmarks[idx]
        eye_points.append((lm.x * w, lm.y * h))
    
    xs = [p[0] for p in eye_points]
    ys = [p[1] for p in eye_points]

    eye_left = min(xs)
    eye_right = max(xs)
    eye_top = min(ys)
    eye_bottom = max(ys)

    eye_width = eye_right - eye_left
    eye_height = eye_bottom - eye_top

    if eye_width < 1 or eye_height < 1:
        return 0.5, 0.5

    gaze_x = (iris_x - eye_left) / eye_width
    gaze_y = (iris_y - eye_top) / eye_height

    return gaze_x, gaze_y

def extract_aus(path, skip_frames):
    """
    Infers facial Action Units from video file. 
    Uses torch.no_grad() to optimize inference and prevent memory leakage.
    """
    with torch.no_grad():
        try:
            video_prediction = detector.detect_video(
                path, 
                skip_frames=skip_frames, 
                face_detection_threshold=0.95
            )
            # Compute temporal mean of Action Units across the segment
            return video_prediction.aus.mean()
        except Exception as e:
            print(f"[ERROR] AU Extraction failed: {e}")
            return None

def process_and_store_analysis(video_path, db_path):
    """
    Worker function: Handles AU extraction, data merging, and SQL persistence.
    Designed to run in a background thread.
    """
    try:
        print(f"[THREAD] Analyzing segment: {video_path}")
        # Analysis sampling: one frame every 5 seconds
        output = extract_aus(video_path, skip_frames=int(FPS * 5))
        if output is not None:
            # Verbindung für diesen Thread öffnen (SQLite Sicherheit)
            conn, cursor = db.connect_db(db_path)
            
            # Prepare payload: Prefix keys to distinguish facial AUs
            data_to_insert = output.to_dict()

            data_to_insert = {
                f"FACE_{k}_mean": v for k, v in data_to_insert.items()
            }

            now = datetime.now()
            ticks = int(time.mktime(now.timetuple()))

            data_to_insert['start_time'] = [ticks]
            data_to_insert = data_to_insert | eye_tracking_features
            
            # making sure that dynamic AU-columns are lists
            # (insert_rows_into_table expects lists for every key)
            final_payload = {k: [v] if not isinstance(v, list) else v for k, v in data_to_insert.items()}


            db.insert_rows_into_table(conn, cursor, "feature_table", final_payload)
            
            db.disconnect_db(conn, cursor)
            print(f"[SUCCESS] Data persisted for {os.path.basename(video_path)}")
            
            # Cleanup temporary files to save disk space
            os.remove(video_path)
            os.remove(video_path.replace(".avi", "_gaze.parquet"))
            

    except Exception as e:
        print(f"[ERROR] Threaded analysis failed for {video_path}: {e}")


class VideoRecorder:
    """Manages the asynchronous writing of video frames to disk."""
    def __init__(self, filename, width, height, db_path):
        self.gaze_data = []
        self.filename = filename
        self.db_path = db_path
        fourcc = cv2.VideoWriter_fourcc(*'XVID')
        self.out = cv2.VideoWriter(filename, fourcc, FPS, (width, height))
        self.frames_to_record = int(VIDEO_DURATION * FPS)
        self.frames_count = 0
        self.is_finished = False

    def write_frame(self, frame):
        if self.frames_count < self.frames_to_record:
            self.out.write(frame)
            self.frames_count += 1
        else:
            self.finish()

    def finish(self):
        if not self.is_finished:
            self.out.release()
            self.is_finished = True
            abs_path = os.path.abspath(self.filename)
            print(f"Video saved: {self.filename}")

            # Trigger background analysis thread
            # Passing a snapshot of eye_tracking_features to avoid race conditions
            analysis_thread = threading.Thread(target=process_and_store_analysis, args=(abs_path, self.db_path))
            analysis_thread.daemon = True  # ends when the program ends
            analysis_thread.start()

class GazeRecorder:
    """Handles the collection and Parquet serialization of oculometric data."""
    def __init__(self, filename):
        self.filename = filename
        self.frames_to_record = int(VIDEO_DURATION * FPS)
        self.frames_count = 0
        self.gaze_data = []
        self.is_finished = False

    def write_frame(self, gaze_row):
        if self.frames_count < self.frames_to_record:
            self.gaze_data.append(gaze_row)
            self.frames_count += 1
        else:
            self.finish()

    def finish(self):
        if not self.is_finished:
            df = pd.DataFrame(self.gaze_data)
            df.to_parquet(self.filename, engine="pyarrow", index=False)

            # Extract high-level features from raw gaze points
            print(f"Gaze-Parquet saved: {self.filename}")
            features = compute_features_from_parquet(self.filename)
            print("Features:", features)
            self.is_finished = True
            eye_tracking_features = features

def main():
    cap = cv2.VideoCapture(CAMERA_INDEX)
    if not cap.isOpened():
        print("[CRITICAL] Could not access camera.")
        return

    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
    active_video_recorders = []
    active_gaze_recorders = []
    last_start_time = 0
    
    print("[INFO] Recording started. Press 'q' to terminate.")

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

            # Pre-processing for MediaPipe
            rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            h, w, _ = frame.shape
            results = face_mesh.process(rgb)

            # Default feature values
            left_valid = 0
            right_valid = 0
            left_diameter = None
            right_diameter = None

            left_gaze_x = None
            left_gaze_y = None
            right_gaze_x = None
            right_gaze_y = None

            if results.multi_face_landmarks:
                face_landmarks = results.multi_face_landmarks[0]

                left_open = eye_openness(
                    face_landmarks.landmark,
                    LEFT_EYE_LIDS[0],
                    LEFT_EYE_LIDS[1],
                    h
                )

                right_open = eye_openness(
                    face_landmarks.landmark,
                    RIGHT_EYE_LIDS[0],
                    RIGHT_EYE_LIDS[1],
                    h
                )

                left_valid = 1 if left_open > EYE_OPEN_THRESHOLD else 0
                right_valid = 1 if right_open > EYE_OPEN_THRESHOLD else 0

                for eye_name, eye_indices in [("left", LEFT_IRIS), ("right", RIGHT_IRIS)]:
                    iris_points = []

                    for idx in eye_indices:
                        lm = face_landmarks.landmark[idx]
                        x_i, y_i = int(lm.x * w), int(lm.y * h)
                        iris_points.append((x_i, y_i))

                    if len(iris_points) == 4:
                        cx = int(sum(p[0] for p in iris_points) / 4)
                        cy = int(sum(p[1] for p in iris_points) / 4)

                        radius = max(
                            ((x - cx) ** 2 + (y - cy) ** 2) ** 0.5
                            for (x, y) in iris_points
                        )

                        diameter = 2 * radius

                        cv2.circle(frame, (cx, cy), int(radius), (0, 255, 0), 2)

                        if eye_name == "left" and left_valid:
                            left_diameter = diameter
                            left_gaze_x, left_gaze_y = compute_gaze(
                                face_landmarks.landmark,
                                (cx, cy),
                                RIGHT_EYE_ALL,
                                w, h
                            )

                        elif eye_name == "right" and right_valid:
                            right_diameter = diameter
                            right_gaze_x, right_gaze_y = compute_gaze(
                                face_landmarks.landmark,
                                (cx, cy),
                                LEFT_EYE_ALL,
                                w, h
                            )

            gaze_row = {
                "timestamp": time.time(),
                "EYE_LEFT_GAZE_POINT_ON_DISPLAY_AREA_X": left_gaze_x,
                "EYE_LEFT_GAZE_POINT_ON_DISPLAY_AREA_Y": left_gaze_y,
                "EYE_RIGHT_GAZE_POINT_ON_DISPLAY_AREA_X": right_gaze_x,
                "EYE_RIGHT_GAZE_POINT_ON_DISPLAY_AREA_Y": right_gaze_y,
                "EYE_LEFT_PUPIL_VALIDITY": left_valid,
                "EYE_RIGHT_PUPIL_VALIDITY": right_valid,
                "EYE_LEFT_PUPIL_DIAMETER": left_diameter,
                "EYE_RIGHT_PUPIL_DIAMETER": right_diameter
            }

            current_time = time.time()

            if current_time - last_start_time >= START_INTERVAL:
                timestamp = datetime.now().strftime("%H%M%S")
                filename = os.path.join(OUTPUT_DIR, f"rec_{timestamp}.avi")
                video_recorder = VideoRecorder(filename, width, height, DB_PATH)

                gaze_filename = filename.replace(".avi", "_gaze.parquet")
                gaze_recorder = GazeRecorder(gaze_filename)

                active_video_recorders.append(video_recorder)
                active_gaze_recorders.append(gaze_recorder)

                last_start_time = current_time

            for v_rec, g_rec in zip(active_video_recorders[:], active_gaze_recorders[:]):

                v_rec.write_frame(frame)
                g_rec.write_frame(gaze_row)

                if v_rec.is_finished:
                    active_video_recorders.remove(v_rec)

                if g_rec.is_finished:
                    active_gaze_recorders.remove(g_rec)

            cv2.imshow('Kamera Livestream', frame)
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break

            time.sleep(1/FPS)

    finally:
        face_mesh.close()
        cap.release()
        cv2.destroyAllWindows()

        print("[INFO] Stream closed. Waiting for background analysis to complete...")

if __name__ == "__main__":
    main()