import os
import json
import glob
import warnings
import argparse

warnings.filterwarnings("ignore")
os.environ['TOKENIZERS_PARALLELISM'] = 'false'
os.environ['PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION'] = 'python'

import onnx
import torch

from utils.model import LlmModel, EmbeddingModel
from utils.tokenizer import LlmTokenizer
from utils.spinner import spinner_run
from utils.custom_op import FakeLinear
from utils.onnx_rebuilder import OnnxRebuilder
from utils.mnn_converter import MNNConverter
from utils.awq_quantizer import AwqQuantizer
from utils.smooth_quantizer import SmoothQuantizer
from utils.omni_quantizer import OmniQuantizer
from utils.torch_utils import onnx_export

class LlmExporter(torch.nn.Module):
    '''
    Base class for all llm model export. Inherits from [`torch.nn.Module`].
    '''
    def __init__(self, args):
        super().__init__()
        self.init_from_args(args)
        self.load_model(args.path)

    def init_from_args(self, args):
        self.args = args
        self.max_new_tokens = 1024
        self.dst_name = 'llm'
        # load config from args
        self.onnx_path = os.path.join(self.args.dst_path, 'onnx')
        if self.args.tokenizer_path is None:
            self.args.tokenizer_path = self.args.path
        if args.lm_quant_bit is None:
            self.args.lm_quant_bit = self.args.quant_bit
        if args.lm_quant_block is None:
            self.args.lm_quant_block = self.args.quant_block
        self.args.tie_word_embeddings = False
        # init export dst dir
        if not os.path.exists(self.args.dst_path):
            os.makedirs(self.args.dst_path)
        if not os.path.exists(self.onnx_path):
            os.makedirs(self.onnx_path)

    @spinner_run(f'load pretrained model ', True)
    def load_model(self, model_path):
        self.model = LlmModel.from_pretrained(model_path, args=self.args)
        self.tokenizer = LlmTokenizer.from_pretrained(
            self.args.tokenizer_path,
            model_type=self.model.config.model_type
        )
        self.model.tokenizer = self.tokenizer
        self.config = self.model.config
        self.model_type = self.config.model_type

        if self.args.awq or self.args.smooth:
            self.model.float()
        if self.args.export is not None:
            # set norm's weight as float for export
            def visit_module(module):
                if not isinstance(module, torch.nn.Linear) and hasattr(module, 'weight'):
                    module.float()
                for name, child in module.named_children():
                    visit_module(child)
            visit_module(self.model)

        self.model_dynamic_axes = {
            "input_ids" : { 0: "seq_len" },
            "attention_mask" : { 2: "seq_len", 3: "seq_len" },
            "position_ids" : { 1: "seq_len" },
        }

        self.llm_config = {
            'model_type': self.config.model_type,
            'hidden_size' : self.config.hidden_size,
            'layer_nums': self.config.num_hidden_layers,
            'attention_mask': 'float', # Will be determined by model later
            'attention_type': self.config.attention_type,
            'is_mrope': self.model.rotary.is_mrope
        }
        self.llm_config.update(self.model.get_config())
        # Attention scaling (gemma4 uses 1.0 instead of 1/sqrt(head_dim))
        if hasattr(self.model, 'blocks') and len(self.model.blocks) > 0:
            attn = self.model.blocks[0].self_attn
            if hasattr(attn, 'attn_scaling') and attn.attn_scaling != 1.0 / (self.config.head_dim ** 0.5):
                self.llm_config['attn_scale'] = attn.attn_scaling
        if self.config.sliding_window > 0:
            self.llm_config['sliding_window'] = self.config.sliding_window
        if hasattr(self.tokenizer, 'get_chat_template'):
             chat_template = self.tokenizer.get_chat_template()
             if chat_template is not None:
                 self.llm_config['jinja'] = {
                     'chat_template': chat_template
                 }
                 if self.tokenizer.bos_token:
                     self.llm_config['jinja']['bos'] = self.tokenizer.bos_token
                 if self.tokenizer.eos_token:
                     self.llm_config['jinja']['eos'] = self.tokenizer.eos_token
        # gemma4's HF template is too complex for minja parser, use simplified version
        if self.model_type == 'gemma4':
            self.llm_config['jinja'] = {
                'chat_template': "{{ bos_token }}{% for message in messages %}{% if message.role == \"system\" %}<|turn>system\n{{ message.content }}<turn|>\n{% elif message.role == \"user\" %}<|turn>user\n{{ message.content }}<turn|>\n{% elif message.role == \"assistant\" %}<|turn>model\n{{ message.content }}<turn|>\n{% endif %}{% endfor %}{% if add_generation_prompt %}<|turn>model\n{% endif %}",
                'bos': '<bos>',
                'eos': '<turn|>'
            }
        # glm_ocr's HF template is too complex for minja parser, use simplified version
        if self.model_type == 'glm_ocr':
            self.llm_config['jinja'] = {
                'chat_template': "[gMASK]<sop>{% for message in messages %}{% if message.role == \"user\" %}<|user|>\n{{ message.content }}{% elif message.role == \"assistant\" %}<|assistant|>\n{{ message.content }}{% elif message.role == \"system\" %}<|system|>\n{{ message.content }}{% endif %}{% endfor %}{% if add_generation_prompt %}<|assistant|>\n{% endif %}",
                'eos': '<|endoftext|>'
            }

        # tie word embeddings
        self.args.tie_word_embeddings = not self.args.seperate_embed and self.model.lm.lm.weight.equal(self.model.embed.embed.weight)
        # Pass properties from model to exporter
        self.visual = self.model.visual
        self.audio = self.model.audio
        self.talker = self.model.talker
        self.mtp = self.model.mtp
        self.scale_emb = self.model.scale_emb

        return model_path

    @torch.no_grad()
    def response(self, query):
        # self.imitate_quant()
        self.model.decode_buffer = []
        messages = [
            {"role": "user", "content": query}
        ]
        prompt = self.tokenizer.apply_chat_template(messages)
        if query not in prompt:
            prompt = query

        # Use model's tokenizer methods for encoding
        # For models with both visual and audio (e.g., gemma4), check content type
        has_audio = self.model.audio is not None and '<audio>' in prompt
        if has_audio:
            # Process audio first, then let visual handle the rest
            input_ids = self.model.audio.str_to_ids(prompt)
        elif self.model.visual is not None:
            input_ids = self.model.visual.str_to_ids(prompt)
        elif self.model.audio is not None:
            input_ids = self.model.audio.str_to_ids(prompt)
        else:
            input_ids = self.tokenizer(prompt, add_special_tokens=False, return_tensors="pt")['input_ids']

        seq_len = input_ids.numel()
        new_tokens = 0

        while new_tokens < self.max_new_tokens:
            attention_mask = self.model.get_attention_mask(seq_len, new_tokens)
            position_ids = self.model.get_position_ids(seq_len, new_tokens, input_ids)
            input_embeds = self.model.embedding(input_ids)
            deepstack_embeds = self.model.visual.deepstacks() if self.model.visual is not None else None

            logits, _, _ = self.model.forward(
                input_ids=input_embeds,
                attention_mask=attention_mask,
                position_ids=position_ids,
                logits_index=torch.tensor([-1], dtype=torch.int32),
                deepstack_embeds=deepstack_embeds
            )

            token_id = torch.argmax(logits[:,-1,:])
            seq_len += 1
            new_tokens += 1
            if token_id in self.tokenizer.stop_ids:
                print("", end='\n')
                break

            # Use tokenizer's method for decoding
            word = self.tokenizer.id_to_str(token_id)
            print(word, end="", flush=True)
            input_ids = token_id

        if hasattr(self.model, 'talker') and self.model.talker is not None:
            self.model.talker.generate()

    def export_mtp(self):
        if self.mtp is None:
            return
        mtp_onnx = self.mtp.export(self.onnx_path)
        if self.mnn_converter:
            self.mtp.unloaded_ops['/lm/lm_head/Linear'] = self.unloaded_ops['/lm/lm_head/Linear']
            MNNConverter(self, self.mtp.unloaded_ops).export(mtp_onnx)

    def export_eagle(self):
        if self.args.eagle_path is None:
            return
        from utils.eagle import Eagle
        self.eagle = Eagle.get_eagle(self.model_type)(self.args.eagle_path, self.model)
        eagle_onnx, eagle_fc_onnx = self.eagle.export(self.onnx_path)
        if self.mnn_converter:
            MNNConverter(self, None).export(eagle_onnx)
            MNNConverter(self, None).export(eagle_fc_onnx)

    def export_dflash(self):
        if not hasattr(self.args, 'dflash_path') or self.args.dflash_path is None:
            return
        from utils.dflash import DFlash
        self.dflash = DFlash(self.args.dflash_path, self.model)
        # Set target layer ids on args so model.forward() can use them
        self.args.dflash_target_layer_ids = self.dflash.target_layer_ids
        dflash_onnx, dflash_fc_onnx = self.dflash.export(self.onnx_path)
        if self.mnn_converter:
            # Disable transformerFuse for dflash model: dflash uses non-causal (bidirectional) attention,
            # but MNN's fused attention assumes causal masking which breaks dflash's attention pattern.
            # Use 8-bit quantization for dflash model to balance quality and size.
            MNNConverter(self, self.dflash.unloaded_ops).export(dflash_onnx, quant_bit=8, transformer_fuse=False)
            # FC model must NOT be quantized: the input (concatenated hidden states from
            # multiple target layers) has very large value ranges during prefill, which
            # causes int8 quantization overflow and produces all-zero outputs.
            MNNConverter(self, None).export(dflash_fc_onnx, quant_bit=0, transformer_fuse=False)


    @spinner_run(f'export embedding to ')
    def export_embed(self):
        import ctypes
        from utils.torch_utils import quant as torch_quant

        if hasattr(self.model, 'word_embeddings'):
            # embedding model's embed
            tensor_data = self.model.word_embeddings.weight.data
        else:
            tensor_data = self.model.embed.embed.weight.data

        format_bit = getattr(self.args, 'embed_bit', 16)
        quant_block = getattr(self.args, 'quant_block', 64)
        symmetric = getattr(self.args, 'sym', False)

        if self.args.skip_weight:
            format_name = f'int{format_bit}' if format_bit < 16 else 'bf16'
            embedding_file = f'{self.args.dst_path}/embeddings_{format_name}.bin'
            # Calculate expected size
            if format_bit == 16:
                file_size = tensor_data.numel() * 2
            else:
                oc, ic = tensor_data.shape
                block_size = ic if quant_block == 0 else quant_block
                block_num = ic // block_size
                q_weight_size = (oc * ic * format_bit + 7) // 8
                alpha_size = oc * block_num * (1 if symmetric else 2) * 4
                file_size = q_weight_size + alpha_size
                self.llm_config['tie_embeddings'] = {
                    "weight_offset": 0,
                    "alpha_offset": q_weight_size,
                    "alpha_size": alpha_size,
                    "quant_bit": format_bit,
                    "quant_block": quant_block,
                    "alpha_dtype": "fp32",
                }

            with open(embedding_file, 'wb') as f:
                if file_size > 0:
                    f.seek(file_size - 1)
                    f.write(b'\0')
            return embedding_file

        if format_bit == 16:
            # BF16 format
            tensor_data = tensor_data.bfloat16()
            data_ptr = tensor_data.untyped_storage().data_ptr()
            buffer = (ctypes.c_byte * (tensor_data.numel() * 2)).from_address(data_ptr)
            embedding_file = f'{self.args.dst_path}/embeddings_bf16.bin'
            with open(embedding_file, 'wb') as f:
                f.write(buffer)
        elif format_bit in [8, 4]:
            # Quantized formats
            quant_bit = format_bit
            format_name = f'int{format_bit}'
            awq = getattr(self.args, 'awq', False)
            hqq = getattr(self.args, 'hqq', False)

            # Apply quantization
            q_weight, alpha = torch_quant(tensor_data.float(), quant_bit, quant_block, symmetric, awq, hqq)

            # Save quantized weights and scales together in one file
            embedding_file = f'{self.args.dst_path}/embeddings_{format_name}.bin'
            with open(embedding_file, 'wb') as f:
                weight_size = f.write(q_weight.numpy().tobytes())
                alpha_size = f.write(alpha.numpy().tobytes())
            self.llm_config['tie_embeddings'] = {
                "weight_offset": 0,
                "alpha_offset": weight_size,
                "alpha_size": alpha_size,
                "quant_bit": quant_bit,
                "quant_block": quant_block,
                "alpha_dtype": "fp32",
            }
        else:
            raise ValueError(f"Unsupported embedding bit precision: {format_bit}")

        return embedding_file

    @spinner_run(f'export config to ')
    def export_config(self, mnn_config = False):
        with open(f'{self.args.dst_path}/export_args.json', 'w', encoding='utf-8') as f:
            json.dump(self.args.__dict__, f, ensure_ascii=False, indent=4)
        config_json = f'{self.args.dst_path}/llm_config.json'
        with open(config_json, 'w', encoding='utf-8') as f:
            json.dump(self.llm_config, f, ensure_ascii=False, indent=4)
        if not mnn_config:
            return config_json
        with open(f'{self.args.dst_path}/config.json', 'w', encoding='utf-8') as f:
            config = {
                "llm_model": f"{self.dst_name}.mnn",
                "llm_weight": f"{self.dst_name}.mnn.weight",
                "backend_type": "cpu",
                "thread_num": 4,
                "precision": "low",
                "memory": "low",
                # "system_prompt": "You are a helpful assistant.",
                "sampler_type": "mixed",
                "temperature": 0.8,
                "top_k": 40,
                "top_p": 0.9,
                "min_p": 0.05,
                "tfs_z": 1.0,
                "typical": 0.95,
                "repetition_penalty": 1.0,
                "presence_penalty": 0.0,
                "frequency_penalty": 0.0,
                "penalty_window": 0,
                "n_gram": 8,
                "ngram_factor": 1.0
            }
            config['tokenizer_file'] = 'tokenizer.mtok'
            if self.args.embed_bit < 16:
                config['embedding_file'] = f"embeddings_int{self.args.embed_bit}.bin"
            if hasattr(self, 'talker') and self.talker is not None:
                config['system_prompt'] = "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, capable of perceiving auditory and visual inputs, as well as generating text and speech."
                config['talker_max_new_tokens'] = 2048
                config['talker_speaker'] = "Chelsie"
                config['dit_steps'] = 5
                config['dit_solver'] = 1
            if self.model_type == "gemma3":
                config.update({'precision': "normal"})
            if (hasattr(self, 'visual') and self.visual is not None) or (hasattr(self, 'visual') and self.audio is not None):
                config['mllm'] = {
                    'backend_type': "cpu",
                    "thread_num": 4,
                    "precision": "normal",
                    "memory": "low"
                }
            if self.args.eagle_path is not None:
                config['speculative_type'] = 'eagle'
                config['hidden_states'] = True
            if hasattr(self.args, 'dflash_path') and self.args.dflash_path is not None:
                config['speculative_type'] = 'dflash'
                config['hidden_states'] = True
                config['dflash_model'] = 'dflash.mnn'
                config['dflash_fc'] = 'dflash_fc.mnn'
                config['dflash_block_size'] = self.dflash.block_size
                config['dflash_mask_token_id'] = self.dflash.mask_token_id
                config['dflash_target_layer_ids'] = self.dflash.target_layer_ids
            json.dump(config, f, ensure_ascii=False, indent=4)
        return config_json

    def imitate_quant(self):
        def quant_dequant(linear, quant_bit = self.args.quant_bit, quant_block = self.args.quant_block):
            weight = linear.weight.data
            oc, ic = weight.shape
            if quant_block == 0:
                block_size = ic
            else:
                block_size = quant_block
            block_num = ic // block_size
            weight = weight.reshape(oc, block_num, block_size)
            max_val = torch.max(weight, axis=-1, keepdims=True).values
            min_val = torch.min(weight, axis=-1, keepdims=True).values
            offset = 1 << (quant_bit - 1)
            clip_max = offset - 1
            clip_min = -offset
            scale = (max_val - min_val) / (clip_max - clip_min)
            q_weight = torch.round((weight - min_val) / scale) + clip_min
            q_weight = torch.clip(q_weight, clip_min, clip_max)
            dq_weight = (q_weight - clip_min) * scale + min_val
            dq_weight = dq_weight.reshape(oc, ic).float()
            linear.weight.data = dq_weight
            return linear
        with torch.no_grad():
            for i in range(self.config.num_hidden_layers):
                for name, child in self.model.blocks[i].self_attn.named_children():
                    if isinstance(child, torch.nn.Linear):
                        setattr(self.model.blocks[i].self_attn, name, quant_dequant(child))
                for name, child in self.model.blocks[i].mlp.named_children():
                    if isinstance(child, torch.nn.Linear):
                        setattr(self.model.blocks[i].mlp, name, quant_dequant(child))
            self.model.lm.lm = quant_dequant(self.model.lm.lm)

    def unload_param(self):
        self.unloaded_ops = {}
        self.experts = []
        self.expert_layer_ids = []
        def build_faker(real, name):
            faker = FakeLinear(real.in_features, real.out_features, real.bias is not None, name)
            self.unloaded_ops[name] = real.cpu()
            return faker
        # replace linear with fakelinear to save export memory and time
        with torch.no_grad():
            for i in range(len(self.model.blocks)):
                # different kv cache shape in different layers
                # if isinstance(self.config.num_attention_heads, list):
                self.model.blocks[i].self_attn.export_fused_attn = True
                is_moe = hasattr(self.model.blocks[i].mlp, 'is_moe') and self.model.blocks[i].mlp.is_moe
                if is_moe:
                    self.model.blocks[i].mlp.export_moe = True
                for name, child in self.model.blocks[i].self_attn.named_children():
                    if isinstance(child, torch.nn.Linear):
                        setattr(self.model.blocks[i].self_attn, name, build_faker(child, f'/layers.{i}/self_attn/{name}/Linear'))
                for name, child in self.model.blocks[i].mlp.named_children():
                    if isinstance(child, torch.nn.Linear):
                        setattr(self.model.blocks[i].mlp, name, build_faker(child, f'/layers.{i}/mlp/{name}/Linear'))
                # PLE per-layer Linear layers (gemma4)
                for name in ['per_layer_input_gate', 'per_layer_projection']:
                    child = getattr(self.model.blocks[i], name, None)
                    if isinstance(child, torch.nn.Linear):
                        setattr(self.model.blocks[i], name, build_faker(child, f'/layers.{i}/{name}/Linear'))
                # shared_expert in MLP-level MoE
                if is_moe and hasattr(self.model.blocks[i].mlp, 'shared_expert'):
                    for name, child in self.model.blocks[i].mlp.shared_expert.named_children():
                        if isinstance(child, torch.nn.Linear):
                            setattr(self.model.blocks[i].mlp.shared_expert, name, build_faker(child, f'/layers.{i}/mlp/shared_expert/{name}/Linear'))
                # MLP-level MoE experts
                if is_moe and hasattr(self.model.blocks[i].mlp, 'experts') and isinstance(self.model.blocks[i].mlp.experts, torch.nn.ModuleList):
                    self.experts.append(self.model.blocks[i].mlp.experts)
                    self.expert_layer_ids.append(i)
                    for j in range(len(self.model.blocks[i].mlp.experts)):
                        for name, cchild in self.model.blocks[i].mlp.experts[j].named_children():
                            if isinstance(cchild, torch.nn.Linear):
                                setattr(self.model.blocks[i].mlp.experts[j], name, build_faker(cchild, f'/expert/{i}_{j}/{name}'))
                # gemma4 decoder-level MoE (parallel to dense MLP)
                has_gemma4_moe = getattr(self.model.blocks[i], 'has_gemma4_moe', False)
                if has_gemma4_moe:
                    self.model.blocks[i].export_moe = True
                    # Unload moe_gate Linear
                    child = self.model.blocks[i].moe_gate
                    if isinstance(child, torch.nn.Linear):
                        self.model.blocks[i].moe_gate = build_faker(child, f'/layers.{i}/moe_gate/Linear')
                    # Unload experts
                    self.experts.append(self.model.blocks[i].experts)
                    self.expert_layer_ids.append(i)
                    for j in range(len(self.model.blocks[i].experts)):
                        for name, cchild in self.model.blocks[i].experts[j].named_children():
                            if isinstance(cchild, torch.nn.Linear):
                                setattr(self.model.blocks[i].experts[j], name, build_faker(cchild, f'/expert/{i}_{j}/{name}'))
            self.model.lm.lm = build_faker(self.model.lm.lm, f'/lm/lm_head/Linear')
            # PLE model-level Linear (gemma4)
            if hasattr(self.model, 'per_layer_model_projection') and isinstance(self.model.per_layer_model_projection, torch.nn.Linear):
                self.model.per_layer_model_projection = build_faker(self.model.per_layer_model_projection, f'/per_layer_model_projection/Linear')

    @spinner_run(f'export model weight to ')
    def onnx_load_param(self, onnx_path):
        return OnnxRebuilder(onnx_path, self.unloaded_ops).rebuild()

    @spinner_run(f'slim the graph of ')
    def slim_onnx(self, onnx_model):
        import onnxslim
        model = onnxslim.slim(onnx_model)
        onnx.save(model, onnx_model)
        return onnx_model

    @spinner_run(f'export onnx model to ')
    def export_onnx(self):
        # unload linear weight to save export memory
        self.unload_param()
        # move entire model to CPU to free GPU memory for quantization
        self.model.cpu()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        model = self.model
        seq_len = 3
        new_tokens = 0
        input_ids = torch.arange(seq_len, dtype=torch.long)
        attention_mask =  model.get_attention_mask(seq_len, new_tokens)
        position_ids = model.get_position_ids(seq_len, new_tokens, input_ids)
        onnx_model = f'{self.onnx_path}/{self.dst_name}.onnx'
        # For export onnx, don't need image or audio's embedding
        input_ids = model.embedding(input_ids)
        logits_index = torch.tensor([-1], dtype=torch.int32)
        if hasattr(model, 'talker') and model.talker is not None:
            output_names = ['logits', 'hidden_states', 'talker_embeds']
        else:
            output_names = ['logits', 'hidden_states']

        # Qwen3-VL
        if self.model_type in ['qwen3_vl', 'qwen3_vl_moe']:
            # add deepstack_embeds input
            deepstack_embeds = torch.randn(3, 1, self.config.hidden_size)
            onnx_export(
                model, (input_ids, attention_mask, position_ids, logits_index, deepstack_embeds),
                onnx_model,
                input_names=[
                    'input_ids', 'attention_mask', 'position_ids', 'logits_index', 'deepstack_embeds'
                ],
                output_names=output_names,
                dynamic_axes=self.model_dynamic_axes)
            return onnx_model

        # gemma4: add ple_embeddings + text_embeds_for_ple for PLE (Per-Layer Embeddings)
        if hasattr(model, 'embed_tokens_per_layer') and model.embed_tokens_per_layer is not None:
            raw_ids = torch.arange(seq_len, dtype=torch.long).unsqueeze(0)
            ple_embeddings = model.embed_tokens_per_layer(raw_ids)
            self.model_dynamic_axes['ple_embeddings'] = {1: 'seq_len'}
            onnx_export(
                model, (input_ids, attention_mask, position_ids, logits_index, None, ple_embeddings),
                onnx_model,
                input_names=[
                    'input_ids', 'attention_mask', 'position_ids', 'logits_index', 'ple_embeddings'
                ],
                output_names=output_names,
                dynamic_axes=self.model_dynamic_axes)
            return onnx_model

        # export to onnx
        onnx_export(
            model, (input_ids, attention_mask, position_ids, logits_index),
            onnx_model,
            input_names=[
                'input_ids', 'attention_mask', 'position_ids', 'logits_index'
            ],
            output_names=output_names,
            dynamic_axes=self.model_dynamic_axes)
        return onnx_model

    def awq_quant(self):
        self.awq_quantizer = AwqQuantizer(self.model)
        self.awq_quantizer.quantize()

    def omni_quant(self):
        default_samples = 128
        total_lines = default_samples

        if self.args.calib_data:
            print(f"检测到 calib_data: {self.args.calib_data}，开始读取...")
            self.model.args.calib_data = self.args.calib_data

            if os.path.exists(self.args.calib_data):
                with open(self.args.calib_data, 'r', encoding='utf-8') as f:
                    # 统计总行数
                    total_lines = sum(1 for _ in f)
            else:
                print(f"错误：找不到文件 {self.args.calib_data}")

        calib_samples = min(total_lines, default_samples)

        print(f"OmniQuant 将使用 {calib_samples} 个样本进行优化 (Epochs={getattr(self.args, 'omni_epochs', 20)})...")

        self.omni_quantizer = OmniQuantizer(
            model=self.model,
            max_calib_samples=calib_samples,
            act_bit=self.args.act_bit,
            act_sym=self.args.act_sym,
            generate_for_npu=self.args.generate_for_npu,

            epochs=getattr(self.args, 'omni_epochs', 20),
            lr=getattr(self.args, 'omni_lr', 5e-3),
            wd=getattr(self.args, 'omni_wd', 1e-4)
        )
        self.omni_quantizer.quantize(self.args.generate_for_npu)

    def smooth_quant(self):
        total_lines = 128
        if self.args.calib_data:
            print(f"检测到 calib_data: {self.args.calib_data}，开始读取...")
            self.model.args.calib_data = self.args.calib_data

            if os.path.exists(self.args.calib_data):
                with open(self.args.calib_data, 'r', encoding='utf-8') as f:
                    total_lines = sum(1 for _ in f)
            else:
                print(f"错误：找不到文件 {self.args.calib_data}")

        calib_samples = min(total_lines, 128)
        self.smooth_quantizer = SmoothQuantizer(model = self.model, max_calib_samples = calib_samples, act_bit=self.args.act_bit, act_sym=self.args.act_sym, generate_for_npu=self.args.generate_for_npu)
        self.smooth_quantizer.quantize()

    def export_vision(self):
        if self.visual is None:
            return
        vision_onnx = self.visual.export(self.onnx_path)
        if self.mnn_converter:
            fuse_transformer = self.visual.transformer_fuse
            native_group_conv = self.visual.group_conv_native
            quant_bit_visual = self.visual.quant_bit
            quant_block_visual = self.visual.quant_block
            if self.args.transformer_fuse:
                fuse_transformer = True
            if self.args.group_conv_native:
                native_group_conv = True
            if self.args.visual_quant_bit is not None:
                quant_bit_visual = self.args.visual_quant_bit
            if self.args.visual_quant_block is not None:
                quant_block_visual = self.args.visual_quant_block
            self.mnn_converter.export(vision_onnx, quant_bit_visual,
                                      quant_block_visual,
                                      transformer_fuse=fuse_transformer,
                                      group_conv_native=native_group_conv,
                                      weight_sym=self.args.visual_sym)

    def export_audio(self):
        if self.audio is None:
            return
        audio_onnx = self.audio.export(self.onnx_path)
        if self.mnn_converter: self.mnn_converter.export(audio_onnx, self.audio.quant_bit)

    def export_talker(self):
        if self.talker is None:
            return
        talker_onnx = self.talker.export(self.onnx_path)
        predit_onnx, dit_onnx, bigvgan_onnx = self.talker.token2wav.export(self.onnx_path)
        if self.mnn_converter:
            self.mnn_converter.export(talker_onnx, self.talker.quant_bit)
            self.mnn_converter.export(predit_onnx, self.talker.token2wav.quant_bit)
            self.mnn_converter.export(dit_onnx, self.talker.token2wav.quant_bit)
            self.mnn_converter.export(bigvgan_onnx, self.talker.token2wav.quant_bit)

    def export_ple_embed(self):
        """Export Per-Layer Embedding weights for gemma4."""
        import ctypes
        from utils.torch_utils import quant as torch_quant
        if not hasattr(self.model, 'embed_tokens_per_layer') or self.model.embed_tokens_per_layer is None:
            return
        embed = self.model.embed_tokens_per_layer
        tensor_data = embed.weight.data
        embed_scale = getattr(embed, 'scalar_embed_scale', 1.0)
        format_bit = getattr(self.args, 'embed_bit', 16)
        quant_block = getattr(self.args, 'quant_block', 64)
        symmetric = getattr(self.args, 'sym', False)
        if format_bit in [4, 8]:
            awq = getattr(self.args, 'awq', False)
            hqq = getattr(self.args, 'hqq', False)
            q_weight, alpha = torch_quant(tensor_data.float(), format_bit, quant_block, symmetric, awq, hqq)
            format_name = f'int{format_bit}'
            ple_file = f'{self.args.dst_path}/ple_embeddings_{format_name}.bin'
            with open(ple_file, 'wb') as f:
                weight_size = f.write(q_weight.numpy().tobytes())
                alpha_size = f.write(alpha.numpy().tobytes())
            self.llm_config['ple_embed_file'] = f'ple_embeddings_{format_name}.bin'
            self.llm_config['ple_quant'] = [0, weight_size, alpha_size, format_bit, quant_block]
        else:
            tensor_data = tensor_data.bfloat16()
            data_ptr = tensor_data.untyped_storage().data_ptr()
            buffer = (ctypes.c_byte * (tensor_data.numel() * 2)).from_address(data_ptr)
            ple_file = f'{self.args.dst_path}/ple_embeddings_bf16.bin'
            with open(ple_file, 'wb') as f:
                f.write(buffer)
            self.llm_config['ple_embed_file'] = 'ple_embeddings_bf16.bin'
        self.llm_config['ple_embed_scale'] = embed_scale
        self.llm_config['ple_embed_dim'] = embed.embedding_dim

    def export_language(self):
        # export_embedding
        if self.mnn_converter and self.args.tie_word_embeddings:
            pass # mnn tie_word_embeddings need't export embedding
        else:
            self.export_embed()
        # export PLE embedding (gemma4)
        self.export_ple_embed()
        # export transformer
        onnx_model = self.export_onnx()

        if self.args.onnx_slim:
            self.slim_onnx(onnx_model)
        if self.mnn_converter:
            tie_embeddings_info = MNNConverter(self, self.unloaded_ops).export(onnx_model)
            if tie_embeddings_info is not None:
                self.llm_config['tie_embeddings'] = tie_embeddings_info
        else:
            self.onnx_load_param(onnx_model)

    def export(self, export_type):
        if not self.args.skip_weight:
            if self.args.omni:
                self.omni_quant()
            if self.args.awq:
                self.awq_quant()
            if self.args.smooth:
                self.smooth_quant()
        export_mnn = export_type == 'mnn'
        self.mnn_converter = MNNConverter(self) if export_mnn else None
        self.export_talker()
        self.export_vision()
        self.export_audio()
        self.export_eagle()
        self.export_dflash()
        self.export_language()
        self.export_mtp()
        self.export_tokenizer()
        self.export_config(export_mnn)
        if export_mnn:
            # delete onnx file
            try:
                for file in glob.glob(f'{self.onnx_path}/*'):
                    os.remove(file)
                os.rmdir(self.onnx_path)
            except Exception as e:
                print(f"remove onnx error: {e}")

    @spinner_run(f'export tokenizer to ')
    def export_tokenizer(self):
        return self.tokenizer.export(self.args.dst_path)

