import base64
import logging
import tempfile
from collections import defaultdict

import cv2
import numpy as np
import requests
from tenacity import retry, stop_after_attempt, wait_random

from paddle_detection import PADDLE_DET
from paddle_detection.deploy.third_engine.onnx.infer import PredictConfig
from paddle_detection.deploy.third_engine.onnx.preprocess import Compose
from util import image_util, util


def predict_image(infer_config, predictor, img_path):
    # load preprocess transforms
    transforms = Compose(infer_config.preprocess_infos)
    # predict image
    inputs = transforms(img_path)
    inputs["image"] = np.array(inputs["image"]).astype('float32')
    inputs_name = [var.name for var in predictor.get_inputs()]
    inputs = {k: inputs[k][None,] for k in inputs_name}

    outputs = predictor.run(output_names=None, input_feed=inputs)

    bboxes = np.array(outputs[0])
    result = defaultdict(list)
    for bbox in bboxes:
        if bbox[0] > -1 and bbox[1] > infer_config.draw_threshold:
            result[bbox[0]].append({"score": bbox[1], "box": bbox[2:]})
    return result


def detect_image(img_path):
    infer_cfg = "model/object_det_model/infer_cfg.yml"
    # load infer config
    infer_config = PredictConfig(infer_cfg)

    return predict_image(infer_config, PADDLE_DET, img_path)


def get_book_areas(image):
    with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
        cv2.imwrite(temp_file.name, image)
    detect_result = detect_image(temp_file.name)
    util.delete_temp_file(temp_file.name)
    book_areas = detect_result[73]
    result = []
    for book_area in book_areas:
        result.append(image_util.capture(image, book_area["box"]))
    return result


@retry(stop=stop_after_attempt(3), wait=wait_random(1, 3), reraise=True,
       after=lambda x: logging.warning("获取文档区域失败！"))
def request_book_areas(image):
    url = "http://det_api:5000/det/detect_books"
    _, encoded_image = cv2.imencode('.jpg', image)
    byte_stream = encoded_image.tobytes()
    files = {"image": ("image.jpg", byte_stream)}
    response = requests.post(url, files=files)
    if response.status_code == 200:
        img_str_list = response.json()
        result = []
        for img_str in img_str_list:
            img_data = base64.b64decode(img_str)
            np_array = np.frombuffer(img_data, np.uint8)
            result.append(cv2.imdecode(np_array, cv2.IMREAD_COLOR))
        return result
    else:
        return []