import torch
import json
import struct
import os
import numpy as np
from collections import OrderedDict
from transformers import AutoModelForCausalLM, AutoConfig, AutoTokenizer


class DQModel:
    """Loader and inference engine for DQ quantized models."""

    def __init__(self, dq_path):
        self.dq_path = dq_path
        self.model = None
        self.tokenizer = None
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self._load_model()

    def _unpack_2bit(self, packed_data, shape, pad):
        """Unpack 4 unsigned 2-bit values per byte."""
        packed = torch.from_numpy(np.frombuffer(packed_data, dtype=np.uint8))
        bits = [
            packed & 0x03,
            (packed >> 2) & 0x03,
            (packed >> 4) & 0x03,
            (packed >> 6) & 0x03
        ]
        flat = torch.stack(bits, dim=1).flatten()
        total = int(np.prod(shape))
        if pad:
            flat = flat[:total + pad]
        return flat[:total].reshape(shape)

    def _unpack_3bit(self, packed_data, shape, pad):
        """Unpack 8 unsigned 3-bit values from 3 bytes."""
        packed = torch.from_numpy(np.frombuffer(packed_data, dtype=np.uint8))
        total = int(np.prod(shape))
        groups = (total + pad + 7) // 8
        packed_groups = packed[:groups * 3].view(-1, 3)
        flat = torch.zeros(groups * 8, dtype=torch.uint8)
        flat[0::8] = packed_groups[:, 0] & 0x07
        flat[1::8] = (packed_groups[:, 0] >> 3) & 0x07
        flat[2::8] = ((packed_groups[:, 0] >> 6) & 0x03) | ((packed_groups[:, 1] & 0x01) << 2)
        flat[3::8] = (packed_groups[:, 1] >> 1) & 0x07
        flat[4::8] = (packed_groups[:, 1] >> 4) & 0x07
        flat[5::8] = ((packed_groups[:, 1] >> 7) & 0x01) | ((packed_groups[:, 2] & 0x03) << 1)
        flat[6::8] = (packed_groups[:, 2] >> 2) & 0x07
        flat[7::8] = (packed_groups[:, 2] >> 5) & 0x07
        if pad:
            flat = flat[:total + pad]
        return flat[:total].reshape(shape)

    def _unpack_4bit(self, packed_data, shape, pad):
        """Unpack 2 unsigned 4-bit values per byte."""
        packed = torch.from_numpy(np.frombuffer(packed_data, dtype=np.uint8))
        low = packed & 0x0F
        high = (packed >> 4) & 0x0F
        flat = torch.stack([low, high], dim=1).flatten()
        total = int(np.prod(shape))
        if pad:
            flat = flat[:total + pad]
        return flat[:total].reshape(shape)

    def _unpack_5bit(self, packed_data, shape, pad):
        """Unpack 8 unsigned 5-bit values from 5 bytes."""
        packed = torch.from_numpy(np.frombuffer(packed_data, dtype=np.uint8))
        total = int(np.prod(shape))
        groups = (total + pad + 7) // 8
        packed_groups = packed[:groups * 5].view(-1, 5)
        flat = torch.zeros(groups * 8, dtype=torch.uint8)
        flat[0::8] = packed_groups[:, 0] & 0x1F
        flat[1::8] = ((packed_groups[:, 0] >> 5) & 0x07) | ((packed_groups[:, 1] & 0x03) << 3)
        flat[2::8] = (packed_groups[:, 1] >> 2) & 0x1F
        flat[3::8] = ((packed_groups[:, 1] >> 7) & 0x01) | ((packed_groups[:, 2] & 0x0F) << 1)
        flat[4::8] = ((packed_groups[:, 2] >> 4) & 0x0F) | ((packed_groups[:, 3] & 0x01) << 4)
        flat[5::8] = (packed_groups[:, 3] >> 1) & 0x1F
        flat[6::8] = ((packed_groups[:, 3] >> 6) & 0x03) | ((packed_groups[:, 4] & 0x07) << 2)
        flat[7::8] = (packed_groups[:, 4] >> 3) & 0x1F
        if pad:
            flat = flat[:total + pad]
        return flat[:total].reshape(shape)

    def _unpack_6bit(self, packed_data, shape, pad):
        """Unpack 4 unsigned 6-bit values from 3 bytes."""
        packed = torch.from_numpy(np.frombuffer(packed_data, dtype=np.uint8))
        total = int(np.prod(shape))
        groups = (total + pad + 3) // 4
        packed_groups = packed[:groups * 3].view(-1, 3)
        flat = torch.zeros(groups * 4, dtype=torch.uint8)
        flat[0::4] = packed_groups[:, 0] & 0x3F
        flat[1::4] = ((packed_groups[:, 0] >> 6) & 0x03) | ((packed_groups[:, 1] & 0x0F) << 2)
        flat[2::4] = ((packed_groups[:, 1] >> 4) & 0x0F) | ((packed_groups[:, 2] & 0x03) << 4)
        flat[3::4] = (packed_groups[:, 2] >> 2) & 0x3F
        if pad:
            flat = flat[:total + pad]
        return flat[:total].reshape(shape)

    def _unpack_8bit(self, packed_data, shape, pad):
        """No packing needed for 8-bit."""
        tensor = torch.from_numpy(np.frombuffer(packed_data, dtype=np.uint8))
        total = int(np.prod(shape))
        return tensor[:total].reshape(shape)

    def _unpack_12bit(self, packed_data, shape, pad):
        """Unpack 2 unsigned 12-bit values from 3 bytes."""
        packed = torch.from_numpy(np.frombuffer(packed_data, dtype=np.uint8))
        total = int(np.prod(shape))
        groups = (total + pad + 1) // 2
        packed_groups = packed[:groups * 3].view(-1, 3)
        flat = torch.zeros(groups * 2, dtype=torch.uint16)
        flat[0::2] = packed_groups[:, 0].to(torch.uint16) | ((packed_groups[:, 1] & 0x0F).to(torch.uint16) << 8)
        flat[1::2] = ((packed_groups[:, 1] >> 4) & 0x0F).to(torch.uint16) | (packed_groups[:, 2].to(torch.uint16) << 4)
        if pad:
            flat = flat[:total + pad]
        return flat[:total].reshape(shape)

    def _unpack_16bit(self, packed_data, shape, pad):
        """No packing for float16."""
        tensor = torch.from_numpy(np.frombuffer(packed_data, dtype=np.float16))
        total = int(np.prod(shape))
        return tensor[:total].reshape(shape)

    def _unpack_32bit(self, packed_data, shape, pad):
        """No packing for float32."""
        tensor = torch.from_numpy(np.frombuffer(packed_data, dtype=np.float32))
        total = int(np.prod(shape))
        return tensor[:total].reshape(shape)

    def _dequantize_tensor(self, data, info):
        """Dequantize tensor back to float32."""
        dtype = info['dtype']
        shape = info['shape']
        pad = info.get('pad', 0)
        unpacked_shape = info.get('packed_shape', shape)

        unpack_functions = {
            'int2_sym': (self._unpack_2bit, True, 'uint8'),
            'int3_sym': (self._unpack_3bit, True, 'uint8'),
            'int4_sym': (self._unpack_4bit, True, 'uint8'),
            'int5_sym': (self._unpack_5bit, True, 'uint8'),
            'int6_sym': (self._unpack_6bit, True, 'uint8'),
            'int8_sym': (self._unpack_8bit, True, 'uint8'),
            'int12_sym': (self._unpack_12bit, True, 'uint16'),
            'float16': (self._unpack_16bit, False, 'float16'),
            'float32': (self._unpack_32bit, False, 'float32'),
        }

        if dtype not in unpack_functions:
            raise ValueError(f"Unknown dtype: {dtype}")

        unpack_fn, is_quantized, storage_dtype = unpack_functions[dtype]

        if is_quantized:
            unsigned = unpack_fn(data, unpacked_shape, pad)
            scale = info['scale']
            level = info.get('level', info.get('offset', 0))
            signed = unsigned.float() - level
            tensor = signed * scale
        else:
            tensor = unpack_fn(data, shape, pad).float()

        return tensor.reshape(shape)

    def _load_model(self):
        with open(self.dq_path, 'rb') as f:
            magic = f.read(4)
            if magic != b'DQMD':
                raise ValueError("Invalid file format: missing DQMD magic")

            version = struct.unpack('<I', f.read(4))[0]
            if version != 1:
                raise ValueError(f"Unsupported format version: {version}")

            metadata_size = struct.unpack('<I', f.read(4))[0]
            metadata = json.loads(f.read(metadata_size).decode('utf-8'))

            config_dict = metadata['config']
            bits_per_param = metadata['bits_per_param']

            config = AutoConfig.for_model(**config_dict)

            tensor_headers_size = struct.unpack('<I', f.read(4))[0]
            tensor_headers = json.loads(f.read(tensor_headers_size).decode('utf-8'))

            state_dict = OrderedDict()

            for name, info in tensor_headers.items():
                data = f.read(info['data_size'])
                tensor = self._dequantize_tensor(data, info)
                state_dict[name] = tensor

            self.model = AutoModelForCausalLM.from_config(config)
            self.model.load_state_dict(state_dict)
            self.model = self.model.to(self.device)
            self.model.eval()
            del state_dict

            tokenizer_path = os.path.dirname(self.dq_path)
            if os.path.exists(os.path.join(tokenizer_path, 'tokenizer.json')):
                self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
            elif os.path.exists(os.path.join(tokenizer_path, 'vocab.json')):
                self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
            else:
                self.tokenizer = AutoTokenizer.from_pretrained('gpt2')

            if self.tokenizer.pad_token is None:
                self.tokenizer.pad_token = self.tokenizer.eos_token

    @torch.no_grad()
    def generate(self, prompt, max_length=100, temperature=0.3, top_p=0.97,
                 top_k=50, repetition_penalty=1.1, do_sample=True):
        if self.model is None or self.tokenizer is None:
            raise ValueError("Model or tokenizer not loaded")

        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.device)

        outputs = self.model.generate(
            inputs.input_ids,
            max_length=max_length,
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            repetition_penalty=repetition_penalty,
            do_sample=do_sample,
            pad_token_id=self.tokenizer.pad_token_id,
            eos_token_id=self.tokenizer.eos_token_id,
            use_cache=True,
        )

        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)


