import gradio as gr
import os
import sys
import torch
import shutil
import urllib.request
import subprocess
import zipfile
import json
from pathlib import Path
from copy import deepcopy
from transformers import WhisperForConditionalGeneration
from peft import PeftModel

# ==========================================
# 1. 写入用户提供的转换脚本到磁盘
# ==========================================
def write_scripts():
    # 写入 convert-pt-to-ggml.py
    ggml_script = r"""
import io
import os
import sys
import struct
import json
import code
import torch
import numpy as np
import base64
from pathlib import Path

def bytes_to_unicode():
    bs = list(range(ord("!"), ord("~")+1))+list(range(ord("¡"), ord("¬")+1))+list(range(ord("®"), ord("ÿ")+1))
    cs = bs[:]
    n = 0
    for b in range(2**8):
        if b not in bs:
            bs.append(b)
            cs.append(2**8+n)
            n += 1
    cs = [chr(n) for n in cs]
    return dict(zip(bs, cs))

if len(sys.argv) < 4:
    print("Usage: convert-pt-to-ggml.py model.pt path-to-whisper-repo dir-output [use-f32]\n")
    sys.exit(1)

fname_inp   = Path(sys.argv[1])
dir_whisper = Path(sys.argv[2])
dir_out     = Path(sys.argv[3])

try:
    model_bytes = open(fname_inp, "rb").read()
    with io.BytesIO(model_bytes) as fp:
        checkpoint = torch.load(fp, map_location="cpu")
except Exception:
    print("Error: failed to load PyTorch model file:" , fname_inp)
    sys.exit(1)

hparams = checkpoint["dims"]
print("hparams:", hparams)

list_vars = checkpoint["model_state_dict"]

n_mels = hparams["n_mels"]
with np.load(dir_whisper / "whisper" / "assets" / "mel_filters.npz") as f:
    filters = torch.from_numpy(f[f"mel_{n_mels}"])

multilingual = hparams["n_vocab"] >= 51865
tokenizer = dir_whisper / "whisper" / "assets" / (multilingual and "multilingual.tiktoken" or "gpt2.tiktoken")
tokenizer_type = "tiktoken"
if not tokenizer.is_file():
    tokenizer = dir_whisper / "whisper" / "assets" / (multilingual and "multilingual" or "gpt2") / "vocab.json"
    tokenizer_type = "hf_transformers"
    if not tokenizer.is_file():
        print("Error: failed to find either tiktoken or hf_transformers tokenizer file:", tokenizer)
        sys.exit(1)

byte_encoder = bytes_to_unicode()
byte_decoder = {v:k for k, v in byte_encoder.items()}

if tokenizer_type == "tiktoken":
    with open(tokenizer, "rb") as f:
        contents = f.read()
        tokens = {base64.b64decode(token): int(rank) for token, rank in (line.split() for line in contents.splitlines() if line)}
elif tokenizer_type == "hf_transformers":
    with open(tokenizer, "r", encoding="utf8") as f:
        _tokens_raw = json.load(f)
        if '<|endoftext|>' in _tokens_raw:
            del _tokens_raw['<|endoftext|>']
        tokens = {bytes([byte_decoder[c] for c in token]): int(idx) for token, idx in _tokens_raw.items()}

fname_out = dir_out / "ggml-model.bin"
use_f16 = True
if len(sys.argv) > 4:
    use_f16 = False
    fname_out = dir_out / "ggml-model-f32.bin"

fout = fname_out.open("wb")

fout.write(struct.pack("i", 0x67676d6c))
fout.write(struct.pack("i", hparams["n_vocab"]))
fout.write(struct.pack("i", hparams["n_audio_ctx"]))
fout.write(struct.pack("i", hparams["n_audio_state"]))
fout.write(struct.pack("i", hparams["n_audio_head"]))
fout.write(struct.pack("i", hparams["n_audio_layer"]))
fout.write(struct.pack("i", hparams["n_text_ctx"]))
fout.write(struct.pack("i", hparams["n_text_state"]))
fout.write(struct.pack("i", hparams["n_text_head"]))
fout.write(struct.pack("i", hparams["n_text_layer"]))
fout.write(struct.pack("i", hparams["n_mels"]))
fout.write(struct.pack("i", use_f16))

fout.write(struct.pack("i", filters.shape[0]))
fout.write(struct.pack("i", filters.shape[1]))
for i in range(filters.shape[0]):
    for j in range(filters.shape[1]):
        fout.write(struct.pack("f", filters[i][j]))

fout.write(struct.pack("i", len(tokens)))

for key in tokens:
    fout.write(struct.pack("i", len(key)))
    fout.write(key)

for name in list_vars.keys():
    data = list_vars[name].squeeze().numpy()
    print("Processing variable: " , name ,  " with shape: ", data.shape)

    if name in ["encoder.conv1.bias", "encoder.conv2.bias"]:
        data = data.reshape(data.shape[0], 1)
        print(f"  Reshaped variable: {name} to shape: ", data.shape)

    n_dims = len(data.shape)

    ftype = 1
    if use_f16:
        if n_dims < 2 or \
                name == "encoder.conv1.bias"   or \
                name == "encoder.conv2.bias"   or \
                name == "encoder.positional_embedding" or \
                name == "decoder.positional_embedding":
            print("  Converting to float32")
            data = data.astype(np.float32)
            ftype = 0
    else:
        data = data.astype(np.float32)
        ftype = 0

    str_ = name.encode('utf-8')
    fout.write(struct.pack("iii", n_dims, len(str_), ftype))
    for i in range(n_dims):
        fout.write(struct.pack("i", data.shape[n_dims - 1 - i]))
    fout.write(str_)

    data.tofile(fout)

fout.close()
print("Done. Output file: " , fname_out)
"""
    with open("convert_pt_to_ggml.py", "w", encoding="utf-8") as f:
        f.write(ggml_script)

    # 写入 convert_whisper_to_coreml.py (包含 mlprogram 精度修复)
    coreml_script = r"""
import argparse
import torch
import torch.nn.functional as F
import coremltools as ct
from torch import Tensor
from torch import nn
from typing import Dict, Optional
from ane_transformers.reference.layer_norm import LayerNormANE as LayerNormANEBase
from whisper.model import Whisper, AudioEncoder, TextDecoder, ResidualAttentionBlock, MultiHeadAttention, ModelDimensions
from whisper import load_model
import whisper.model
whisper.model.MultiHeadAttention.use_sdpa = False

def linear_to_conv2d_map(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
    for k in state_dict:
        is_attention = all(substr in k for substr in ['attn', '.weight'])
        is_mlp = any(k.endswith(s) for s in ['mlp.0.weight', 'mlp.2.weight'])
        if (is_attention or is_mlp) and len(state_dict[k].shape) == 2:
            state_dict[k] = state_dict[k][:, :, None, None]

def correct_for_bias_scale_order_inversion(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
    state_dict[prefix + 'bias'] = state_dict[prefix + 'bias'] / state_dict[prefix + 'weight']
    return state_dict

class LayerNormANE(LayerNormANEBase):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._register_load_state_dict_pre_hook(correct_for_bias_scale_order_inversion)

class MultiHeadAttentionANE(MultiHeadAttention):
    def __init__(self, n_state: int, n_head: int):
        super().__init__(n_state, n_head)
        self.query =  nn.Conv2d(n_state, n_state, kernel_size=1)
        self.key = nn.Conv2d(n_state, n_state, kernel_size=1, bias=False)
        self.value = nn.Conv2d(n_state, n_state, kernel_size=1)
        self.out = nn.Conv2d(n_state, n_state, kernel_size=1)
    def forward(self, x: Tensor, xa: Optional[Tensor] = None, mask: Optional[Tensor] = None, kv_cache: Optional[dict] = None):
        q = self.query(x)
        if kv_cache is None or xa is None or self.key not in kv_cache:
            k = self.key(x if xa is None else xa)
            v = self.value(x if xa is None else xa)
        else:
            k = kv_cache[self.key]
            v = kv_cache[self.value]
        wv, qk = self.qkv_attention_ane(q, k, v, mask)
        return self.out(wv), qk
    def qkv_attention_ane(self, q: Tensor, k: Tensor, v: Tensor, mask: Optional[Tensor] = None):
        _, dim, _, seqlen = q.size()
        dim_per_head = dim // self.n_head
        scale = float(dim_per_head)**-0.5
        q = q * scale
        mh_q = q.split(dim_per_head, dim=1)
        mh_k = k.transpose(1,3).split(dim_per_head, dim=3)
        mh_v = v.split(dim_per_head, dim=1)
        mh_qk = [torch.einsum('bchq,bkhc->bkhq', [qi, ki]) for qi, ki in zip(mh_q, mh_k)]
        if mask is not None:
            for head_idx in range(self.n_head):
                mh_qk[head_idx] = mh_qk[head_idx] + mask[:, :seqlen, :, :seqlen]
        attn_weights = [aw.softmax(dim=1) for aw in mh_qk]
        attn = [torch.einsum('bkhq,bchk->bchq', wi, vi) for wi, vi in zip(attn_weights, mh_v)]
        attn = torch.cat(attn, dim=1)
        return attn, torch.cat(mh_qk, dim=1).float().detach()

class ResidualAttentionBlockANE(ResidualAttentionBlock):
    def __init__(self, n_state: int, n_head: int, cross_attention: bool = False):
        super().__init__(n_state, n_head, cross_attention)
        self.attn =  MultiHeadAttentionANE(n_state, n_head)
        self.attn_ln = LayerNormANE(n_state)
        self.cross_attn =  MultiHeadAttentionANE(n_state, n_head) if cross_attention else None
        self.cross_attn_ln =  LayerNormANE(n_state) if cross_attention else None
        n_mlp = n_state * 4
        self.mlp =  nn.Sequential(nn.Conv2d(n_state, n_mlp, kernel_size=1), nn.GELU(), nn.Conv2d(n_mlp, n_state, kernel_size=1))
        self.mlp_ln = LayerNormANE(n_state)

class AudioEncoderANE(AudioEncoder):
    def __init__(self, n_mels: int, n_ctx: int, n_state: int, n_head: int, n_layer: int):
        super().__init__(n_mels, n_ctx, n_state, n_head, n_layer)
        self.blocks = nn.ModuleList([ResidualAttentionBlockANE(n_state, n_head) for _ in range(n_layer)])
        self.ln_post = LayerNormANE(n_state)
    def forward(self, x: Tensor):
        x = F.gelu(self.conv1(x))
        x = F.gelu(self.conv2(x))
        x = (x + self.positional_embedding.transpose(0,1)).to(x.dtype).unsqueeze(2)
        for block in self.blocks: x = block(x)
        x = self.ln_post(x)
        x = x.squeeze(2).transpose(1, 2)
        return x

class TextDecoderANE(TextDecoder):
    def __init__(self, n_vocab: int, n_ctx: int, n_state: int, n_head: int, n_layer: int):
        super().__init__(n_vocab, n_ctx, n_state, n_head, n_layer)
        self.blocks= nn.ModuleList([ResidualAttentionBlockANE(n_state, n_head, cross_attention=True) for _ in range(n_layer)])
        self.ln= LayerNormANE(n_state)
    def forward(self, x: Tensor, xa: Tensor, kv_cache: Optional[dict] = None):
        offset = next(iter(kv_cache.values())).shape[3] if kv_cache else 0
        x = self.token_embedding(x) + self.positional_embedding[offset : offset + x.shape[-1]]
        x = x.to(xa.dtype)
        mask = self.mask[None, None, :, :].permute(0,3,1,2)
        x = x.transpose(1,2).unsqueeze(2)
        for block in self.blocks: x = block(x, xa, mask=mask, kv_cache=kv_cache)
        x = self.ln(x)
        x = x.permute(0,2,3,1).squeeze(0)
        if self.token_embedding.weight.shape[0] >= 51865:
            splits = self.token_embedding.weight.split(self.token_embedding.weight.shape[0]//11, dim=0)
            logits = torch.cat([torch.einsum('bid,jd->bij', x, split) for split in splits]).view(*x.shape[:2], -1)
        else:
            splits = self.token_embedding.weight.split(self.token_embedding.weight.shape[0]//12, dim=0)
            logits = torch.cat([torch.einsum('bid,jd->bij', x, split) for split in splits]).view(*x.shape[:2], -1)
        return logits

class WhisperANE(Whisper):
    def __init__(self, dims: ModelDimensions):
        super().__init__(dims)
        self.encoder = AudioEncoderANE(self.dims.n_mels, self.dims.n_audio_ctx, self.dims.n_audio_state, self.dims.n_audio_head, self.dims.n_audio_layer)
        self.decoder = TextDecoderANE(self.dims.n_vocab, self.dims.n_text_ctx, self.dims.n_text_state, self.dims.n_text_head, self.dims.n_text_layer)
        self._register_load_state_dict_pre_hook(linear_to_conv2d_map)
    def forward(self, mel: torch.Tensor, tokens: torch.Tensor) -> Dict[str, torch.Tensor]:
        return self.decoder(tokens, self.encoder(mel))
    def install_kv_cache_hooks(self, cache: Optional[dict] = None):
        pass 

def convert_encoder(hparams, model, quantize=False):
    model.eval()
    input_shape = (1, hparams.n_mels, 3000)
    input_data = torch.randn(input_shape)
    traced_model = torch.jit.trace(model, input_data)
    
    precision = ct.precision.FLOAT16 if quantize else ct.precision.FLOAT32
    
    model = ct.convert(
        traced_model, 
        convert_to="mlprogram", 
        inputs=[ct.TensorType(name="logmel_data", shape=input_shape)], 
        outputs=[ct.TensorType(name="output")], 
        compute_units=ct.ComputeUnit.ALL,
        compute_precision=precision
    )
    return model

def convert_decoder(hparams, model, quantize=False):
    model.eval()
    tokens_shape = (1, 1)
    audio_shape = (1, hparams.n_audio_ctx, hparams.n_audio_state)
    audio_data = torch.randn(audio_shape)
    token_data = torch.randint(hparams.n_vocab, tokens_shape).long()
    traced_model = torch.jit.trace(model, (token_data, audio_data))
    
    precision = ct.precision.FLOAT16 if quantize else ct.precision.FLOAT32

    model = ct.convert(
        traced_model, 
        convert_to="mlprogram", 
        inputs=[ct.TensorType(name="token_data", shape=tokens_shape, dtype=int), ct.TensorType(name="audio_data", shape=audio_shape)],
        compute_precision=precision
    )
    return model
"""
    with open("convert_whisper_to_coreml.py", "w", encoding="utf-8") as f:
        f.write(coreml_script)