class EmbeddingExporter(LlmExporter):
    def __init__(self, args):
        super().__init__(args)

    def response(self, query):
        self.model.eval()
        prompt = self.build_prompt(query)
        input_ids = self.tokenizer(prompt)['input_ids']
        seq_len = len(input_ids)
        input_ids = torch.tensor(input_ids)
        position_ids = self.model.get_position_ids(seq_len)
        attention_mask = self.model.get_attention_mask(seq_len)
        inputs_embeds = self.model.word_embed(input_ids)
        res = self.model.forward(inputs_embeds, attention_mask, position_ids)
        print(res, res.shape)
        return res

    def build_prompt(self, content):
        if self.config.model_type == 'bert':
            return f'[CLS]{content}[SEP]'
        if self.config.model_type == 'new':
            return f'<s> {content}</s>'
        if self.config.model_type == 'qwen3':
            return f'{content}<|endoftext|>'

    @spinner_run(f'load pretrained model ', True)
    def load_model(self, model_path):
        self.model = EmbeddingModel.from_pretrained(model_path, args=self.args)
        self.config = self.model.config
        self.model_type = self.config.model_type
        self.tokenizer = LlmTokenizer(model_path, self.model_type)
        self.llm_config = {
            'model_type': self.config.model_type,
            'hidden_size' : self.config.hidden_size,
            # qwen3 embedding is a causal decoder (last-token pooling) and needs a
            # causal mask; bert/gte encoders are bidirectional and use the all-ones
            # ('int') mask. Using 'int' for qwen3 makes attention bidirectional and
            # degrades the embeddings (see issue: identical/low-quality vectors).
            'attention_mask': 'float' if self.config.model_type == 'qwen3' else 'int',
            "jinja": {
                "chat_template": self.build_prompt("{{ messages | map(attribute='content') | join('') }}")
            },
            'is_visual': False
        }
        return model_path

    def export_reranker(self):
        seq_len = 4
        input_ids = torch.arange(12, dtype=torch.long)
        position_ids = self.model.get_position_ids(seq_len)
        attention_mask = self.model.get_attention_mask(seq_len)
        inputs_embeds = self.model.word_embed(input_ids)
        inputs_embeds = inputs_embeds.reshape(3, 4, self.config.hidden_size)
        attention_mask = torch.zeros(3, 1, 1, 4, dtype=torch.float)
        onnx_model = f'{self.onnx_path}/{self.dst_name}.onnx'
        onnx_export(
            self.model, (inputs_embeds, attention_mask, position_ids),
            onnx_model,
            input_names=[
                'input_ids',
                'attention_mask',
                'position_ids'
            ],
            output_names=['sentence_embeddings'],
            dynamic_axes={
                "input_ids" : { 0: "batch", 1: "seq_len" },
                "position_ids" : { 1: "seq_len" },
                "attention_mask" : { 0: "batch", 3: "seq_len" }
            })
        return onnx_model

    @spinner_run(f'export onnx model to ')
    def export_onnx(self):
        if self.model_type == 'qwen3':
            self.unload_param()
        else:
            self.unloaded_ops = None
        if self.model.is_reranker:
            return self.export_reranker()
        seq_len = 3
        input_ids = torch.arange(seq_len, dtype=torch.long)
        position_ids = self.model.get_position_ids(seq_len)
        attention_mask = self.model.get_attention_mask(seq_len)
        inputs_embeds = self.model.word_embed(input_ids)
        onnx_model = f'{self.onnx_path}/{self.dst_name}.onnx'
        onnx_export(
            self.model, (inputs_embeds, attention_mask, position_ids),
            onnx_model,
            input_names=[
                'input_ids',
                'attention_mask',
                'position_ids'
            ],
            output_names=['sentence_embeddings'],
            dynamic_axes={
                "input_ids" : { 1: "seq_len" },
                "position_ids" : { 1: "seq_len" },
                "attention_mask" : { 2: "seq_len", 3: "seq_len" }
            })
        return onnx_model

    def export(self, export_type):
        export_mnn = 'mnn' in export_type
        self.export_tokenizer()
        self.export_embed()
        self.export_config(export_mnn)
        onnx_model = self.export_onnx()
        if self.args.onnx_slim:
            self.slim_onnx(onnx_model)
        if export_mnn:
            transformer_fuse = not self.model.is_reranker
            tie_embeddings_info = MNNConverter(self, self.unloaded_ops).export(onnx_model, transformer_fuse=transformer_fuse)
            if tie_embeddings_info is not None:
                self.llm_config['tie_embeddings'] = tie_embeddings_info
            # delete onnx file
            try:
                for file in glob.glob(f'{self.onnx_path}/*'):
                    os.remove(file)
                os.rmdir(self.onnx_path)
            except Exception as e:
                print(f"remove onnx error: {e}")

