#!/usr/bin/env python3
"""
Script to process assets containing a specific person and align faces.
"""

import os
import io
import requests
import concurrent.futures
import argparse
from datetime import datetime

from PIL import Image, ImageOps
import numpy as np
import cv2
import dlib
from tqdm import tqdm


def get_assets_with_person(api_key, base_url, person_id):
    headers = {
        'Content-Type': 'application/json',
        'Accept': 'application/json',
        'x-api-key': api_key,
    }
    url = f"{base_url}/search/metadata"
    all_assets = []
    payload = {
        "page": 1,
        "type": "IMAGE",
        "personIds": [person_id],
        "withArchived": False,
        "withDeleted": True,
        "withExif": True,
        "withPeople": True,
        "withStacked": True,
    }
    while payload["page"] is not None:
        response = requests.post(url, headers=headers, json=payload)
        if response.status_code != 200:
            print(f"Error fetching page {payload['page']}: {response.status_code} - {response.text}")
            break
        data = response.json()
        if not data:
            break
        all_assets.extend(data['assets']['items'])
        print(f"Fetched page {payload['page']} with {len(data['assets']['items'])} assets")
        payload["page"] = data['assets'].get('nextPage')
    return all_assets


def download_asset(api_key, base_url, asset_id):
    headers = {'x-api-key': api_key}
    response = requests.get(f'{base_url}/assets/{asset_id}/original', headers=headers)
    response.raise_for_status()
    return response.content


def format_timestamp(timestamp):
    dt = datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
    return dt.strftime("%Y%m%d_%H%M%S")


def crop_face_from_metadata(image, face_data, padding_percent):
    face_img_width = face_data.get("imageWidth")
    face_img_height = face_data.get("imageHeight")
    img_width, img_height = image.size
    scale_x = img_width / face_img_width
    scale_y = img_height / face_img_height
    x1 = int(face_data.get("boundingBoxX1", 0) * scale_x)
    x2 = int(face_data.get("boundingBoxX2", 0) * scale_x)
    y1 = int(face_data.get("boundingBoxY1", 0) * scale_y)
    y2 = int(face_data.get("boundingBoxY2", 0) * scale_y)
    w = x2 - x1
    h = y2 - y1
    padding = int(max(w, h) * padding_percent)
    new_x1 = max(x1 - padding, 0)
    new_y1 = max(y1 - padding, 0)
    new_x2 = min(x2 + padding, img_width)
    new_y2 = min(y2 + padding, img_height)
    return image.crop((new_x1, new_y1, new_x2, new_y2))


def get_head_pose(shape, img_size):
    image_points = np.array([
        (shape.part(30).x, shape.part(30).y),
        (shape.part(8).x, shape.part(8).y),
        (shape.part(36).x, shape.part(36).y),
        (shape.part(45).x, shape.part(45).y),
        (shape.part(48).x, shape.part(48).y),
        (shape.part(54).x, shape.part(54).y)
    ], dtype="double")
    model_points = np.array([
        (0.0, 0.0, 0.0),
        (0.0, -330.0, -65.0),
        (-225.0, 170.0, -135.0),
        (225.0, 170.0, -135.0),
        (-150.0, -150.0, -125.0),
        (150.0, -150.0, -125.0)
    ])
    w, h = img_size
    focal_length = w
    center = (w / 2, h / 2)
    camera_matrix = np.array(
        [[focal_length, 0, center[0]],
         [0, focal_length, center[1]],
         [0, 0, 1]], dtype="double"
    )
    dist_coeffs = np.zeros((4, 1))
    success, rotation_vector, translation_vector = cv2.solvePnP(
        model_points, image_points, camera_matrix, dist_coeffs, flags=cv2.SOLVEPNP_ITERATIVE
    )
    rotation_mat, _ = cv2.Rodrigues(rotation_vector)
    proj_matrix = np.hstack((rotation_mat, translation_vector))
    _, _, _, _, _, _, eulerAngles = cv2.decomposeProjectionMatrix(proj_matrix)
    pitch, yaw, roll = [float(angle) for angle in eulerAngles]
    return pitch, yaw, roll


