import argparse
import cv2
import gradio as gr
import json
import logging
import nh3
import numpy as np
import os
import re
import sys
import torch
import torch.nn.functional as F
from fastapi import FastAPI, File, UploadFile, Request
from fastapi.responses import HTMLResponse, RedirectResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from transformers import AutoTokenizer, BitsAndBytesConfig, CLIPImageProcessor
from typing import Callable

from model.LISA import LISAForCausalLM
from model.llava import conversation as conversation_lib
from model.llava.mm_utils import tokenizer_image_token
from model.segment_anything.utils.transforms import ResizeLongestSide
from utils import session_logger
from utils.utils import (DEFAULT_IM_END_TOKEN, DEFAULT_IM_START_TOKEN,
                         DEFAULT_IMAGE_TOKEN, IMAGE_TOKEN_INDEX)

session_logger.change_logging(logging.DEBUG)


CUSTOM_GRADIO_PATH = "/"
app = FastAPI(title="lisa_app", version="1.0")

FASTAPI_STATIC = os.getenv("FASTAPI_STATIC")
os.makedirs(FASTAPI_STATIC, exist_ok=True)
app.mount("/static", StaticFiles(directory=FASTAPI_STATIC), name="static")
templates = Jinja2Templates(directory="templates")


@app.get("/health")
@session_logger.set_uuid_logging
def health() -> str:
    try:
        logging.info("health check")
        return json.dumps({"msg": "ok"})
    except Exception as e:
        logging.error(f"exception:{e}.")
        return json.dumps({"msg": "request failed"})


@session_logger.set_uuid_logging
def get_cleaned_input(input_str):
    logging.info(f"start cleaning of input_str: {input_str}.")
    input_str = nh3.clean(
        input_str,
        tags={
            "a",
            "abbr",
            "acronym",
            "b",
            "blockquote",
            "code",
            "em",
            "i",
            "li",
            "ol",
            "strong",
            "ul",
        },
        attributes={
            "a": {"href", "title"},
            "abbr": {"title"},
            "acronym": {"title"},
        },
        url_schemes={"http", "https", "mailto"},
        link_rel=None,
    )
    logging.info(f"cleaned input_str: {input_str}.")
    return input_str


@session_logger.set_uuid_logging
def get_inference_model_by_args(args_to_parse):
    logging.info(f"args_to_parse:{args_to_parse}.")

    @session_logger.set_uuid_logging
    def inference(input_str, input_image):
        ## filter out special chars

        input_str = get_cleaned_input(input_str)
        logging.info(f"input_str type: {type(input_str)}, input_image type: {type(input_image)}.")
        logging.info(f"input_str: {input_str}.")

        return output_image, output_str

    return inference


@session_logger.set_uuid_logging
def get_gradio_interface(fn_inference: Callable):
    return gr.Interface(
        fn_inference,
    inputs=[
            gr.Textbox(lines=1, placeholder=None, label="Text Instruction"),
            gr.Image(type="filepath", label="Input Image")
    ],
    outputs=[
            gr.Image(type="pil", label="Segmentation Output"),
        gr.Textbox(lines=1, placeholder=None, label="Text Output"),
    ],
        title=title,
        description=description,
        article=article,
        examples=examples,
        allow_flagging="auto",
)


args = parse_args(sys.argv[1:])
inference_fn = get_inference_model_by_args(args)
io = get_gradio_interface(inference_fn)
app = gr.mount_gradio_app(app, io, path=CUSTOM_GRADIO_PATH)