'''
Copyright (C) 2024-2025 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, Conversation
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]):
        # Used by backend.py for serving 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 Mistral7B(AbstractLLM):
    '''
    https://docs.mistral.ai/models/
    https://huggingface.co/docs/transformers/model_doc/mistral
    https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2
    https://huggingface.co/docs/transformers/main/chat_templating
    https://huggingface.co/blog/mixtral
    '''
    NAME = 'Mistral7B'
    model_id = 'mistralai/Mistral-7B-Instruct-v0.2'
    is_pipeline = True

    def __init__(self, *, 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.device = device  # override abstract class
        console.log(
            f'{self.NAME}> Loading {self.model_id} ({device}/{precision})')
        self.tok = AutoTokenizer.from_pretrained(self.model_id)
        llm_kwargs = {
            'torch_dtype': th.float16,
            'load_in_8bit': False,
            'load_in_4bit': False
        }
        if precision == 'fp16':
            llm_kwargs['torch_dtype'] = th.float16
        elif precision == 'fp32':
            llm_kwargs['torch_dtype'] = th.float32
        elif precision == 'bf16':
            llm_kwargs['torch_dtype'] = 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]):
        if self.is_pipeline:
            templated = self.tok.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=True)
            outputs = self.llm(templated, **self.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)
            input_length = model_inputs.shape[1]
            generated_ids = self.llm.generate(model_inputs, **self.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

    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


class Mixtral8x7B(Mistral7B):
    NAME = 'Mixtral8x7B'
    model_id = 'mistralai/Mixtral-8x7B-Instruct-v0.1'
    is_pipeline = True


def create_llm(args) -> AbstractLLM:
    # factory
    if args.llm == 'Mistral7B':
        model = Mistral7B(device=args.device, precision=args.precision)
        model.kwargs['max_new_tokens'] = args.max_new_tokens
    elif args.llm == 'Mixtral8x7B':
        model = Mixtral8x7B(device=args.device, precision=args.precision)
        model.kwargs['max_new_tokens'] = args.max_new_tokens
    else:
        raise NotImplementedError(f'{args.llm} is not yet implemented')
    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('--llm',
                    type=str,
                    default='Mistral7B',
                    choices=('Mistral7B', 'Mixtral8x7B'))
    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}')