write_scripts()
# 引入CoreML脚本
import convert_whisper_to_coreml as cml
from whisper import load_model as load_openai_model

# ==========================================
# 2. 全局设置 & 资源下载
# ==========================================

TEMP_DIR = Path("temp_workspace")
ASSETS_DIR = Path("whisper/assets")

# 清理重建
if TEMP_DIR.exists(): shutil.rmtree(TEMP_DIR)
TEMP_DIR.mkdir()
ASSETS_DIR.mkdir(parents=True, exist_ok=True)

def setup_assets():
    print("Checking Assets...")
    
    # 1. Mel Filters
    if not (ASSETS_DIR / "mel_filters.npz").exists():
        urllib.request.urlretrieve(
            "https://raw.githubusercontent.com/openai/whisper/main/whisper/assets/mel_filters.npz",
            ASSETS_DIR / "mel_filters.npz"
        )
        
    # 2. GPT2 Tokenizer (English)
    if not (ASSETS_DIR / "gpt2.tiktoken").exists():
        urllib.request.urlretrieve(
            "https://raw.githubusercontent.com/openai/whisper/main/whisper/assets/gpt2.tiktoken",
            ASSETS_DIR / "gpt2.tiktoken"
        )
        
    # 3. Multilingual Tokenizer (FIXED: Added this for small/medium/large models)
    if not (ASSETS_DIR / "multilingual.tiktoken").exists():
        print("Downloading multilingual tokenizer...")
        urllib.request.urlretrieve(
            "https://raw.githubusercontent.com/openai/whisper/main/whisper/assets/multilingual.tiktoken",
            ASSETS_DIR / "multilingual.tiktoken"
        )
        
    print("Assets Ready.")

