import base64
import hashlib
import io
import os
from pathlib import Path
from threading import Thread
from typing import Iterator, Optional, List, Union

import gradio as gr
import spaces
import torch
from PIL import Image
from pydantic import BaseModel
from qwen_vl_utils import process_vision_info
from swift.llm import (
    ModelType,
    get_model_tokenizer,
    get_default_template_type,
    get_template,
    inference,
    inference_stream,
)
from transformers import (
    Qwen2VLForConditionalGeneration,
    PreTrainedTokenizer,
    Qwen2VLProcessor,
    TextIteratorStreamer,
    AutoTokenizer,
)

MAX_MAX_NEW_TOKENS = 2048
DEFAULT_MAX_NEW_TOKENS = 1024
MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))

DESCRIPTION = """\
# Reasoning Paths Optimization: Learning to Reason and Explore From Diverse Paths

This Space demonstrates the reasoning paths optimization (RPO) framework with a Llama 3 model with 8B parameters fine-tuned for math reasoning. Feel free to play with it, or duplicate to run generations without a queue!

🔎 For more details about the RPO training framework, check out the [paper](https://arxiv.org/abs/2410.10858) or [code](https://github.com/DAMO-NLP-SG/reasoning-paths-optimization).
"""

LICENSE = """
<p/>

---
As a derivate work of [Llama-3-8b-chat](https://huggingface.co./meta-llama/Meta-Llama-3-8B) by Meta,
this demo is governed by the original [license](https://huggingface.co./meta-llama/Meta-Llama-3-8B/blob/main/LICENSE) and [acceptable use policy](https://huggingface.co./meta-llama/Meta-Llama-3-8B/blob/main/USE_POLICY.md).
"""


def convert_image_to_text(image: Image) -> str:
    # This is also how OpenAI encodes images: https://platform.openai.com/docs/guides/vision
    with io.BytesIO() as output:
        image.save(output, format="PNG")
        data = output.getvalue()
    return base64.b64encode(data).decode("utf-8")


def convert_text_to_image(text: str) -> Image:
    data = base64.b64decode(text.encode("utf-8"))
    return Image.open(io.BytesIO(data))


def save_image(image: Image.Image, folder: str) -> str:
    image_hash = hashlib.md5(image.tobytes()).hexdigest()
    path = Path(folder, f"{image_hash}.png")
    path.parent.mkdir(exist_ok=True, parents=True)
    if not path.exists():
        image.save(path)
    return str(path)


def resize_image(image: Image.Image, max_size: int) -> Image.Image:
    # Same as modeling.py resize_image
    width, height = image.size
    if width <= max_size and height <= max_size:
        return image
    if width > height:
        new_width = max_size
        new_height = round(height * max_size / width)
    else:
        new_height = max_size
        new_width = round(width * max_size / height)
    return image.resize((new_width, new_height), Image.LANCZOS)


class EvalModel(BaseModel, arbitrary_types_allowed=True):
    engine: str
    timeout: int = 60
    temperature: float = 0.0
    max_output_tokens: int = 512

    def run(self, inputs: List[Union[str, Image.Image]]) -> str:
        raise NotImplementedError

    def run_many(self, inputs: List[Union[str, Image.Image]], num: int) -> List[str]:
        raise NotImplementedError


class SwiftQwenModel(EvalModel):
    # https://github.com/modelscope/ms-swift/blob/main/docs/source_en/Multi-Modal/qwen2-vl-best-practice.md
    path: str = ""
    model: Optional[Qwen2VLForConditionalGeneration] = None
    tokenizer: Optional[PreTrainedTokenizer] = None
    engine: str = ModelType.qwen2_vl_7b_instruct
    image_size: int = 768
    image_dir: str = "data/qwen_images"

    def load(self):
        if self.model is None or self.tokenizer is None:
            self.model, self.tokenizer = get_model_tokenizer(
                self.engine,
                torch.bfloat16,
                model_kwargs={"device_map": "auto"},
                model_id_or_path=self.path or None,
            )

    def run(self, inputs: List[Union[str, Image.Image]]) -> str:
        self.load()
        template_type = get_default_template_type(self.engine)
        self.model.generation_config.max_new_tokens = self.max_output_tokens
        template = get_template(template_type, self.tokenizer)

        text = "\n\n".join([x for x in inputs if isinstance(x, str)])
        content = []
        for x in inputs:
            if isinstance(x, Image.Image):
                path = save_image(resize_image(x, self.image_size), self.image_dir)
                content.append(f"<img>{path}</img>")
        content.append(text)

        query = "".join(content)
        response, history = inference(self.model, template, query)
        return response

    def run_stream(self, inputs: List[Union[str, Image.Image]]) -> Iterator[str]:
        self.load()
        template_type = get_default_template_type(self.engine)
        self.model.generation_config.max_new_tokens = self.max_output_tokens
        template = get_template(template_type, self.tokenizer)

        text = "\n\n".join([x for x in inputs if isinstance(x, str)])
        content = []
        for x in inputs:
            if isinstance(x, Image.Image):
                path = save_image(resize_image(x, self.image_size), self.image_dir)
                content.append(f"<img>{path}</img>")
        content.append(text)

        query = "".join(content)
        generator = inference_stream(self.model, template, query)
        print_idx = 0
        print(f"query: {query}\nresponse: ", end="")
        for response, history in generator:
            delta = response[print_idx:]
            print(delta, end="", flush=True)
            print_idx = len(response)
            yield delta


