#!/usr/bin/env python3
import json
import sys
import os
import io
import argparse
import os
os.environ["MPLCONFIGDIR"] = "/tmp/matplotlib"

import uuid
import base64
import logging
import time
import copy
import cv2
import insightface
import numpy as np
from typing import List, Union
from PIL import Image
from restoration import *
from flask import Flask, request, jsonify, make_response
from waitress import serve

LOG_LEVEL = logging.DEBUG
TMP_PATH = '/tmp/inswapper'
script_dir = os.path.dirname(os.path.abspath(__file__))
log_path = ''

# Mac does not have permission to /var/log for example
if sys.platform == 'linux':
    log_path = '/var/log/'

logging.basicConfig(
    filename=f'{log_path}inswapper.log',
    format='%(asctime)s : %(levelname)s : %(message)s',
    level=LOG_LEVEL
)

logging.getLogger().addHandler(logging.StreamHandler(sys.stdout))


def process_request(request_obj):
    try:
        logging.debug('Swapping face')
        face_swap_timer = Timer()
        result_image = face_swap(request_obj['source_image'], request_obj['target_image'])
        face_swap_time = face_swap_timer.get_elapsed_time()
        logging.info(f'Time taken to swap face: {face_swap_time} seconds')

        response = {
            'status': 'ok',
            'image': result_image
        }
    except Exception as e:
        logging.error(e)
        response = {
            'status': 'error',
            'msg': 'Face swap failed',
            'detail': str(e)
        }

    return response


class Timer:
    def __init__(self):
        self.start = time.time()

    def restart(self):
        self.start = time.time()

    def get_elapsed_time(self):
        end = time.time()
        return round(end - self.start, 1)


def get_args():
    parser = argparse.ArgumentParser(
        description='Inswapper REST API'
    )

    parser.add_argument(
        '-p', '--port',
        help='Port to listen on',
        type=int,
        default=80
    )

    parser.add_argument(
        '-H', '--host',
        help='Host to bind to',
        default='0.0.0.0'
    )

    return parser.parse_args()


def determine_file_extension(image_data):
    try:
        if image_data.startswith('/9j/'):
            image_extension = '.jpg'
        elif image_data.startswith('iVBORw0Kg'):
            image_extension = '.png'
        else:
            # Default to png if we can't figure out the extension
            image_extension = '.png'
    except Exception as e:
        image_extension = '.png'

    return image_extension


def write_base64_to_disk(file_b64: str, file_path: str):
    with open(file_path, 'wb') as file:
        file.write(base64.b64decode(file_b64))


def get_face_swap_model(model_path: str):
    model = insightface.model_zoo.get_model(model_path)
    return model


def get_face_analyser(model_path: str,
                      det_size=(320, 320)):
    face_analyser = insightface.app.FaceAnalysis(name="buffalo_l", root="./checkpoints")
    face_analyser.prepare(ctx_id=0, det_size=det_size)
    return face_analyser


def get_one_face(face_analyser,
                 frame:np.ndarray):
    face = face_analyser.get(frame)
    try:
        return min(face, key=lambda x: x.bbox[0])
    except ValueError:
        return None


def get_many_faces(face_analyser,
                   frame:np.ndarray):
    """
    get faces from left to right by order
    """
    try:
        face = face_analyser.get(frame)
        return sorted(face, key=lambda x: x.bbox[0])
    except IndexError:
        return None


def swap_face(face_swapper,
              source_faces,
              target_faces,
              source_index,
              target_index,
              temp_frame):
    """
    paste source_face on target image
    """
    source_face = source_faces[source_index]
    target_face = target_faces[target_index]

    return face_swapper.get(temp_frame, target_face, source_face, paste_back=True)