setup_assets()

WHISPER_CPP_DIR = Path("whisper.cpp")

def setup_quantize_tool():
    # 检查仓库是否存在，不存在则克隆
    if not WHISPER_CPP_DIR.exists():
        print("Cloning whisper.cpp...")
        subprocess.run(["git", "clone", "https://github.com/ggerganov/whisper.cpp.git", str(WHISPER_CPP_DIR)], check=True)
    
    # 检查 quantize 工具是否已编译 (CMake 默认输出到 build/bin/quantize)
    quantize_bin = WHISPER_CPP_DIR / "build" / "bin" / "quantize"
    
    if not quantize_bin.exists():
        print("Building whisper.cpp quantization tool...")
        # 1. 配置构建目录
        subprocess.run(["cmake", "-B", "build"], cwd=WHISPER_CPP_DIR, check=True)
        # 2. 编译 quantize 目标
        # 注意：这里我们指定 --target quantize 只编译需要的工具，节省时间
        subprocess.run(["cmake", "--build", "build", "-j", "--config", "Release", "--target", "quantize"], cwd=WHISPER_CPP_DIR, check=True)
        
    print(f"Quantize tool ready at {quantize_bin}")

setup_quantize_tool()

# ==========================================
# 3. 核心转换逻辑 (HF -> OpenAI PT)
# ==========================================

