"""
Copyright (C) 2024-2026 Mo Zhou <lumin@debian.org>

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU Lesser General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public License
along with this program.  If not, see <https://www.gnu.org/licenses/>.
"""
import json
import os

from prompt_toolkit import prompt
from prompt_toolkit.styles import Style
from transformers import pipeline

try:
    from transformers import Conversation
except ImportError:

    class Conversation:
        def __init__(self, messages=None):
            self.messages = messages or []

        def add_message(self, message):
            self.messages.append(message)
import argparse
from typing import Dict, List, Union

import torch as th
import transformers
from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer

from .defaults import console


class AbstractLLM(object):
    def __init__(self):
        self.device = "cuda" if th.cuda.is_available() else "cpu"

    @th.no_grad()
    def generate(self, messages: Union[list, str], **kwargs):
        # Used by backend.py for serving a client
        raise NotImplementedError

    @th.no_grad()
    def generate_stream(self, messages: Union[list, str], **kwargs):
        # Used by backend.py for streaming response to a client
        raise NotImplementedError

    @th.no_grad()
    def __call__(self, *args, **kwargs):
        return self.generate(*args, **kwargs)

    def chat(self):
        # chat with LLM locally
        raise NotImplementedError


class HuggingFaceLLM(AbstractLLM):
    """
    Generic Hugging Face LLM backend wrapper.
    """

    NAME = "HuggingFaceLLM"
    model_id = None
    is_pipeline = True

    def __init__(self, model_id: str, *, device: str, precision: str):
        """
        torch_dtype: th.float32 requires 32GB CUDA memory.
                     th.float16/th.bfloat16 requires 16GB CUDA memory.
                     th.float16 has better hardware compatibility than bfloat16.
                     th.float16 has better compatibility than bfloat16.
        """
        super().__init__()
        self.model_id = model_id
        if self.model_id is None:
            raise ValueError("model_id must be specified")
        self.device = device  # override abstract class
        console.log(f"{self.NAME}> Loading {self.model_id} ({device}/{precision})")
        self.tok = AutoTokenizer.from_pretrained(self.model_id)
        # In transformers >= 5.0.0, torch_dtype is deprecated in favor of dtype.
        try:
            from packaging.version import Version

            is_v5 = Version(transformers.__version__) >= Version("5.0.0.dev0")
        except Exception:
            is_v5 = not transformers.__version__.startswith(("3.", "4."))
        dtype_key = "dtype" if is_v5 else "torch_dtype"

        llm_kwargs = {
            dtype_key: th.float16,
        }
        if precision == "fp16":
            llm_kwargs[dtype_key] = th.float16
        elif precision == "fp32":
            llm_kwargs[dtype_key] = th.float32
        elif precision == "bf16":
            llm_kwargs[dtype_key] = th.bfloat16
        elif precision == "8bit":
            llm_kwargs["load_in_8bit"] = True
        elif precision == "4bit":
            llm_kwargs["load_in_4bit"] = True
            llm_kwargs["bnb_4bit_compute_dtype"] = th.float16
        else:
            raise NotImplementedError(precision)
        if self.is_pipeline:
            self.llm = transformers.pipeline(
                "text-generation",
                model=self.model_id,
                model_kwargs=llm_kwargs,
                tokenizer=self.tok,
                device_map="auto" if self.device == "cuda" else self.device,
            )
        else:
            self.llm = AutoModelForCausalLM.from_pretrained(self.model_id, **llm_kwargs)
        if precision in ("fp16", "fp32", "bf16") and not self.is_pipeline:
            self.llm.to(self.device)
        else:
            pass
        self.kwargs = {
            "max_new_tokens": 512,
            "do_sample": True,
            "pad_token_id": 2,
            # default parameters for mixtral
            # https://huggingface.co/blog/mixtral
            "temperature": 0.7,
            "top_k": 50,
            "top_p": 0.95,
        }

    @th.no_grad()
    def generate(self, messages: List[Dict], **kwargs):
        gen_kwargs = self.kwargs.copy()
        for k, v in kwargs.items():
            if v is not None:
                gen_kwargs[k] = v
        if self.is_pipeline:
            templated = self.tok.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=True
            )
            outputs = self.llm(templated, **gen_kwargs)
            generated = outputs[0]["generated_text"][len(templated) :].lstrip()
            messages.append({"role": "assistant", "content": generated})
            return messages
        else:
            encoded = self.tok.apply_chat_template(
                messages, tokenize=True, return_tensors="pt", add_generation_prompt=True
            ).to(0)
            model_inputs = encoded.to(self.device)
            if hasattr(model_inputs, "keys"):
                input_length = model_inputs["input_ids"].shape[1]
                generated_ids = self.llm.generate(**model_inputs, **gen_kwargs)
            else:
                input_length = model_inputs.shape[1]
                generated_ids = self.llm.generate(model_inputs, **gen_kwargs)
            generated_new_ids = generated_ids[:, input_length:]
            generated_new_text = self.tok.batch_decode(
                generated_new_ids,
                skip_special_tokens=True,
                clean_up_tokenization_spaces=True,
            )[0]
            new_message = {"role": "assistant", "content": generated_new_text}
            messages.append(new_message)
            return messages

    @th.no_grad()
    def generate_stream(self, messages: List[Dict], **kwargs):
        from threading import Thread

        from transformers import TextIteratorStreamer

        streamer = TextIteratorStreamer(
            self.tok,
            skip_prompt=True,
            clean_up_tokenization_spaces=True,
            skip_special_tokens=True,
        )
        gen_kwargs = self.kwargs.copy()
        for k, v in kwargs.items():
            if v is not None:
                gen_kwargs[k] = v
        gen_kwargs["streamer"] = streamer

        if self.is_pipeline:
            templated = self.tok.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=True
            )
            thread = Thread(target=self.llm, args=(templated,), kwargs=gen_kwargs)
            thread.start()
        else:
            encoded = self.tok.apply_chat_template(
                messages, tokenize=True, return_tensors="pt", add_generation_prompt=True
            ).to(self.device)
            generate_kwargs = {**gen_kwargs}
            if hasattr(encoded, "keys"):
                generate_kwargs.update(encoded)
            else:
                generate_kwargs["input_ids"] = encoded
            thread = Thread(target=self.llm.generate, kwargs=generate_kwargs)
            thread.start()

        for new_text in streamer:
            yield new_text
        thread.join()

    def chat(self, chat=Conversation()):
        """
        https://huggingface.co/docs/transformers/main/en/main_classes/pipelines#transformers.ConversationalPipeline
        """
        if self.is_pipeline:
            pipe = self.llm
        else:
            pipe = pipeline(
                "conversational", model=self.llm, tokenizer=self.tok, device=self.device
            )

        streamer = TextStreamer(
            self.tok,
            skip_prompt=True,
            clean_up_tokenization_spaces=True,
            skip_special_tokens=True,
        )
        prompt_style = Style([("prompt", "bold fg:ansibrightcyan"), ("", "ansiwhite")])
        try:
            while text := prompt(f"Prompt[{len(chat.messages)}]> ", style=prompt_style):
                chat.add_message({"role": "user", "content": text})
                if True:
                    if self.is_pipeline:
                        templated = self.tok.apply_chat_template(
                            chat.messages, tokenize=False, add_generation_prompt=True
                        )
                        print(
                            f"\x1b[1;32mStream[{len(chat.messages)}]>\x1b[0m \x1b[0;32m",
                            end="",
                        )
                        outputs = pipe(
                            templated,
                            **self.kwargs,
                            streamer=streamer,
                            clean_up_tokenization_spaces=True,
                        )
                        print("\x1b[0m", end="")
                        generated = outputs[0]["generated_text"][len(templated) :]
                        chat.add_message({"role": "assistant", "content": generated})
                    else:
                        chat = pipe(chat, **self.kwargs)
        except EOFError:
            pass
        except KeyboardInterrupt:
            pass
        return chat