def process(source_img: Union[Image.Image, List],
            target_img: Image.Image,
            source_indexes: str,
            target_indexes: str,
            model: str):

    # load face_analyser
    face_analyser = get_face_analyser(model)

    # load face_swapper
    model_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), model)
    face_swapper = get_face_swap_model(model_path)

    # read target image
    target_img = cv2.cvtColor(np.array(target_img), cv2.COLOR_RGB2BGR)

    # detect faces that will be replaced in target_img
    target_faces = get_many_faces(face_analyser, target_img)
    num_target_faces = len(target_faces)
    num_source_images = len(source_img)

    if target_faces is not None:
        temp_frame = copy.deepcopy(target_img)
        if isinstance(source_img, list) and num_source_images == num_target_faces:
            logging.debug('Replacing the faces in the target image from left to right by order')
            for i in range(num_target_faces):
                source_faces = get_many_faces(face_analyser, cv2.cvtColor(np.array(source_img[i]), cv2.COLOR_RGB2BGR))
                source_index = i
                target_index = i

                if source_faces is None:
                    raise Exception('No source faces found!')

                temp_frame = swap_face(
                    face_swapper,
                    source_faces,
                    target_faces,
                    source_index,
                    target_index,
                    temp_frame
                )
        elif num_source_images == 1:
            # detect source faces that will be replaced into the target image
            source_faces = get_many_faces(face_analyser, cv2.cvtColor(np.array(source_img[0]), cv2.COLOR_RGB2BGR))
            num_source_faces = len(source_faces)
            logging.debug(f'Source faces: {num_source_faces}')
            logging.debug(f'Target faces: {num_target_faces}')

            if source_faces is None:
                raise Exception('No source faces found!')

            if target_indexes == "-1":
                if num_source_faces == 1:
                    logging.debug('Replacing all faces in target image with the same face from the source image')
                    num_iterations = num_target_faces
                elif num_source_faces < num_target_faces:
                    logging.debug('There are less faces in the source image than the target image, replacing as many as we can')
                    num_iterations = num_source_faces
                elif num_target_faces < num_source_faces:
                    logging.debug('There are less faces in the target image than the source image, replacing as many as we can')
                    num_iterations = num_target_faces
                else:
                    logging.debug('Replacing all faces in the target image with the faces from the source image')
                    num_iterations = num_target_faces

                for i in range(num_iterations):
                    source_index = 0 if num_source_faces == 1 else i
                    target_index = i

                    temp_frame = swap_face(
                        face_swapper,
                        source_faces,
                        target_faces,
                        source_index,
                        target_index,
                        temp_frame
                    )
            elif source_indexes == '-1' and target_indexes == '-1':
                logging.debug('Replacing specific face(s) in the target image with the face from the source image')
                target_indexes = target_indexes.split(',')
                source_index = 0

                for target_index in target_indexes:
                    target_index = int(target_index)

                    temp_frame = swap_face(
                        face_swapper,
                        source_faces,
                        target_faces,
                        source_index,
                        target_index,
                        temp_frame
                    )
            else:
                logging.debug('Replacing specific face(s) in the target image with specific face(s) from the source image')

                if source_indexes == "-1":
                    source_indexes = ','.join(map(lambda x: str(x), range(num_source_faces)))

                if target_indexes == "-1":
                    target_indexes = ','.join(map(lambda x: str(x), range(num_target_faces)))

                source_indexes = source_indexes.split(',')
                target_indexes = target_indexes.split(',')
                num_source_faces_to_swap = len(source_indexes)
                num_target_faces_to_swap = len(target_indexes)

                if num_source_faces_to_swap > num_source_faces:
                    raise Exception('Number of source indexes is greater than the number of faces in the source image')

                if num_target_faces_to_swap > num_target_faces:
                    raise Exception('Number of target indexes is greater than the number of faces in the target image')

                if num_source_faces_to_swap > num_target_faces_to_swap:
                    num_iterations = num_source_faces_to_swap
                else:
                    num_iterations = num_target_faces_to_swap

                if num_source_faces_to_swap == num_target_faces_to_swap:
                    for index in range(num_iterations):
                        source_index = int(source_indexes[index])
                        target_index = int(target_indexes[index])

                        if source_index > num_source_faces-1:
                            raise ValueError(f'Source index {source_index} is higher than the number of faces in the source image')

                        if target_index > num_target_faces-1:
                            raise ValueError(f'Target index {target_index} is higher than the number of faces in the target image')

                        temp_frame = swap_face(
                            face_swapper,
                            source_faces,
                            target_faces,
                            source_index,
                            target_index,
                            temp_frame
                        )
        else:
            logging.error('Unsupported face configuration')
            raise Exception('Unsupported face configuration')
        result = temp_frame
    else:
        logging.error('No target faces found')
        raise Exception('No target faces found!')

    result_image = Image.fromarray(cv2.cvtColor(result, cv2.COLOR_BGR2RGB))
    return result_image


def face_swap(src_img_path,
              target_img_path,
              source_indexes,
              target_indexes,
              background_enhance,
              face_restore,
              face_upsample,
              upscale,
              codeformer_fidelity,
              output_format):

    source_img_paths = src_img_path.split(';')
    source_img = [Image.open(img_path) for img_path in source_img_paths]
    target_img = Image.open(target_img_path)

    # download from https://huggingface.co/ashleykleynhans/inswapper/tree/main
    model = os.path.join(script_dir, 'checkpoints/inswapper_128.onnx')
    logging.debug(f'Face swap model: {model}')

    try:
        logging.debug('Performing face swap')
        result_image = process(
            source_img,
            target_img,
            source_indexes,
            target_indexes,
            model
        )
        logging.debug('Face swap complete')
    except Exception as e:
        raise

    # make sure the ckpts downloaded successfully
    check_ckpts()

    if face_restore:
        # https://huggingface.co/spaces/sczhou/CodeFormer
        logging.debug('Setting upsampler to RealESRGAN_x2plus')
        upsampler = set_realesrgan()

        if torch.cuda.is_available():
            torch_device = 'cuda'
        else:
            torch_device = 'cpu'

        logging.debug(f'Torch device: {torch_device.upper()}')
        device = torch.device(torch_device)

        codeformer_net = ARCH_REGISTRY.get('CodeFormer')(
            dim_embd=512,
            codebook_size=1024,
            n_head=8,
            n_layers=9,
            connect_list=['32', '64', '128', '256'],
        ).to(device)

        ckpt_path = os.path.join(script_dir, 'CodeFormer/CodeFormer/weights/CodeFormer/codeformer.pth')
        logging.debug(f'Loading CodeFormer model: {ckpt_path}')
        checkpoint = torch.load(ckpt_path)['params_ema']
        codeformer_net.load_state_dict(checkpoint)
        codeformer_net.eval()
        result_image = cv2.cvtColor(np.array(result_image), cv2.COLOR_RGB2BGR)
        logging.debug('Performing face restoration using CodeFormer')

        try:
            result_image = face_restoration(
                result_image,
                background_enhance,
                face_upsample,
                upscale,
                codeformer_fidelity,
                upsampler,
                codeformer_net,
                device
            )
        except Exception as e:
            raise

        logging.debug('CodeFormer face restoration completed successfully')
        result_image = Image.fromarray(result_image)

    output_buffer = io.BytesIO()
    result_image.save(output_buffer, format=output_format)
    image_data = output_buffer.getvalue()

    return base64.b64encode(image_data).decode('utf-8')


