import gradio as gr

from typing import List, Union
from pathlib import Path

from lib.platform import PM
from lib.info.log import get_logger
from .backend import HSMRBackend

class HSMRService:
    # ===== Initialization Part =====

    def __init__(self, backend:HSMRBackend) -> None:
        self.example_imgs_root = PM.inputs / 'example_imgs'
        self.description = self._load_description()
        self.backend:HSMRBackend = backend

    def _load_description(self) -> str:
        description_fn = PM.inputs / 'description.md'
        with open(description_fn, 'r') as f:
            description = f.read()
        return description


    # ===== Funcitonal Private Part =====

    def _inference_img(
        self,
        raw_img_path:Union[str, Path],
        max_instances:int = 5,
    ) -> List:
        get_logger(brief=True).info(f'Image Path: {raw_img_path}')
        get_logger(brief=True).info(f'max_instances: {max_instances}')
        if raw_img_path is None:
            gr.Warning('No image uploaded yet. Please upload an image first.')
            return '<NONE>'
        if isinstance(raw_img_path, str):
            raw_img_path = Path(raw_img_path)

        args = {
                'max_instances': max_instances,
                'rec_bs': 1,
            }
        output_paths = self.backend(raw_img_path, args)

        # bbx_img_path  = output_paths['bbx_img_path']
        # mesh_img_path = output_paths['mesh_img_path']
        # skel_img_path = output_paths['skel_img_path']
        blend_img_path  = output_paths['front_blend']

        return blend_img_path


    # ===== Service Part =====

    def serve(self) -> None:
        ''' Build UI and set up the service. '''
        with gr.Blocks() as demo:
            # 1a. Setup UI.
            gr.Markdown(self.description)

            with gr.Tab(label='HSMR-IMG'):
                gr.Markdown('> Demo for recoverying human mesh and skeleton from a single image. (For **Pure CPU** demo, inference may take **about 1~3 minutes**.)')
                with gr.Row(equal_height=False):
                    with gr.Column():
                        input_image = gr.Image(
                                label = 'Input',
                                type  = 'filepath',
                            )
                        with gr.Row(equal_height=True):
                            run_button_image = gr.Button(
                                    value     = 'Inference',
                                    variant   = 'primary',
                                )

                    with gr.Column():
                        output_blend = gr.Image(
                                label       = 'Output',
                                type        = 'filepath',
                                interactive = False,
                            )

            # 1b. Add examples sections after setting I/O policy.
            example_fns = sorted(self.example_imgs_root.glob('*'))
            gr.Examples(
                examples       = example_fns,
                fn             = self._inference_img,
                inputs         = input_image,
                outputs        = output_blend,
            )

            # 2b. Continue binding I/O logic.
            run_button_image.click(
                fn      = self._inference_img,
                inputs  = input_image,
                outputs = output_blend,
            )


            # 3. Launch the service.
            demo.queue(max_size=20).launch(server_name='0.0.0.0', server_port=7860)