import sys
import os
import io
import gradio as gr
import cv2
import json
import requests
import numpy as np
from gradio.components import Image
from PIL import Image, ExifTags
import logging

from id_ocr.engine.header import *
from id_live.engine.header import *
import id_ocr.engine.header as id_ocr_header
import id_live.engine.header as id_live_header

css = """
.example-image img{
    display: flex; /* Use flexbox to align items */
    justify-content: center; /* Center the image horizontally */
    align-items: center; /* Center the image vertically */
    height: 300px; /* Set the height of the container */
    object-fit: contain; /* Preserve aspect ratio while fitting the image within the container */
}
.example-image img{
    display: flex; /* Use flexbox to align items */
    text-align: center;
    justify-content: center; /* Center the image horizontally */
    align-items: center; /* Center the image vertically */
    height: 350px; /* Set the height of the container */
    object-fit: contain; /* Preserve aspect ratio while fitting the image within the container */
}

.markdown-success-container {
    background-color: #F6FFED;
    padding: 20px;
    margin: 20px;
    border-radius: 1px;
    border: 2px solid green;
    text-align: center;
}

.markdown-fail-container {
    background-color: #FFF1F0;
    padding: 20px;
    margin: 20px;
    border-radius: 1px;
    border: 2px solid red;
    text-align: center;
}

.block-background {
    # background-color: #202020; /* Set your desired background color */
    border-radius: 5px;
}
"""

file_path = os.path.abspath(__file__)
root_path = os.path.dirname(file_path)

g_id_ocr_activation_result = -1
g_id_live_activation_result = -1

screenReplayThreshold = 0.5
portraitReplaceThreshold = 0.5
printedCopyThreshold = 0.5

def activate_id_ocr_sdk():
    id_ocr_key = os.environ.get("LICENSE_KEY")
    id_ocr_dict_path = os.path.join(root_path, "id_ocr/engine/bin")

    ret = -1
    if id_ocr_key is None:
        print_warning("ID OCR online license key not found!")
        return ret
    else:
        activate_ret = id_ocr_header.set_activation(id_ocr_key.encode('utf-8')).decode('utf-8')
        ret = json.loads(activate_ret).get("errorCode", None)

    if ret == 0:
        init_ret = id_ocr_header.init_sdk(id_ocr_dict_path.encode('utf-8')).decode('utf-8')
        ret = json.loads(init_ret).get("errorCode", None)
        if ret == 0:
            print_log("Successfully init ID OCR SDK!")
        else:
            print_log("Failed to init ID OCR SDK!")
    else:
        print_error(f"Falied to activate ID OCR SDK, Error code {ret}")

    return ret

def activate_id_live_sdk():
    print("id live sdk activation start")
    sys.stdout.flush()
    id_live_key = os.environ.get("ID_LIVE_LICENSE_KEY")
    print(id_live_key)
    id_live_dict_path = os.path.join(root_path, "id_live/engine/model")
    print(id_live_dict_path)
    ret = -1
    if id_live_key is None:
        print_warning("ID LIVE license key not found!")
        return ret
    else:
        ret = id_live_header.set_activation(id_live_key.encode('utf-8'))

    print(f"activation result   {ret}")
    if ret == 0:
        ret = id_live_header.init_sdk(id_live_dict_path.encode('utf-8'))
        print(f"init result   {ret}")
        
        if ret == 0:
            print_log("Successfully init ID LIVE SDK!")
        else:
            print_log("Failed to init ID LIVE SDK!")
    else:
        print_error(f"Falied to activate ID LIVE SDK, Error code {ret}")

    return ret

def find_key_in_dict(d, target_key):
  for key, value in d.items():
    if key == target_key:
      return value
    elif isinstance(value, dict):  # If the value is a dictionary, search recursively
      result = find_key_in_dict(value, target_key)
      if result is not None:
        return result
  return None