class QwenModel(EvalModel):
    path: str = "models/qwen"
    engine: str = "Qwen/Qwen2-VL-7B-Instruct"
    model: Optional[Qwen2VLForConditionalGeneration] = None
    processor: Optional[Qwen2VLProcessor] = None
    tokenizer: Optional[AutoTokenizer] = None
    device: str = "cuda"
    image_size: int = 768
    lora_path: str = ""

    def load(self):
        if self.model is None:
            path = self.path if os.path.exists(self.path) else self.engine
            print(dict(load_path=path))
            # noinspection PyTypeChecker
            self.model = Qwen2VLForConditionalGeneration.from_pretrained(
                path, torch_dtype="auto", device_map="auto"
            )
            self.tokenizer = AutoTokenizer.from_pretrained(self.engine)

            if self.lora_path:
                print("Loading LORA from", self.lora_path)
                self.model.load_adapter(self.lora_path)

            self.model = self.model.to(self.device).eval()
            self.processor = Qwen2VLProcessor.from_pretrained(self.engine)
            torch.manual_seed(0)
            torch.cuda.manual_seed_all(0)

    def make_messages(self, inputs: List[Union[str, Image.Image]]) -> List[dict]:
        text = "\n\n".join([x for x in inputs if isinstance(x, str)])
        content = [
            dict(
                type="image",
                image=f"data:image;base64,{convert_image_to_text(resize_image(x, self.image_size))}",
            )
            for x in inputs
            if isinstance(x, Image.Image)
        ]
        content.append(dict(type="text", text=text))
        return [dict(role="user", content=content)]

    def run(self, inputs: List[Union[str, Image.Image]]) -> str:
        self.load()
        messages = self.make_messages(inputs)
        text = self.processor.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )

        image_inputs, video_inputs = process_vision_info(messages)
        # noinspection PyTypeChecker
        model_inputs = self.processor(
            text=[text],
            images=image_inputs,
            videos=video_inputs,
            padding=True,
            return_tensors="pt",
        ).to(self.device)

        with torch.inference_mode():
            generated_ids = self.model.generate(
                **model_inputs, max_new_tokens=self.max_output_tokens
            )

        generated_ids_trimmed = [
            out_ids[len(in_ids) :]
            for in_ids, out_ids in zip(model_inputs.input_ids, generated_ids)
        ]
        output_text = self.processor.batch_decode(
            generated_ids_trimmed,
            skip_special_tokens=True,
            clean_up_tokenization_spaces=False,
        )
        return output_text[0]

    def run_stream(self, inputs: List[Union[str, Image.Image]]) -> Iterator[str]:
        self.load()
        messages = self.make_messages(inputs)
        text = self.processor.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )

        image_inputs, video_inputs = process_vision_info(messages)
        # noinspection PyTypeChecker
        model_inputs = self.processor(
            text=[text],
            images=image_inputs,
            videos=video_inputs,
            padding=True,
            return_tensors="pt",
        ).to(self.device)

        streamer = TextIteratorStreamer(
            self.tokenizer,
            timeout=10.0,
            skip_prompt=True,
            skip_special_tokens=True,
        )

        generate_kwargs = dict(
            **model_inputs,
            streamer=streamer,
            max_new_tokens=self.max_output_tokens,
        )
        t = Thread(target=self.model.generate, kwargs=generate_kwargs)
        t.start()

        outputs = []
        for text in streamer:
            outputs.append(text)
            yield "".join(outputs)


if not torch.cuda.is_available():
    DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"


if torch.cuda.is_available():
    model = QwenModel()
    model.load()


@spaces.GPU
def generate(
    message: str,
    chat_history: list[dict],
    system_prompt: str = "",
    max_new_tokens: int = 1024,
    temperature: float = 0.6,
    top_p: float = 0.9,
    top_k: int = 50,
    repetition_penalty: float = 1.2,
) -> Iterator[str]:
    for text in model.run_stream([message]):
        yield text


chat_interface = gr.ChatInterface(
    fn=generate,
    additional_inputs=[
        gr.Textbox(label="System prompt", lines=6),
        gr.Slider(
            label="Max new tokens",
            minimum=1,
            maximum=MAX_MAX_NEW_TOKENS,
            step=1,
            value=DEFAULT_MAX_NEW_TOKENS,
        ),
        gr.Slider(
            label="Temperature",
            minimum=0.1,
            maximum=4.0,
            step=0.1,
            value=0.6,
        ),
        gr.Slider(
            label="Top-p (nucleus sampling)",
            minimum=0.05,
            maximum=1.0,
            step=0.05,
            value=0.9,
        ),
        gr.Slider(
            label="Top-k",
            minimum=1,
            maximum=1000,
            step=1,
            value=50,
        ),
        gr.Slider(
            label="Repetition penalty",
            minimum=1.0,
            maximum=2.0,
            step=0.05,
            value=1.2,
        ),
    ],
    stop_btn=None,
    examples=[
        [
            "Roger has 5 tennis balls. He buys 2 more cans of tennis balls. Each can has 3 tennis balls. How many tennis balls does he have now?"
        ],
        [
            "Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May?"
        ],
        [
            "Weng earns $12 an hour for babysitting. Yesterday, she just did 50 minutes of babysitting. How much did she earn?"
        ],
    ],
    cache_examples=False,
    type="messages",
)

with gr.Blocks(css_paths="style.css", fill_height=True) as demo:
    gr.Markdown(DESCRIPTION)
    gr.DuplicateButton(
        value="Duplicate Space for private use", elem_id="duplicate-button"
    )
    chat_interface.render()
    gr.Markdown(LICENSE)

if __name__ == "__main__":
    demo.queue(max_size=20).launch()