def hf_to_openai_state_dict(model):
    state_dict = model.model.state_dict()
    new_state_dict = {}
    
    mapping = {
        "encoder.layers": "encoder.blocks",
        "decoder.layers": "decoder.blocks",
        "encoder.layer_norm": "encoder.ln_post",
        "decoder.layer_norm": "decoder.ln",
        "embed_tokens": "token_embedding",
        "encoder.embed_positions.weight": "encoder.positional_embedding",
        "decoder.embed_positions.weight": "decoder.positional_embedding",
        "fc1": "mlp.0",
        "fc2": "mlp.2",
        "final_layer_norm": "mlp_ln",
        "self_attn.q_proj": "attn.query",
        "self_attn.k_proj": "attn.key",
        "self_attn.v_proj": "attn.value",
        "self_attn.out_proj": "attn.out",
        "self_attn_layer_norm": "attn_ln",
        "encoder_attn.q_proj": "cross_attn.query",
        "encoder_attn.k_proj": "cross_attn.key",
        "encoder_attn.v_proj": "cross_attn.value",
        "encoder_attn.out_proj": "cross_attn.out",
        "encoder_attn_layer_norm": "cross_attn_ln",
    }

    for key, value in state_dict.items():
        new_key = key
        for k, v in mapping.items():
            new_key = new_key.replace(k, v)
        new_state_dict[new_key] = value

    return new_state_dict