def align_face(image, predictor, detector, desired_face_width, desired_face_height,
               desired_left_eye, pose_threshold):
    image_np = np.array(image)
    gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
    rects = detector(gray, 2)
    if not rects:
        print("No face detected in the crop. Discarding.")
        return None
    rect = rects[0]
    shape = predictor(gray, rect)
    img_size = (image_np.shape[1], image_np.shape[0])
    pitch, yaw, roll = get_head_pose(shape, img_size)
    if abs(abs(pitch) - 180) > pose_threshold or abs(yaw) > pose_threshold:
        print(f"Face not frontal enough: pitch={pitch:.2f}, yaw={yaw:.2f}, roll={roll:.2f}. Discarding.")
        return None
    shape_np = np.array([(shape.part(i).x, shape.part(i).y) for i in range(68)], dtype="int")
    left_eye_center = shape_np[36:42].mean(axis=0).astype("int")
    right_eye_center = shape_np[42:48].mean(axis=0).astype("int")
    dY = right_eye_center[1] - left_eye_center[1]
    dX = right_eye_center[0] - left_eye_center[0]
    angle = np.degrees(np.arctan2(dY, dX))
    eye_distance = np.linalg.norm(right_eye_center - left_eye_center)
    desired_right_eye_x = 1.0 - desired_left_eye[0]
    desired_eye_distance = (desired_right_eye_x - desired_left_eye[0]) * desired_face_width
    scale = desired_eye_distance / eye_distance
    eyes_center = ((left_eye_center[0] + right_eye_center[0]) / 2.0,
                   (left_eye_center[1] + right_eye_center[1]) / 2.0)
    adjusted_scale = scale * 0.8
    M = cv2.getRotationMatrix2D(eyes_center, angle, adjusted_scale)
    extra_offset_x = 10
    tX = desired_face_width * 0.5 + extra_offset_x
    tY = desired_face_height * desired_left_eye[1]
    M[0, 2] += (tX - eyes_center[0])
    M[1, 2] += (tY - eyes_center[1])
    aligned_face_np = cv2.warpAffine(
        image_np,
        M,
        (desired_face_width, desired_face_height),
        flags=cv2.INTER_CUBIC,
        borderMode=cv2.BORDER_REPLICATE
    )
    return Image.fromarray(aligned_face_np)


def process_asset_worker(asset, api_key, base_url, person_id, output_folder,
                         padding_percent, min_face_width, min_face_height,
                         resize_width, resize_height, pose_threshold, desired_left_eye):
    local_detector = dlib.get_frontal_face_detector()
    predictor_path = "shape_predictor_68_face_landmarks.dat"
    local_predictor = dlib.shape_predictor(predictor_path)
    try:
        asset_id = asset['id']
        timestamp = format_timestamp(asset['fileCreatedAt'])
        image_bytes = download_asset(api_key, base_url, asset_id)
        image = Image.open(io.BytesIO(image_bytes))
        image = ImageOps.exif_transpose(image)
        image = image.convert("RGB")
    except Exception as e:
        print(f"Error processing asset {asset.get('id')}: {e}")
        return None
    matching_person = next((p for p in asset.get('people', []) if p.get('id') == person_id), None)
    if not matching_person:
        print("Subject not in image.")
        return None
    faces = matching_person.get('faces', [])
    if not faces:
        print("No face data available.")
        return None
    face_data = faces[0]
    cropped_face = crop_face_from_metadata(image, face_data, padding_percent)
    face_width, face_height = cropped_face.size
    if face_width < min_face_width or face_height < min_face_height:
        print(f"Face resolution too low ({face_width}x{face_height}).")
        return None
    aligned_face = align_face(cropped_face, local_predictor, local_detector,
                              desired_face_width=resize_width,
                              desired_face_height=resize_height,
                              desired_left_eye=desired_left_eye,
                              pose_threshold=pose_threshold)
    if aligned_face is None:
        return None
    filename = os.path.join(output_folder, f"{timestamp}.jpg")
    aligned_face.save(filename)
    return filename


def main():
    parser = argparse.ArgumentParser(description="Process and align faces from assets.")
    parser.add_argument("--api-key", required=True, help="API key for authentication")
    parser.add_argument("--base-url", required=True, help="Base URL for the API")
    parser.add_argument("--person-id", required=True, help="ID of the person to search for")
    parser.add_argument("--output-folder", default="output", help="Folder to save output images")
    parser.add_argument("--padding-percent", type=float, default=0.3, help="Padding percentage for face crop")
    parser.add_argument("--resize-width", type=int, default=512, help="Output image width")
    parser.add_argument("--resize-height", type=int, default=512, help="Output image height")
    parser.add_argument("--min-face-width", type=int, default=128, help="Minimum face width")
    parser.add_argument("--min-face-height", type=int, default=128, help="Minimum face height")
    parser.add_argument("--pose-threshold", type=float, default=25, help="Threshold for acceptable head pose")
    parser.add_argument("--desired-left-eye", type=float, nargs=2, default=[0.35, 0.45],
                        help="Desired left eye position as fraction (x y) in the output image")
    parser.add_argument("--max-workers", type=int, default=4, help="Maximum number of parallel workers")
    args = parser.parse_args()

    os.makedirs(args.output_folder, exist_ok=True)

    assets = get_assets_with_person(args.api_key, args.base_url, args.person_id)
    print(f"Found {len(assets)} assets containing the person.")

    process_args = (
        args.api_key, args.base_url, args.person_id, args.output_folder,
        args.padding_percent, args.min_face_width, args.min_face_height,
        args.resize_width, args.resize_height, args.pose_threshold, tuple(args.desired_left_eye)
    )

    with concurrent.futures.ProcessPoolExecutor(max_workers=args.max_workers) as executor:
        results = list(tqdm(
            executor.map(lambda asset: process_asset_worker(asset, *process_args), assets),
            total=len(assets)
        ))
    processed_files = [r for r in results if r is not None]
    print(f"Finished processing. {len(processed_files)} images saved.")


if __name__ == "__main__":
    main()