import os
import cv2
import torch
import base64
import argparse
import numpy as np
import insightface
import onnxruntime
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

from face_swapper import Inswapper, paste_to_whole
from face_analyser import get_analysed_data
from face_parsing import init_parsing_model, get_parsed_mask


# Глобальные константы и переменные
USE_COLAB = user_args.colab
USE_CUDA = user_args.cuda
PROVIDER = ["CPUExecutionProvider"]
DETECT_SIZE = 640
DETECT_THRESH = 0.6
FACE_ANALYSER = None
FACE_SWAPPER = None
FACE_PARSER = None

## ------------------------------ USER ARGS ------------------------------

parser = argparse.ArgumentParser(description="Swap-Mukham Face Swapper")
parser.add_argument("--out_dir", help="Default Output directory", default=os.getcwd())
parser.add_argument("--batch_size", help="Gpu batch size", default=32)
parser.add_argument("--cuda", action="store_true", help="Enable cuda", default=False)
parser.add_argument(
    "--colab", action="store_true", help="Enable colab mode", default=False
)
user_args = parser.parse_args()

## ------------------------------ SET EXECUTION PROVIDER ------------------------------
# Note: Non CUDA users may change settings here

PROVIDER = ["CPUExecutionProvider"]

if USE_CUDA:
    available_providers = onnxruntime.get_available_providers()
    if "CUDAExecutionProvider" in available_providers:
        print("\n********** Running on CUDA **********\n")
        PROVIDER = ["CUDAExecutionProvider", "CPUExecutionProvider"]
    else:
        USE_CUDA = False
        print("\n********** CUDA unavailable running on CPU **********\n")
else:
    USE_CUDA = False
    print("\n********** Running on CPU **********\n")

device = "cuda" if USE_CUDA else "cpu"
EMPTY_CACHE = lambda: torch.cuda.empty_cache() if device == "cuda" else None


# Функции загрузки моделей
def load_face_analyser_model(name="buffalo_l"):
    global FACE_ANALYSER
    if FACE_ANALYSER is None:
        FACE_ANALYSER = insightface.app.FaceAnalysis(name=name, providers=PROVIDER)
        FACE_ANALYSER.prepare(
            ctx_id=0, det_size=(DETECT_SIZE, DETECT_SIZE), det_thresh=DETECT_THRESH
        )

def load_face_swapper_model(path="./assets/pretrained_models/inswapper_128.onnx"):
    global FACE_SWAPPER
    if FACE_SWAPPER is None:
        FACE_SWAPPER = Inswapper(model_file=path, batch_size=1, providers=PROVIDER)

def load_face_parser_model(path="./assets/pretrained_models/79999_iter.pth"):
    global FACE_PARSER
    if FACE_PARSER is None:
        FACE_PARSER = init_parsing_model(path, device=device)

# Загрузка всех моделей
load_face_analyser_model()
load_face_swapper_model()
load_face_parser_model()

def base64_to_image(base64_string):
    img_data = base64.b64decode(base64_string)
    nparr = np.frombuffer(img_data, np.uint8)
    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
    return img

def image_to_base64(image):
    _, buffer = cv2.imencode('.png', image)
    return base64.b64encode(buffer).decode('utf-8')

def process_images(source_img_base64, target_img_base64):
    # Декодирование base64 в изображения
    source_img = base64_to_image(source_img_base64)
    target_img = base64_to_image(target_img_base64)
    
    # Анализ лиц
    analysed_targets, analysed_sources, _, _ = get_analysed_data(
        FACE_ANALYSER,
        [target_img],
        source_img,
        swap_condition="First",
        detect_condition="best detection"
    )

    # Замена лица
    preds = []
    matrs = []
    for batch_pred, batch_matr in FACE_SWAPPER.batch_forward([target_img], analysed_targets, analysed_sources):
        preds.extend(batch_pred)
        matrs.extend(batch_matr)

    # Парсинг лица и создание маски
    masks = get_parsed_mask(FACE_PARSER, preds, classes=["skin", "l_brow", "r_brow", "l_eye", "r_eye", "nose", "u_lip", "l_lip", "mouth"], device='cpu')

    # Наложение результата обратно на изображение
    result_img = paste_to_whole(preds[0], target_img, matrs[0], mask=masks[0])

    # Кодирование результата в base64
    result_base64 = image_to_base64(result_img)

    return result_base64

# Создание FastAPI приложения
app = FastAPI(title="Faceswap API", description="API для замены лица. Отправьте два изображения в формате base64.")

# Определение модели запроса
class SwapRequest(BaseModel):
    source_img: str
    target_img: str

@app.post("/swap_face")
async def swap_face(request: SwapRequest):
    try:
        result = process_images(request.source_img, request.target_img)
        return {"result": result}
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

# Запуск сервера
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)