def prepare_pt_file(base_model_id, lora_model_id):
    print(f"Loading HF Model: {base_model_id}")
    try:
        model = WhisperForConditionalGeneration.from_pretrained(base_model_id)
        
        if lora_model_id and lora_model_id.strip():
            print(f"Loading LoRA: {lora_model_id}")
            model = PeftModel.from_pretrained(model, lora_model_id)
            model = model.merge_and_unload()
        
        config = model.config
        dims = {
            'n_mels': config.num_mel_bins,
            'n_vocab': config.vocab_size,
            'n_audio_ctx': config.max_source_positions,
            'n_audio_state': config.d_model,
            'n_audio_head': config.encoder_attention_heads,
            'n_audio_layer': config.encoder_layers,
            'n_text_ctx': config.max_target_positions,
            'n_text_state': config.d_model,
            'n_text_head': config.decoder_attention_heads,
            'n_text_layer': config.decoder_layers
        }
        
        print("Converting State Dict...")
        state_dict = hf_to_openai_state_dict(model)
        
        pt_path = TEMP_DIR / "model.pt"
        torch.save({"dims": dims, "model_state_dict": state_dict}, pt_path)
        return pt_path, None
    except Exception as e:
        return None, str(e)

# ==========================================
# 4. 执行流程
# ==========================================