def format_chat_prompt(messages):
    formatted = []
    for role, content in messages:
        if role == "user":
            formatted.append(f"<|user|>\n{content}")
        elif role == "assistant":
            formatted.append(f"<|assistant|>\n{content}")
    return "\n".join(formatted) + "\n<|assistant|>\n"


def chat_with_dq_model(model_path, max_history=5):
    print("=" * 60)
    print("Chat with DQ Model")
    print("Format: <|user|> and <|assistant|>")
    print("=" * 60)

    dq_model = DQModel(model_path)
    conversation_history = []
    temperature = 0.7
    top_p = 0.9
    top_k = 50
    system_message = None

    print("\n" + "=" * 60)
    print("Model ready. Type /help for commands.")
    print("=" * 60)

    while True:
        try:
            user_input = input("\nYou: ").strip()
            if not user_input:
                continue

            if user_input.startswith('/'):
                command = user_input.lower()

                if command == '/exit':
                    print("Goodbye!")
                    break

                elif command == '/clear':
                    conversation_history = []
                    print("History cleared.")
                    continue

                elif command == '/help':
                    print("Commands:")
                    print("  /clear      Clear conversation history")
                    print("  /temp X     Set temperature (0.1-2.0)")
                    print("  /top_p X    Set top_p (0.1-1.0)")
                    print("  /top_k X    Set top_k (1-100)")
                    print("  /system X   Set system message")
                    print("  /stats      Show current parameters")
                    print("  /exit       Exit")
                    continue

                elif command == '/stats':
                    print(f"Temperature: {temperature}, Top-p: {top_p}, Top-k: {top_k}")
                    print(f"System message: {system_message or 'None'}")
                    print(f"History: {len(conversation_history)} messages")
                    continue

                elif command.startswith('/temp'):
                    try:
                        val = float(command.split()[1])
                        if 0.1 <= val <= 2.0:
                            temperature = val
                            print(f"Temperature: {temperature}")
                        else:
                            print("Invalid range. Use 0.1-2.0")
                    except:
                        print("Usage: /temp 0.7")
                    continue

                elif command.startswith('/top_p'):
                    try:
                        val = float(command.split()[1])
                        if 0.1 <= val <= 1.0:
                            top_p = val
                            print(f"Top-p: {top_p}")
                        else:
                            print("Invalid range. Use 0.1-1.0")
                    except:
                        print("Usage: /top_p 0.9")
                    continue

                elif command.startswith('/top_k'):
                    try:
                        val = int(command.split()[1])
                        if 1 <= val <= 100:
                            top_k = val
                            print(f"Top-k: {top_k}")
                        else:
                            print("Invalid range. Use 1-100")
                    except:
                        print("Usage: /top_k 50")
                    continue

                elif command.startswith('/system'):
                    system_message = user_input[8:].strip() or None
                    print(f"System message: {system_message or 'cleared'}")
                    continue

                else:
                    print(f"Unknown command: {command}")
                    continue

            conversation_history.append(("user", user_input))
            if len(conversation_history) > max_history:
                conversation_history = conversation_history[-max_history:]

            messages = []
            if system_message:
                messages.append(("system", system_message))
            messages.extend(conversation_history)

            prompt = format_chat_prompt(messages)

            print("\nAssistant: ", end="", flush=True)

            import time
            start = time.time()
            response = dq_model.generate(
                prompt,
                max_length=len(dq_model.tokenizer.encode(prompt)) + 200,
                temperature=temperature,
                top_p=top_p,
                top_k=top_k,
                repetition_penalty=1.1,
                do_sample=True
            )

            if prompt in response:
                response = response[len(prompt):]
            if "<|user|>" in response:
                response = response.split("<|user|>")[0]
            if "<|assistant|>" in response:
                response = response.split("<|assistant|>")[0]
            response = response.strip()

            print(response)
            print(f"[{time.time() - start:.2f}s]")

            conversation_history.append(("assistant", response))

        except KeyboardInterrupt:
            print("\nGoodbye!")
            break
        except Exception as e:
            print(f"Error: {e}")
            continue


if __name__ == "__main__":
    chat_with_dq_model("VDrontV1.5-Flash_Q16.dq", 256)