def json_to_html_table(data, image_keys):
  html = "<table border='1' style='border-collapse: collapse; width: 100%;'>"
  for key, value in data.items():
    if isinstance(value, dict):
      html += f"<tr><td colspan='2'><strong>{key}</strong></td></tr>"
      for sub_key, sub_value in value.items():
        if sub_key in image_keys:
          html += f"<tr><td>{sub_key}</td><td><img src='data:image/png;base64,{sub_value}' width = '200'  height= '100' /></td></tr>"
        else:
          html += f"<tr><td>{sub_key}</td><td>{sub_value}</td></tr>"
    else:
      if key in image_keys:
        html += f"<tr><td>{key}</td><td><img src='data:image/png;base64,{value}' width = '200'  height= '100' /></td></tr>"
      else:
        html += f"<tr><td>{key}</td><td>{value}</td></tr>"
      
  html += "</table>"
  return html

def apply_exif_rotation(image):
    try:
        exif = image._getexif()
        if exif is not None:
            for orientation in ExifTags.TAGS.keys():
                if ExifTags.TAGS[orientation] == 'Orientation':
                    break

            # Get the orientation value
            orientation = exif.get(orientation, None)

            # Apply the appropriate rotation based on the orientation
            if orientation == 3:
                image = image.rotate(180, expand=True)
            elif orientation == 6:
                image = image.rotate(270, expand=True)
            elif orientation == 8:
                image = image.rotate(90, expand=True)

    except AttributeError:
        print("No EXIF data found")

    return image

def check_liveness(frame):
  # if frame is None:
  #   liveness_result = f"""<div class="markdown-fail-container"><p style="text-align: center; font-size: 20px; color: red;">Liveness Check Failed</p></div>"""
  #   return [liveness_result, {"status": "error", "result": "select image file!"}]
  
  # img_bytes = io.BytesIO()
  # Image.open(frame).save(img_bytes, format="JPEG")
  # img_bytes.seek(0)    
  
  # url = "https://recognito-iddocumentlivenessdetection.p.rapidapi.com/process_image"
  
  # try:
  #   files = {'image': img_bytes}
  #   headers = {"X-RapidAPI-Key": os.environ.get("API_KEY")}

  #   result = requests.post(url=url, files=files, headers=headers)
  # except:
  #   liveness_result = f"""<div class="markdown-fail-container"><p style="text-align: center; font-size: 20px; color: red;">Liveness Check Failed</p></div>"""
  #   return [liveness_result, {"status": "error", "result": "failed to open file!"}]

  # if result.ok:
  #   json_result = result.json()
  #   if json_result.get("resultCode") == "Error":
  #     liveness_result = f"""<div class="markdown-fail-container"><p style="text-align: center; font-size: 20px; color: red;">Liveness Check Failed</p></div>"""
  #     return [liveness_result, {"status": "error", "result": "server error!"}]

  #   process_results = json_result.get("result")
  #   status = process_results.get("status")
  #   if status == "Ok":
  #     screenReply = process_results.get("screenReply")
  #     portraitReplace = process_results.get("portraitReplace")
  #     printedCopy = process_results.get("printedCopy")
  #     liveness_result = f"""<div class="markdown-success-container"><p style="text-align: center; font-size: 20px; color: green;">Liveness Check:  GENUINE</p></div>"""

  #     # Check for "Spoof" condition
  #     if screenReply < screenReplayThreshold or portraitReplace < portraitReplaceThreshold or printedCopy < printedCopyThreshold:
  #       liveness_result = f"""<div class="markdown-fail-container"><p style="text-align: center; font-size: 20px; color: red;">Liveness Check:  SPOOF</p></div>"""

  #     json_output = {"Screen Replay Check": "Failed" if screenReply < screenReplayThreshold else "Success", 
  #                    "Portrait Replace Check": "Failed" if portraitReplace < portraitReplaceThreshold else "Success",
  #                    "Printed Cutout Check": "Failed" if printedCopy < printedCopyThreshold else "Success"}
  #     # Update json_result with the modified process_results
  #     return [liveness_result, json_output]
    
  #   liveness_result = f"""<div class="markdown-fail-container"><p style="text-align: center; font-size: 20px; color: red;">Liveness Check Failed</p></div>"""
  #   return [liveness_result, {"status": "error", "result": "document not found!"}]
  # else:
  #   liveness_result = f"""<div class="markdown-fail-container"><p style="text-align: center; font-size: 20px; color: red;">Liveness Check Failed</p></div>"""
  #   return [liveness_result, {"status": "error", "result": f"{result.text}"}]
    global g_id_live_activation_result
    if g_id_live_activation_result != 0:
        gr.Warning("ID LIVE SDK Activation Failed!")
        return {"status": "error", "result": "activation error!"}
    try:
        image = Image.open(frame)
        image = apply_exif_rotation(image.convert('RGB'))
    except:
        raise gr.Error("Please select image file!")
    
    image_np = np.asarray(image)
    result = id_live_header.processImage(image_np, image_np.shape[1], image_np.shape[0])
    result_dict = json.loads(result.decode('utf-8'))
    status = result_dict["status"]
    if status == "Ok":
        screenReply = float(result_dict["screenReply"])
        portraitReplace = float(result_dict["portraitReplace"])
        printedCopy = float(result_dict["printedCopy"])
        detResult = "genuine"

        # Check for "Spoof" condition
        if screenReply < screenReplayThreshold or portraitReplace < portraitReplaceThreshold or printedCopy < printedCopyThreshold:
            detResult = "spoof"

        # Update json_result with the modified process_results
        return {"status": "ok", "data": {"result": detResult, "screenreplay_integrity_score": screenReply, "portraitreplace_integrity_score": portraitReplace, "printedcutout_integrity_score": printedCopy}}
            
    return {"status": "error", "result": "document not found!"}

    