def build_args(parser):
    parser.add_argument('--path', type=str, required=True,
                        help='path(`str` or `os.PathLike`):\nCan be either:'
                        '\n\t- A string, the *model id* of a pretrained model like `THUDM/chatglm-6b`. [TODO]'
                        '\n\t- A path to a *directory* clone from repo like `../chatglm-6b`.')
    parser.add_argument('--type', type=str, default=None,
                        help='type(`str`, *optional*):'
                        '\n\tThe pretrain llm model type.'
                        )
    parser.add_argument('--tokenizer_path', type=str, default=None, help='tokenizer path, default is `None` mean using `--path` value.')
    parser.add_argument('--eagle_path', type=str, default=None, help='eagle model path, default is `None`')
    parser.add_argument('--dflash_path', type=str, default=None, help='dflash draft model path, default is `None`')
    parser.add_argument('--lora_path', type=str, default=None, help='lora path, default is `None` mean not apply lora.')
    parser.add_argument('--gptq_path', type=str, default=None, help='gptq path, default is `None` mean not apply gptq.')
    parser.add_argument('--dst_path', type=str, default='./model', help='export onnx/mnn model to path, default is `./model`.')
    parser.add_argument('--verbose', action='store_true', help='Whether or not to print verbose.')
    parser.add_argument('--test', type=str, help='test model inference with query `TEST`.')
    parser.add_argument('--export', type=str, default=None, help='export model to an onnx/mnn model.')
    parser.add_argument('--onnx_slim', action='store_true', help='Whether or not to use onnx-slim.')
    parser.add_argument('--quant_bit', type=int, default=4, help='mnn quant bit, 2/3/4/8 (2 and 3 require ARMV86 i8mm + FP16), default is 4.')
    parser.add_argument('--quant_block', type=int, default=64, help='mnn quant block, 0 mean channel-wise, default is 64.')
    parser.add_argument('--visual_quant_bit', type=int, default=None, help='mnn visual quant bit, 4 or 8, default is setting in utils/vision.py by different vit model.')
    parser.add_argument('--visual_quant_block', type=int, default=None, help='mnn quant block, default is setting in utils/vision.py by different vit model.')
    parser.add_argument('--lm_quant_bit', type=int, default=None, help='mnn lm_head quant bit, 4 or 8, default is `quant_bit`.')
    parser.add_argument('--lm_quant_block', type=int, default=None, help='mnn lm_head quant block, 0 mean channle-wise, default is `quant_block`.')
    parser.add_argument('--mnnconvert', type=str, default='../../../build/MNNConvert', help='local mnnconvert path, if invalid, using pymnn.')
    parser.add_argument('--ppl', action='store_true', help='Whether or not to get all logits of input tokens.')
    parser.add_argument('--awq', action='store_true', help='Whether or not to use awq quant.')
    parser.add_argument('--hqq', action='store_true', help='Whether or not to use hqq quant.')
    parser.add_argument('--omni', action='store_true', help='Whether or not to use omni quant.')
    parser.add_argument('--transformer_fuse', action='store_true', help='Whether or not to fuse vision transformer op.')
    parser.add_argument('--enable_transformer_c4', dest='transformer_c4', action='store_true', default=False,
                        help='Whether or not to enable LLM C4 graph fusion.')
    parser.add_argument('--group_conv_native', action='store_true', help='Whether or not to keep native group_conv.')
    parser.add_argument('--smooth', action='store_true', help='Whether or not to use smooth quant.')
    parser.add_argument('--sym', action='store_true', help='Whether or not to using symmetric quant (without zeropoint), default is False.')
    parser.add_argument('--scale_bit', type=int, default=16, choices=[16, 32], help='Bit-width for quant scale/zero-point storage. Currently supports 16 (fp16, default) and 32 (fp32); 8/4 reserved for future.')
    parser.add_argument('--visual_sym', action='store_true', help='Whether or not to using symmetric quant (without zeropoint) for visual model, default is False.')
    parser.add_argument('--seperate_embed', action='store_true', help='For lm and embed shared model, whether or not to sepearte embed to avoid quant, default is False, if True, embed weight will be seperate to embedding bf16.bin.')
    parser.add_argument('--lora_split', action='store_true', help='Whether or not export lora split, default is False.')
    parser.add_argument('--calib_data', type=str, default=None, help='calibration data path, default is `None` mean not use calib data.')
    parser.add_argument('--act_bit', type=int, default=16, help='smooth quant act bit, 8 or 16, default is 16.')
    parser.add_argument('--embed_bit', type=int, default=16, choices=[16, 8, 4], help='embedding export bit precision, choices are 16 (bf16), 8 (int8), 4 (int4), default is 16.')
    parser.add_argument('--act_sym', action='store_true', help='smooth quant act us sym or not, default asym.')
    parser.add_argument('--quant_config', type=str, default=None, help='path to the JSON file for op-wise quantization configuration.')
    parser.add_argument('--generate_for_npu', action='store_true', help='Whether or not to generate model for NPU deployment, default is False.')
    parser.add_argument('--skip_weight', action='store_true', help='Whether or not to skip loading model weights, useful for testing export flow.')
    # omni quant
    parser.add_argument('--omni_epochs', type=int, default=20, help='OmniQuant 优化的轮数')
    parser.add_argument('--omni_lr', type=float, default=5e-3, help='OmniQuant 的学习率')
    parser.add_argument('--omni_wd', type=float, default=1e-4, help='OmniQuant 的权重衰减')

def export(path, **kwargs):
    parser = argparse.ArgumentParser()
    build_args(parser)
    args = parser.parse_args(['--path', path])
    for k, v in kwargs.items():
        setattr(args, k, v)
    if 'bge' in path:
        llm_exporter = EmbeddingExporter(args)
    else:
        llm_exporter = LlmExporter(args)
    # export
    llm_exporter.export(args.export)

def main():
    parser = argparse.ArgumentParser(description='llm_exporter', formatter_class=argparse.RawTextHelpFormatter)
    build_args(parser)
    args = parser.parse_args()

    model_path = args.path

    embedding_models = ['bge', 'gte', 'Qwen3-Embedding']
    if any(model in model_path for model in embedding_models):
        llm_exporter = EmbeddingExporter(args)
    else:
        llm_exporter = LlmExporter(args)

    # some actions
    if args.test is not None:
        llm_exporter.response(args.test)

    if args.export is not None:
        llm_exporter.export(args.export)

if __name__ == '__main__':
    main()