def run_process(base_model, lora_model, output_format, quant_type):
    logs = []
    
    out_dir = TEMP_DIR / "output"
    if out_dir.exists(): shutil.rmtree(out_dir)
    out_dir.mkdir()
    
    logs.append(f"Starting process for {base_model}...")
    
    # 1. 准备 PT 文件
    pt_file, err = prepare_pt_file(base_model, lora_model)
    if err:
        return None, f"Error preparing model: {err}"
    
    logs.append("Intermediate PT file created.")
    result_files = []

    # 2. GGML 转换
    if output_format == "GGML":
        # === 第一步：先转成 F16 (作为中间文件) ===
        temp_f16_path = out_dir / f"ggml-{base_model.split('/')[-1]}-f16.bin"
        
        # [关键修复] 使用 whisper.cpp 仓库自带的脚本，确保与 quantize 工具版本兼容
        official_script = WHISPER_CPP_DIR / "models" / "convert-pt-to-ggml.py"
        
        if not official_script.exists():
            return None, f"Error: Could not find official script at {official_script}. Make sure whisper.cpp cloned successfully."

        cmd = [
            sys.executable, 
            str(official_script),
            str(pt_file),
            ".",           # 指定 whisper/assets 的根目录 (当前目录)
            str(out_dir)   # 输出目录
        ]
        
        # 官方脚本逻辑：如果不传第4个参数，默认就是 F16。
        # 我们始终生成 F16 作为基底，所以这里不需要额外参数。
        
        logs.append(f"Running Official Python conversion: {' '.join(cmd)}")
        try:
            proc = subprocess.run(cmd, capture_output=True, text=True, check=True)
            logs.append(proc.stdout)
            
            # 官方脚本默认输出可能是 ggml-model.bin，我们需要找到它
            generated_bins = list(out_dir.glob("*.bin"))
            if not generated_bins:
                return None, f"Conversion failed. No .bin files found.\nSTDOUT:{proc.stdout}\nSTDERR:{proc.stderr}"
                
            raw_bin = generated_bins[0]
            os.rename(raw_bin, temp_f16_path)
            
            # === 第二步：判断是否需要二次量化 ===
            final_path = temp_f16_path
            
            if quant_type != "F16":
                # 定义量化后的文件名
                quantized_filename = f"ggml-{base_model.split('/')[-1]}-{quant_type.lower()}.bin"
                quantized_path = out_dir / quantized_filename
                
                # 调用 C++ quantize 工具
                quant_bin_path = WHISPER_CPP_DIR / "build" / "bin" / "quantize"
                
                # 命令格式：./quantize source dest type
                quant_cmd = [
                    str(quant_bin_path),
                    str(temp_f16_path),
                    str(quantized_path),
                    quant_type.lower() # 必须小写，如 q5_0
                ]
                
                logs.append(f"Running C++ Quantization ({quant_type})...")
                q_proc = subprocess.run(quant_cmd, capture_output=True, text=True, check=True)
                logs.append(q_proc.stdout)
                
                final_path = quantized_path
                
                # (可选) 删除中间的 F16 文件以节省空间
                if temp_f16_path.exists():
                    os.remove(temp_f16_path)

            result_files.append(str(final_path))
            
        except subprocess.CalledProcessError as e:
            logs.append(f"Conversion/Quantization Failed!\nSTDOUT: {e.stdout}\nSTDERR: {e.stderr}")

    # 3. CoreML 转换 (保持不变)
    elif output_format == "CoreML":
        logs.append("Running CoreML conversion (Imported logic)...")
        do_quantize_coreml = True 
        
        try:
            print("Loading converted PT file for CoreML...")
            whisper_model = load_openai_model(str(pt_file)).cpu()
            hparams = whisper_model.dims
            
            model_name = base_model.split('/')[-1]
            encoder = whisper_model.encoder
            decoder = whisper_model.decoder

            logs.append("Converting Encoder...")
            coreml_encoder = cml.convert_encoder(hparams, encoder, quantize=do_quantize_coreml)
            enc_path = out_dir / f"coreml-encoder-{model_name}.mlpackage"
            coreml_encoder.save(str(enc_path))
            
            logs.append("Converting Decoder...")
            coreml_decoder = cml.convert_decoder(hparams, decoder, quantize=do_quantize_coreml)
            dec_path = out_dir / f"coreml-decoder-{model_name}.mlpackage"
            coreml_decoder.save(str(dec_path))
            
            zip_filename = out_dir / f"coreml-{model_name}.zip"
            logs.append("Zipping files...")
            with zipfile.ZipFile(zip_filename, 'w', zipfile.ZIP_DEFLATED) as zf:
                for pkg in [enc_path, dec_path]:
                    for root, dirs, files in os.walk(pkg):
                        for file in files:
                            abs_path = os.path.join(root, file)
                            rel_path = os.path.relpath(abs_path, out_dir)
                            zf.write(abs_path, rel_path)
            
            result_files.append(str(zip_filename))
            logs.append("CoreML conversion done.")
            
        except Exception as e:
            import traceback
            err_msg = traceback.format_exc()
            logs.append(f"CoreML Error:\n{err_msg}")

    return result_files, "\n".join(logs)