def idcard_recognition(frame1, frame2):
  # url = "https://recognito-iddocumentrecognition.p.rapidapi.com/api/read_idcard"
  
  # files = None
  # if frame1 is not None and frame2 is not None:
  #   files = {'image': open(frame1, 'rb'), 'image2': open(frame2, 'rb')}
  # elif frame1 is not None and frame2 is None:
  #   files = {'image': open(frame1, 'rb')}
  # elif frame1 is None and frame2 is not None:
  #   files = {'image': open(frame2, 'rb')}
  # else:
  #   return ['', None, None]
  
  # headers = {"X-RapidAPI-Key": os.environ.get("API_KEY")}
    
  # r = requests.post(url=url, files=files, headers=headers)

  # images = None   
  # rawValues = {}
  # image_table_value = ""
  # result_table_dict = {
  #   'portrait':'', 
  #   'documentName':'',
  #   'score':'',
  #   'countryName':'',
  #   'name':'', 
  #   'sex':'', 
  #   'address':'', 
  #   'dateOfBirth':'', 
  #   'dateOfIssue':'', 
  #   'dateOfExpiry':'',
  #   'documentNumber':'',
  #   }

  # if 'data' in r.json():
  #     for key, value in r.json()['data'].items():
  #       if key == 'image':
  #         for image_key, image_value in value.items():
  #           row_value = ("<tr>"
  #                         "<td>{key}</td>"
  #                         "<td><img src=""data:image/png;base64,{base64_image} width = '200'  height= '100' /></td>"
  #                       "</tr>".format(key=image_key, base64_image=image_value))
  #           image_table_value = image_table_value + row_value

  # images = ("<table>"
  #           "<tr>"
  #             "<th>Field</th>"
  #             "<th>Image</th>"
  #           "</tr>"
  #           "{image_table_value}"
  #           "</table>".format(image_table_value=image_table_value))

  # for key, value in r.json().items():
  #   if key == 'data':
  #     if 'image' in value:
  #       del value['image']
  #     rawValues[key] = value
  #   else:
  #     rawValues[key] = value

  
  # for result_key in result_table_dict.keys():
  #   result_table_dict[result_key] =  find_key_in_dict(r.json(), result_key)
  
  # result = json_to_html_table(result_table_dict, {'portrait'})
  # json_result = json.dumps(rawValues, indent=6)
  # return [result, json_result, images]

  global g_id_ocr_activation_result
  if g_id_ocr_activation_result != 0:
    gr.Warning("ID OCR SDK Activation Failed!")
    return ['', None]
  
  image1, image2 = '', ''
  if frame1 is not None and frame2 is not None:
    image1, image2 = frame1, frame2
  elif frame1 is not None and frame2 is None:
    image1 = frame1
  elif frame1 is None and frame2 is not None:
    image1 = frame2
  elif frame1 is None and frame2 is None:
    raise gr.Error("Please select images files!")

  ocrResult = id_ocr_header.ocr_id_card(image1.encode('utf-8'), image2.encode('utf-8'))
  ocrResDict = json.loads(ocrResult)  
  images = None   
  rawValues = {}
  image_table_value = ""
  result_table_dict = {
    'portrait':'', 
    'documentName':'',
    'score':'',
    'countryName':'',
    'name':'', 
    'sex':'', 
    'address':'', 
    'dateOfBirth':'', 
    'dateOfIssue':'', 
    'dateOfExpiry':'',
    'documentNumber':'',
    }

  for key, value in ocrResDict.items():
    if key == 'image':
      for image_key, image_value in value.items():
        row_value = ("<tr>"
                      "<td>{key}</td>"
                      "<td><img src=""data:image/png;base64,{base64_image} width = '200'  height= '100' /></td>"
                    "</tr>".format(key=image_key, base64_image=image_value))
        image_table_value = image_table_value + row_value

  images = ("<table>"
            "<tr>"
              "<th>Field</th>"
              "<th>Image</th>"
            "</tr>"
            "{image_table_value}"
            "</table>".format(image_table_value=image_table_value))

  value = ocrResDict.copy()
  if 'image' in value:
    del value['image']
    rawValues = value
  else:
    rawValues = value

  
  for result_key in result_table_dict.keys():
    result_table_dict[result_key] =  find_key_in_dict(ocrResDict, result_key)
  
  result = json_to_html_table(result_table_dict, {'portrait'})
  json_result = json.dumps(rawValues, indent=6)
  return [result, json_result, images]