def create_llm(args) -> AbstractLLM:
    model = HuggingFaceLLM(args.model_id, device=args.device, precision=args.precision)
    model.kwargs["max_new_tokens"] = args.max_new_tokens
    if hasattr(args, "temperature") and args.temperature is not None:
        model.kwargs["temperature"] = args.temperature
    if hasattr(args, "top_p") and args.top_p is not None:
        model.kwargs["top_p"] = args.top_p
    return model


if __name__ == "__main__":
    ag = argparse.ArgumentParser("Chat with LLM locally.")
    ag.add_argument(
        "--max_new_tokens",
        type=int,
        default=512,
        help="max length of new token sequences added by llm",
    )
    ag.add_argument(
        "--debgpt_home", type=str, default=os.path.expanduser("~/.cache/debgpt")
    )
    ag.add_argument("--model_id", type=str, default="Qwen/Qwen3.5-0.8B")
    ag.add_argument("-i", "--ipython", action="store_true")
    ag.add_argument(
        "--device", type=str, default="cuda" if th.cuda.is_available() else "cpu"
    )
    ag.add_argument(
        "--precision", type=str, default="fp16" if th.cuda.is_available() else "4bit"
    )
    ag = ag.parse_args()
    console.log(ag)

    if not os.path.exists(ag.debgpt_home):
        os.mkdir(ag.debgpt_home)

    # for debugging
    if ag.ipython:
        llm = create_llm(ag)
        # XXX: ipython here is for debugging
        msg = [{"role": "user", "content": "hi!"}]
        import IPython

        IPython.embed(colors="neutral")
        exit()

    # load model and start chat
    llm = create_llm(ag)
    log = llm.chat()

    # save a record
    fpath = os.path.join(ag.debgpt_home, str(log.uuid) + ".json")
    with open(fpath, "wt") as f:
        json.dump(log.messages, f)

    console.log(f"LLM chat session saved at {fpath}")
