import html
import json
import os
import socket
import sys
import tempfile
import threading
import traceback
from pathlib import Path

import gradio as gr
import numpy as np
from huggingface_hub import hf_hub_download
from PIL import Image

try:
    import spaces
except ImportError:
    class _SpacesFallback:
        @staticmethod
        def GPU(*decorator_args, **decorator_kwargs):
            if decorator_args and callable(decorator_args[0]) and len(decorator_args) == 1:
                return decorator_args[0]

            def _decorate(fn):
                return fn

            return _decorate

    spaces = _SpacesFallback()


ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT))
sys.path.insert(0, str(ROOT / "TDATR" / "eval"))

os.chdir(ROOT)
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
os.environ.setdefault("HYDRA_FULL_ERROR", "1")
os.environ.setdefault("NCCL_DEBUG", "WARN")

MODEL_REPO = os.environ.get("TDATR_MODEL_REPO", "CCWM/TDATR")
MODEL_FILE = os.environ.get("TDATR_MODEL_FILE", "model.pt")

_runner = None
_runner_lock = threading.Lock()
_predict_lock = threading.Lock()


def _free_port() -> int:
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
        sock.bind(("127.0.0.1", 0))
        return int(sock.getsockname()[1])


def _checkpoint_path() -> str:
    local_path = os.environ.get("TDATR_CKPT_PATH")
    if local_path and Path(local_path).exists():
        return local_path
    return hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)


def _build_cfg(max_new_tokens: int, temperature: float, seed: int):
    from omegaconf import OmegaConf, open_dict

    from TDATR_utils.dataclass import HulkConfig
    from TDATR_utils.utils import add_defaults

    base_cfg = OmegaConf.structured(HulkConfig)
    file_cfg = OmegaConf.load(ROOT / "configs" / "config.yaml")
    cfg = OmegaConf.merge(base_cfg, file_cfg)
    add_defaults(cfg)

    with open_dict(cfg):
        cfg.common.user_dir = str(ROOT / "TDATR")
        cfg.common.npu = False
        cfg.common.npu_jit_compile = False
        cfg.common.suppress_crashes = False

        cfg.distributed_training.distributed_world_size = 1
        cfg.distributed_training.distributed_num_procs = 1
        cfg.distributed_training.distributed_rank = 0
        cfg.distributed_training.distributed_local_rank = 0
        cfg.distributed_training.distributed_no_spawn = True
        cfg.distributed_training.distributed_backend = "nccl"
        cfg.distributed_training.distributed_master_addr = "127.0.0.1"
        cfg.distributed_training.distributed_master_port = _free_port()
        cfg.distributed_training.distributed_init_method = (
            f"tcp://127.0.0.1:{cfg.distributed_training.distributed_master_port}"
        )
        cfg.distributed_training.zero_sharding = "none"
        cfg.model_parallel.recompute_granularity = None

        with open_dict(cfg.model), open_dict(cfg.model.lora):
            cfg.model.ckpt = _checkpoint_path()
            cfg.model.lora.apply_lora = False
            cfg.model.use_naiive = True
            cfg.model.use_vit_encoder = False
            cfg.model.use_donut_encoder = True
            cfg.model.use_cfgi = True
            cfg.model.use_flash = False
            cfg.model.cross_flash_attn = False
            cfg.model.CFGI_CFG.use_flash = False
            cfg.model.CFGI_CFG.cross_flash_attn = False
            cfg.model.rectification_rotate_flag = False
            cfg.model.rectification_textline_height_flag = False

        cfg.task.use_ocr_plug = False
        cfg.task.seed = int(seed)

        cfg.generation.max_len = int(max_new_tokens)
        cfg.generation.temperature = float(temperature)
        cfg.generation.sampling_topk = 1
        cfg.generation.sampling_topp = 0.0

    return cfg