def launch_demo():
  with gr.Blocks(css=css) as demo:
    gr.Markdown(
      f"""
      <a href="https://recognito.vision" style="display: flex; align-items: center;">
        <img src="https://recognito.vision/wp-content/uploads/2024/03/Recognito-modified.png" style="width: 8%; margin-right: 15px;"/>
        <div>
          <p style="font-size: 32px; font-weight: bold; margin: 0;">Recognito</p>
          <p style="font-size: 18px; margin: 0;">www.recognito.vision</p>
        </div>
      </a>
      <p style="font-size: 20px; font-weight: bold;">📘 Product Documentation</p>
      <div style="display: flex; align-items: center;">          
        &emsp;&emsp;<a href="https://docs.recognito.vision" style="display: flex; align-items: center;"><img src="https://recognito.vision/wp-content/uploads/2024/05/book.png" style="width: 48px; margin-right: 5px;"/></a>
      </div>
      <p style="font-size: 20px; font-weight: bold;">🏠 Visit Recognito</p>
      <div style="display: flex; align-items: center;">
        &emsp;&emsp;<a href="https://recognito.vision" style="display: flex; align-items: center;"><img src="https://recognito.vision/wp-content/uploads/2024/03/recognito_64_cl.png" style="width: 32px; margin-right: 5px;"/></a>
        &nbsp;&nbsp;&nbsp;&nbsp;<a href="https://www.linkedin.com/company/recognito-vision" style="display: flex; align-items: center;"><img src="https://recognito.vision/wp-content/uploads/2024/03/linkedin_64_cl.png" style="width: 32px; margin-right: 5px;"/></a>
        &nbsp;&nbsp;&nbsp;&nbsp;<a href="https://huggingface.co/recognito" style="display: flex; align-items: center;"><img src="https://recognito.vision/wp-content/uploads/2024/03/hf_64_cl.png" style="width: 32px; margin-right: 5px;"/></a>
        &nbsp;&nbsp;&nbsp;&nbsp;<a href="https://github.com/recognito-vision" style="display: flex; align-items: center;"><img src="https://recognito.vision/wp-content/uploads/2024/03/github_64_cl.png" style="width: 32px; margin-right: 5px;"/></a>
        &nbsp;&nbsp;&nbsp;&nbsp;<a href="https://hub.docker.com/u/recognito" style="display: flex; align-items: center;"><img src="https://recognito.vision/wp-content/uploads/2024/03/docker_64_cl.png" style="width: 32px; margin-right: 5px;"/></a>
        &nbsp;&nbsp;&nbsp;&nbsp;<a href="https://www.youtube.com/@recognito-vision" style="display: flex; align-items: center;"><img src="https://recognito.vision/wp-content/uploads/2024/04/youtube_64_cl.png" style="width: 32px; margin-right: 5px;"/></a>
      </div>
      <p style="font-size: 20px; font-weight: bold;">🤝 Contact us for our on-premise ID Document Verification SDKs deployment</p>
      <div style="display: flex; align-items: center;">
        &emsp;&emsp;<a target="_blank" href="mailto:hello@recognito.vision"><img src="https://img.shields.io/badge/email-hassan@recognito.vision-blue.svg?logo=gmail " alt="www.recognito.vision"></a>
        &nbsp;&nbsp;&nbsp;&nbsp;<a target="_blank" href="https://wa.me/+14158003112"><img src="https://img.shields.io/badge/whatsapp-+14158003112-blue.svg?logo=whatsapp " alt="www.recognito.vision"></a>
        &nbsp;&nbsp;&nbsp;&nbsp;<a target="_blank" href="https://t.me/recognito_vision"><img src="https://img.shields.io/badge/telegram-@recognito__vision-blue.svg?logo=telegram " alt="www.recognito.vision"></a>
        &nbsp;&nbsp;&nbsp;&nbsp;<a target="_blank" href="https://join.slack.com/t/recognito-workspace/shared_invite/zt-2d4kscqgn-"><img src="https://img.shields.io/badge/slack-recognito__workspace-blue.svg?logo=slack " alt="www.recognito.vision"></a>
      </div>
      <br/>
      """
    )

    with gr.Tabs():
      with gr.Tab("ID Document Recognition"):
        with gr.Row():
          with gr.Column(scale=6):
            with gr.Row():
              with gr.Column(scale=3):
                id_image_input1 = gr.Image(type='filepath', label='Front', elem_classes="example-image")
              with gr.Column(scale=3):
                id_image_input2 = gr.Image(type='filepath', label='Back', elem_classes="example-image")
            
            with gr.Row():
              id_examples = gr.Examples(
                examples=[['examples/1_f.png', 'examples/1_b.png'], 
                          ['examples/2_f.png', 'examples/2_b.png'], 
                          ['examples/3_f.png', 'examples/3_b.png'], 
                          ['examples/4.png', None]],
                inputs=[id_image_input1, id_image_input2],
                outputs=None,
                fn=idcard_recognition
              )
            
          with gr.Blocks():
            with gr.Column(scale=4, min_width=400, elem_classes="block-background"):     
              id_recognition_button = gr.Button("ID Card Recognition", variant="primary", size="lg")

              with gr.Tab("Key Fields"):
                id_result_output = gr.HTML()
              with gr.Tab("Raw JSON"):
                json_result_output = gr.JSON()
              with gr.Tab("Images"):
                image_result_output = gr.HTML()
              
              id_recognition_button.click(idcard_recognition, inputs=[id_image_input1, id_image_input2], outputs=[id_result_output, json_result_output, image_result_output])

      with gr.Tab("ID Document Liveness Detection"):
        with gr.Row():
          with gr.Column(scale=1):
            id_image_input = gr.Image(label="Image", type='filepath', elem_classes="example-image")
            gr.Examples(examples=['examples/1_f.png', 'examples/2_f.png', 'examples/3_f.png', 'examples/4.png'], inputs=id_image_input)

          with gr.Blocks():
            with gr.Column(scale=1, elem_classes="block-background"):     
              check_liveness_button = gr.Button("Check Document Liveness", variant="primary", size="lg")
              
              liveness_result = gr.Markdown("")
              json_output = gr.JSON()
              
              check_liveness_button.click(check_liveness, inputs=id_image_input, outputs=[liveness_result, json_output])
    
    gr.HTML('<a href="https://visitorbadge.io/status?path=https%3A%2F%2Fhuggingface.co%2Fspaces%2Frecognito%2FID-Document-Verification"><img src="https://api.visitorbadge.io/api/combined?path=https%3A%2F%2Fhuggingface.co%2Fspaces%2Frecognito%2FID-Document-Verification&labelColor=%2337d67a&countColor=%23263759&style=flat" /></a>')
    
  demo.launch(server_name="0.0.0.0", server_port=7860, show_api=False)

if __name__ == '__main__':
  logging.info("This is an info log")
  g_id_ocr_activation_result = activate_id_ocr_sdk()
  g_id_live_activation_result = activate_id_live_sdk()
  launch_demo()