'''
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)
from transformers import AutoModelForCausalLM, AutoTokenizer
import transformers
from transformers import TextStreamer
import torch as th
from typing import Union, List, Dict
import argparse
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 transformers import TextIteratorStreamer
        from threading import Thread

        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('~/.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}')
