import gradio as gr
from PIL import Image as PIL_Image # Renaming to avoid conflict with Image from gltflib
from io import BytesIO
import json
import os

import src.constants as constants

from src.front_card_image import create_card_image
from src.front_card_image_historic_site import create_historic_site_card_image
from src.picture_box_model import create_picture_box_model
from src.extracted_objects_model import create_extracted_objects_model
from src.card_model import create_card_model
from src.manhole_model import create_manhole_model

language = os.environ.get('LANGUAGE', 'en')
print(f'LANGUAGE: {language}')

with open(f'src/display_message_{language}.json', 'r', encoding='utf-8') as f:
    display_message_dict = json.load(f)
    model_type_dict = display_message_dict['model_type_dict']
    color_dict = display_message_dict['color_dict']

def create_3dmodel(model_no, title, color, mark, historic_site_type, difficulty, description, is_thick, image):

    img_bytearray = BytesIO()
    image['background'].save(img_bytearray, "JPEG", quality=95)
    img_bytearray.seek(0) # Seek to the beginning of the image, otherwise it results in empty data.

    option_dict = {
        'タイトル': title,
        '色': color,
        'マーク': mark,
        '史跡種類': historic_site_type,
        '訪問難度': difficulty,
        '説明文': description,
        '厚み': '有' if is_thick else '無',
        'language': language
    }

    # Consider implementing model_no categorization rules or creating a factory class
    if model_no not in ['A', 'B']:
        # Create the card image (front side)
        if model_no == '1':
            front_img_bytearray = create_historic_site_card_image(img_bytearray, option_dict)
        else:
            front_img_bytearray = create_card_image(model_no, img_bytearray, option_dict)

        # Retrieve the card image (back side)
        back_path = constants.back_card_img_dict[model_no]
        back_img = PIL_Image.open(back_path)
        back_img_bytearray = BytesIO()
        back_img.convert('RGB').save(back_img_bytearray, "JPEG", quality=95)
        back_img_bytearray.seek(0) # Seek to the beginning of the image, otherwise it results in empty data

        # Create a 3D model of the card (return value is the path of the created model)
        model_path = create_card_model(front_img_bytearray, back_img_bytearray, option_dict)
        
    else:
        # Create a 3D model of the card (return value is the path of the created model)
        if model_no == 'A':
            model_path = create_picture_box_model(img_bytearray)
        if model_no == 'B':
            model_path = create_extracted_objects_model(img_bytearray)
    
    return model_path

def create_manhole(image):
    img_bytearray = BytesIO()
    image['background'].save(img_bytearray, "JPEG", quality=95)
    img_bytearray.seek(0) # Seek to the beginning of the image, otherwise it results in empty data.
    
    model_path = create_manhole_model(img_bytearray)
    
    return model_path

with gr.Blocks() as demo:

    gr.Markdown(display_message_dict['header'])
    
    with gr.Tab(display_message_dict['tab_label_card_general']):
        with gr.Row():
            with gr.Column():
                title = gr.Textbox(label=display_message_dict['label_title'], placeholder=display_message_dict['placeholder_title'])
                model_type = gr.Radio([(model_type_dict[key], key) for key in model_type_dict], value='2', label=display_message_dict['label_card_type'])
            with gr.Column():
                description = gr.Textbox(lines=2, label=display_message_dict['label_description'], placeholder=display_message_dict['placeholder_description'])
                is_thick = gr.Checkbox(label=display_message_dict['label_is_thick'], value=False, info=display_message_dict['info_is_thick'])
                image = gr.ImageEditor(image_mode='RGB', sources="upload", type="pil", crop_size="1:1", label=display_message_dict['label_image'])

        button = gr.Button(display_message_dict['label_button'])
        button.click(
                fn=lambda model_type, title, description, is_thick, image:
                    create_3dmodel(model_type, title, None, None, None, None, description, is_thick, image),
            inputs=[model_type, title, description, is_thick, image],
            outputs=[gr.Model3D(camera_position=(90, 90, 5))]
        )
    
    with gr.Tab(display_message_dict['tab_label_historic_site_card']):
        with gr.Row():
            with gr.Column():
                title = gr.Textbox(label=display_message_dict['label_title'], placeholder=display_message_dict['placeholder_title'])
                color = gr.Radio([(color_dict[key], key) for key in color_dict], value=list(color_dict)[0], label=display_message_dict['label_color'])
                
                # Mark(Designation Type) is only visible for Japanese(Because there are no Designation Type Images for other languages).
                # In other languages, mark is only 'No Designation'(No Designation Type Image).
                is_mark_visible = True if language == 'ja' else False
                mark = gr.Radio(display_message_dict['mark_list'], visible=is_mark_visible,
                    value=list(display_message_dict['mark_list'])[len(display_message_dict['mark_list'])-1], label=display_message_dict['label_mark'])
                
                historic_site_type = gr.Textbox(label=display_message_dict["label_historic_site_type"], placeholder=display_message_dict["placeholder_historic_site_type"])
            with gr.Column():
                difficulty = gr.Slider(1, 5, 3, step=1, label=display_message_dict['label_difficulty'])
                description = gr.Textbox(lines=2, label=display_message_dict['label_description'], placeholder=display_message_dict['placeholder_description'])
                is_thick = gr.Checkbox(label=display_message_dict['label_is_thick'], value=False, info=display_message_dict['info_is_thick'])
                image = gr.ImageEditor(image_mode='RGB', sources="upload", type="pil", crop_size="1:1", label=display_message_dict['label_image'])

        button = gr.Button(display_message_dict['label_button'])
        button.click(
                fn=lambda title, color, mark, historic_site_type, difficulty, description, is_thick, image:
                    create_3dmodel('1', title, color, mark, historic_site_type, difficulty, description, is_thick, image),
            inputs=[title, color, mark, historic_site_type, difficulty, description, is_thick, image],
            outputs=[gr.Model3D(camera_position=(90, 90, 5))]
        )

        gr.Markdown(display_message_dict['footer_historic_site_card'])
        
    with gr.Tab(display_message_dict["tab_label_manhole"]):
        image = gr.ImageEditor(image_mode='RGB', sources="upload", type="pil", label=display_message_dict['label_image'])
        button = gr.Button(display_message_dict['label_button'])
        button.click(
            create_manhole,
            inputs=[image],
            outputs=[gr.Model3D(camera_position=(90, 90, 5))]
        )

demo.launch()