class TDATRRunner:
    def __init__(self, max_new_tokens: int, temperature: float, seed: int):
        import torch

        from TDATR.eval.infer import Dataset_infer
        from TDATR.models.detect.structures_.det_data_sample import DetDataSample
        from TDATR.models.detect.structures_.instance_data import InstanceData
        from TDATR.models.mini_gpt4_ipt_v2 import MiniGPT4
        from TDATR_utils.global_context import global_context as gpc
        from TDATR_utils.initialize import initialize_hulk

        self.torch = torch
        self.DetDataSample = DetDataSample
        self.InstanceData = InstanceData
        self.dataset = Dataset_infer()
        self.cfg = _build_cfg(max_new_tokens, temperature, seed)

        if not torch.distributed.is_initialized():
            initialize_hulk(self.cfg)
        elif gpc.config is None:
            gpc.config = self.cfg

        torch.cuda.set_device(0)
        self.device = torch.device("cuda")
        self.model = MiniGPT4(self.cfg).half()
        self.tokenizer = self.model.ipt_tokenizer
        self.model.eval()
        self.model = self.model.to(device=self.device)
        self.model.cfgi_decoder.neck.train()
        self.model.cfgi_decoder.encoder.train()

    def predict(self, image_path: str, max_new_tokens: int, temperature: float, seed: int):
        import cv2

        from TDATR.eval.infer import encode_img, single_prompt_process_cfgi
        from TDATR_utils.global_context import global_context as gpc

        self.cfg.generation.max_len = int(max_new_tokens)
        self.cfg.generation.temperature = float(temperature)
        self.cfg.task.seed = int(seed)
        gpc.config = self.cfg

        image_embed, det_input = encode_img(self.model, self.dataset, image_path, self.device)
        _, scale, img_shape, raw_scale, image_padding_shape = det_input

        data_sample = self.DetDataSample()
        data_sample.gt_instances = self.InstanceData()
        data_sample.set_metainfo({"img_shape": img_shape, "batch_input_shape": img_shape})

        query_text = "将图片中的表格转换为HTML语言。<iflytek_ret>"
        (
            _raw_query,
            raw_answer,
            clear_answer,
            _gen_tokens,
            _embs,
            _context_length,
            cell_boxes_pred,
            cell_span_html,
            cell_texts,
        ) = single_prompt_process_cfgi(
            self.model,
            self.tokenizer,
            "<end>",
            query_text,
            image_embed,
            max_new_tokens=int(max_new_tokens),
            sampling_topk=1,
            sampling_topp=0.0,
            temperature=float(temperature),
            max_length=int(max_new_tokens),
            random_seed=int(seed),
            image_shape=image_padding_shape,
            donuts_out=det_input[0],
            gt_inst=data_sample,
        )

        cell_boxes_pred = self.dataset.recover_pred_cell_box2raw_image(
            img_shape, scale, raw_scale, cell_boxes_pred
        )
        cells = self.dataset.process_cell_info(cell_boxes_pred, cell_texts, cell_span_html)

        image = cv2.imread(image_path)
        for cell in cell_boxes_pred:
            image = cv2.rectangle(
                image,
                cell[:2].tolist(),
                cell[2:].tolist(),
                (0, 0, 255),
                3,
            )
        vis_path = str(Path(tempfile.mkdtemp()) / "tdatr_cells.png")
        cv2.imwrite(vis_path, image)

        result = {
            "image_path": image_path,
            "answer": {
                "query": query_text,
                "clear_answer": clear_answer,
                "raw_answer": raw_answer,
                "cells": cells,
            },
        }
        return clear_answer, result, vis_path


def _get_runner(max_new_tokens: int, temperature: float, seed: int) -> TDATRRunner:
    global _runner
    with _runner_lock:
        if _runner is None:
            _runner = TDATRRunner(max_new_tokens, temperature, seed)
        return _runner


def _save_upload(image: Image.Image) -> str:
    if image is None:
        raise gr.Error("画像をアップロードしてください。")
    image = image.convert("RGB")
    tmp_dir = Path(tempfile.mkdtemp())
    image_path = tmp_dir / "input.png"
    image.save(image_path)
    return str(image_path)


@spaces.GPU(size="xlarge", duration=600)
def run_tdatr(image, max_new_tokens, temperature, seed):
    try:
        import torch

        if not torch.cuda.is_available():
            return (
                "<p>This Space is currently running without CUDA. "
                "TDATR needs ZeroGPU or another CUDA GPU hardware to run inference.</p>",
                "",
                None,
            )
        image_path = _save_upload(image)
        runner = _get_runner(int(max_new_tokens), float(temperature), int(seed))
        with _predict_lock:
            table_html, result, vis_path = runner.predict(
                image_path,
                int(max_new_tokens),
                float(temperature),
                int(seed),
            )
        return table_html, json.dumps(result, ensure_ascii=False, indent=2), vis_path
    except gr.Error:
        raise
    except Exception:
        escaped = html.escape(traceback.format_exc())
        return f"<pre>{escaped}</pre>", "", None


with gr.Blocks(title="TDATR") as demo:
    gr.Markdown("# TDATR")
    with gr.Row():
        image_input = gr.Image(type="pil", label="Table image")
        with gr.Column():
            max_tokens = gr.Slider(256, 4096, value=2048, step=128, label="Max tokens")
            temperature = gr.Slider(0.1, 1.0, value=0.5, step=0.1, label="Temperature")
            seed = gr.Number(value=42, precision=0, label="Seed")
            run_button = gr.Button("Run", variant="primary")
    with gr.Row():
        html_output = gr.HTML(label="HTML")
        vis_output = gr.Image(type="filepath", label="Cell boxes")
    json_output = gr.Code(language="json", label="Result JSON")

    run_button.click(
        fn=run_tdatr,
        inputs=[image_input, max_tokens, temperature, seed],
        outputs=[html_output, json_output, vis_output],
    )


if __name__ == "__main__":
    demo.queue(max_size=8).launch(show_error=True)