# ==========================================
# 5. UI 构建
# ==========================================

with gr.Blocks(title="Whisper Converter") as demo:
    gr.Markdown("# Whisper Model Converter (HF -> GGML / CoreML)")
    gr.Markdown("Convert HuggingFace Whisper models (and optionally merge LoRA) to formats compatible with whisper.cpp and CoreML.")
    
    with gr.Row():
        with gr.Column():
            base_model = gr.Textbox(label="Base Model (HF Repo ID)", value="openai/whisper-tiny")
            lora_model = gr.Textbox(label="LoRA Model (HF Repo ID, Optional)", placeholder="Leave empty if none")
            fmt = gr.Dropdown(["GGML", "CoreML"], label="Output Format", value="GGML")
            quant_type = gr.Dropdown(["F16", "Q4_0", "Q4_1", "Q5_0", "Q5_1", "Q8_0"], label="GGML Quantization Type", value="F16", info="F16 is default. Q5_0 is recommended for balance.")
            btn = gr.Button("Convert", variant="primary")
        
        with gr.Column():
            output_files = gr.File(label="Download")
            logs = gr.Textbox(label="Logs", lines=15)
            
    btn.click(run_process, inputs=[base_model, lora_model, fmt, quant_type], outputs=[output_files, logs])

if __name__ == "__main__":
    demo.queue().launch()