app = Flask(__name__)


@app.errorhandler(400)
def not_found(error):
    return make_response(jsonify(
        {
            'status': 'error',
            'msg': f'Bad Request',
            'detail': str(error)
        }
    ), 400)


@app.errorhandler(404)
def not_found(error):
    return make_response(jsonify(
        {
            'status': 'error',
            'msg': f'{request.url} not found',
            'detail': str(error)
        }
    ), 404)


@app.errorhandler(500)
def internal_server_error(error):
    return make_response(jsonify(
        {
            'status': 'error',
            'msg': 'Internal Server Error',
            'detail': str(error)
        }
    ), 500)


@app.route('/', methods=['GET'])
def ping():
    return make_response(jsonify(
        {
            'status': 'ok'
        }
    ), 200)


@app.route('/faceswap', methods=['POST'])
def face_swap_api():
    total_timer = Timer()
    logging.debug('Received face swap API request')
    payload = request.get_json()

    if not os.path.exists(TMP_PATH):
        logging.debug(f'Creating temporary directory: {TMP_PATH}')
        os.makedirs(TMP_PATH)

    unique_id = uuid.uuid4()
    source_image_data = payload['source_image']
    target_image_data = payload['target_image']

    # Decode the source image data
    source_image = base64.b64decode(source_image_data)
    source_file_extension = determine_file_extension(source_image_data)
    source_image_path = f'{TMP_PATH}/source_{unique_id}{source_file_extension}'

    # Save the source image to disk
    with open(source_image_path, 'wb') as source_file:
        source_file.write(source_image)

    # Decode the target image data
    target_image = base64.b64decode(target_image_data)
    target_file_extension = determine_file_extension(target_image_data)
    target_image_path = f'{TMP_PATH}/target_{unique_id}{target_file_extension}'

    # Save the target image to disk
    with open(target_image_path, 'wb') as target_file:
        target_file.write(target_image)

    # Set defaults if they are not specified in the payload
    if 'source_indexes' not in payload:
        payload['source_indexes'] = '-1'

    if 'target_indexes' not in payload:
        payload['target_indexes'] = '-1'

    if 'background_enhance' not in payload:
        payload['background_enhance'] = True

    if 'face_restore' not in payload:
        payload['face_restore'] = True

    if 'face_upsample' not in payload:
        payload['face_upsample'] = True

    if 'upscale' not in payload:
        payload['upscale'] = 1

    if 'codeformer_fidelity' not in payload:
        payload['codeformer_fidelity'] = 0.5

    if 'output_format' not in payload:
        payload['output_format'] = 'JPEG'

    try:
        logging.debug(f'Source indexes: {payload["source_indexes"]}')
        logging.debug(f'Target indexes: {payload["target_indexes"]}')
        logging.debug(f'Background enhance: {payload["background_enhance"]}')
        logging.debug(f'Face Restoration: {payload["face_restore"]}')
        logging.debug(f'Face Upsampling: {payload["face_upsample"]}')
        logging.debug(f'Upscale: {payload["upscale"]}')
        logging.debug(f'Codeformer Fidelity: {payload["codeformer_fidelity"]}')
        logging.debug(f'Output Format: {payload["output_format"]}')

        result_image = face_swap(
            source_image_path,
            target_image_path,
            payload['source_indexes'],
            payload['target_indexes'],
            payload['background_enhance'],
            payload['face_restore'],
            payload['face_upsample'],
            payload['upscale'],
            payload['codeformer_fidelity'],
            payload['output_format']
        )

        status_code = 200

        response = {
            'status': 'ok',
            'image': result_image
        }
    except Exception as e:
        logging.error(e)

        response = {
            'status': 'error',
            'msg': 'Face swap failed',
            'detail': str(e)
        }

        status_code = 500

    # Clean up temporary images
    os.remove(source_image_path)
    os.remove(target_image_path)

    total_time = total_timer.get_elapsed_time()
    logging.info(f'Total time taken for face swap API call {total_time} seconds')

    return make_response(jsonify(response), status_code)


if __name__ == '__main__':
    args = get_args()

    serve(
        app,
        host=args.host,
        port=args.port
    )