import gradio as gr
import torch
import os
import gc
import shutil
import requests
import json
import struct
import numpy as np
import re
from pathlib import Path
from typing import Dict, Any, Optional, List
from huggingface_hub import HfApi, hf_hub_download, list_repo_files, login
from safetensors.torch import load_file, save_file
from tqdm import tqdm

# --- Memory Efficient Safetensors ---
class MemoryEfficientSafeOpen:
    """
    Reads safetensors metadata and tensors without mmap, keeping RAM usage low.
    """
    def __init__(self, filename):
        self.filename = filename
        self.file = open(filename, "rb")
        self.header, self.header_size = self._read_header()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.file.close()

    def keys(self) -> list[str]:
        return [k for k in self.header.keys() if k != "__metadata__"]

    def metadata(self) -> Dict[str, str]:
        return self.header.get("__metadata__", {})

    def get_tensor(self, key):
        if key not in self.header:
            raise KeyError(f"Tensor '{key}' not found in the file")
        metadata = self.header[key]
        offset_start, offset_end = metadata["data_offsets"]
        self.file.seek(self.header_size + 8 + offset_start)
        tensor_bytes = self.file.read(offset_end - offset_start)
        return self._deserialize_tensor(tensor_bytes, metadata)

    def _read_header(self):
        header_size = struct.unpack("<Q", self.file.read(8))[0]
        header_json = self.file.read(header_size).decode("utf-8")
        return json.loads(header_json), header_size

    def _deserialize_tensor(self, tensor_bytes, metadata):
        dtype_map = {
            "F32": torch.float32, "F16": torch.float16, "BF16": torch.bfloat16,
            "I64": torch.int64, "I32": torch.int32, "I16": torch.int16, "I8": torch.int8,
            "U8": torch.uint8, "BOOL": torch.bool
        }
        dtype = dtype_map[metadata["dtype"]]
        shape = metadata["shape"]
        return torch.frombuffer(tensor_bytes, dtype=torch.uint8).view(dtype).reshape(shape)

# --- Constants & Setup ---
# Use /tmp/temp_tool if possible for better ephemeral handling, 
# or fall back to ./temp_tool in working dir.
try:
    TempDir = Path("/tmp/temp_tool")
    os.makedirs(TempDir, exist_ok=True)
except:
    TempDir = Path("./temp_tool")
    os.makedirs(TempDir, exist_ok=True)

api = HfApi()

def cleanup_temp():
    if TempDir.exists():
        shutil.rmtree(TempDir)
    os.makedirs(TempDir, exist_ok=True)
    gc.collect()

def download_file(input_path, token, filename=None):
    local_path = TempDir / (filename if filename else "model.safetensors")
    if input_path.startswith("http"):
        print(f"Downloading {filename} from URL...")
        try:
            response = requests.get(input_path, stream=True, timeout=30)
            response.raise_for_status()
            with open(local_path, 'wb') as f:
                for chunk in response.iter_content(chunk_size=8192):
                    f.write(chunk)
        except Exception as e: raise ValueError(f"Download failed: {e}")
    else:
        print(f"Downloading {filename} from Hub...")
        if not filename:
            try:
                files = list_repo_files(repo_id=input_path, token=token)
                safetensors = [f for f in files if f.endswith(".safetensors")]
                filename = safetensors[0] if safetensors else "adapter_model.safetensors"
            except: filename = "adapter_model.safetensors"
        
        try:
            hf_hub_download(repo_id=input_path, filename=filename, token=token, local_dir=TempDir, local_dir_use_symlinks=False)
            # Handle default download path logic if specific filename wasn't requested
            if not (TempDir / filename).exists():
                # HF might download to a nested folder structure
                found = list(TempDir.rglob(filename))
                if found: shutil.move(found[0], local_path)
        except Exception as e: raise ValueError(f"Hub download failed: {e}")
    
    return local_path

def get_key_stem(key):
    key = key.replace(".weight", "").replace(".bias", "")
    key = key.replace(".lora_down", "").replace(".lora_up", "")
    key = key.replace(".lora_A", "").replace(".lora_B", "")
    key = key.replace(".alpha", "")
    prefixes = [
        "model.diffusion_model.", "diffusion_model.", "model.", 
        "transformer.", "text_encoder.", "lora_unet_", "lora_te_", "base_model.model."
    ]
    changed = True
    while changed:
        changed = False
        for p in prefixes:
            if key.startswith(p):
                key = key[len(p):]
                changed = True
    return key

# =================================================================================
# TAB 1: GREEDY STREAMING RESHARDER
# =================================================================================

def load_lora_to_memory(lora_path, precision_dtype=torch.bfloat16):
    print(f"Loading LoRA from {lora_path}...")
    state_dict = load_file(lora_path, device="cpu")
    pairs = {} 
    alphas = {}
    for k, v in state_dict.items():
        stem = get_key_stem(k)
        if "alpha" in k:
            alphas[stem] = v.item() if isinstance(v, torch.Tensor) else v
        else:
            if stem not in pairs: pairs[stem] = {}
            if "lora_down" in k or "lora_A" in k:
                pairs[stem]["down"] = v.to(dtype=precision_dtype)
                pairs[stem]["rank"] = v.shape[0]
            elif "lora_up" in k or "lora_B" in k:
                pairs[stem]["up"] = v.to(dtype=precision_dtype)
    for stem in pairs:
        pairs[stem]["alpha"] = alphas.get(stem, float(pairs[stem].get("rank", 1.0)))
    return pairs

class ShardBuffer:
    def __init__(self, max_size_gb, output_dir, output_repo, hf_token):
        self.max_bytes = int(max_size_gb * 1024**3)
        self.output_dir = output_dir
        self.output_repo = output_repo
        self.hf_token = hf_token
        self.buffer = [] # List of (key, bytes, dtype_str, shape)
        self.current_bytes = 0
        self.shard_count = 0
        self.index_map = {}
        
    def add_tensor(self, key, tensor):
        # Convert to bytes
        if tensor.dtype == torch.bfloat16:
            # View as int16 to get raw bytes
            raw_bytes = tensor.view(torch.int16).numpy().tobytes()
            dtype_str = "BF16"
        elif tensor.dtype == torch.float16:
            raw_bytes = tensor.numpy().tobytes()
            dtype_str = "F16"
        else:
            raw_bytes = tensor.numpy().tobytes()
            dtype_str = "F32"
            
        size = len(raw_bytes)
        self.buffer.append({
            "key": key,
            "data": raw_bytes,
            "dtype": dtype_str,
            "shape": tensor.shape
        })
        self.current_bytes += size
        
        # Flush if full
        if self.current_bytes >= self.max_bytes:
            self.flush()
            
    def flush(self):
        if not self.buffer: return
        
        self.shard_count += 1
        # Placeholder filename, will rename later or use sequential numbering
        shard_name = f"model-{self.shard_count:05d}.safetensors" # Suffix to be fixed at end? 
        # Actually, standard is model-00001-of-XXXXX. 
        # Since we don't know total count yet, we use a temp naming scheme, 
        # OR we just use model-00001.safetensors and fix the index.json later.
        # Diffusers accepts model-xxxxx-of-xxxxx.
        # We will use "model-xxxxx.safetensors" and rename locally if needed, 
        # but for simple uploading we can just assume we don't know the total yet.
        # Actually, let's just count up. model-00001.safetensors is fine if we update index.
        
        print(f"Flushing Shard {self.shard_count} ({self.current_bytes / 1024**3:.2f} GB)...")
        
        # Construct Header
        header = {"__metadata__": {"format": "pt"}}
        current_offset = 0
        for item in self.buffer:
            header[item["key"]] = {
                "dtype": item["dtype"],
                "shape": item["shape"],
                "data_offsets": [current_offset, current_offset + len(item["data"])]
            }
            current_offset += len(item["data"])
            self.index_map[item["key"]] = shard_name
            
        header_json = json.dumps(header).encode('utf-8')
        
        # Write File
        out_path = self.output_dir / shard_name
        with open(out_path, 'wb') as f:
            f.write(struct.pack('<Q', len(header_json)))
            f.write(header_json)
            for item in self.buffer:
                f.write(item["data"])
                
        # Upload
        print(f"Uploading {shard_name}...")
        api.upload_file(path_or_fileobj=out_path, path_in_repo=shard_name, repo_id=self.output_repo, token=self.hf_token)
        
        # Cleanup
        os.remove(out_path)
        self.buffer = []
        self.current_bytes = 0
        gc.collect()

def task_merge(hf_token, base_repo, base_subfolder, lora_input, scale, precision, shard_size, output_repo, structure_repo, private, progress=gr.Progress()):
    cleanup_temp()
    login(hf_token)
    
    # 1. Output Setup
    try:
        api.create_repo(repo_id=output_repo, private=private, exist_ok=True, token=hf_token)
    except Exception as e: return f"Error creating repo: {e}"
    
    # Clone structure
    if structure_repo:
        print("Cloning structure...")
        try:
            files = list_repo_files(repo_id=structure_repo, token=hf_token)
            for f in files:
                if not f.endswith(".safetensors") and not f.endswith(".bin"):
                    try:
                        path = hf_hub_download(repo_id=structure_repo, filename=f, token=hf_token)
                        api.upload_file(path_or_fileobj=path, path_in_repo=f, repo_id=output_repo, token=hf_token)
                    except: pass
        except: pass

    # 2. Load LoRA
    dtype = torch.bfloat16 if precision == "bf16" else torch.float16 if precision == "fp16" else torch.float32
    try:
        progress(0.1, desc="Downloading LoRA...")
        lora_path = download_file(lora_input, hf_token, filename="adapter.safetensors")
        lora_pairs = load_lora_to_memory(lora_path, precision_dtype=dtype)
    except Exception as e: return f"Error loading LoRA: {e}"

    # 3. Stream Process
    progress(0.2, desc="Fetching File List...")
    files = list_repo_files(repo_id=base_repo, token=hf_token)
    input_shards = [f for f in files if f.endswith(".safetensors")]
    if base_subfolder:
        input_shards = [f for f in input_shards if f.startswith(base_subfolder)]
        
    if not input_shards: return "No base safetensors found."
    
    # Sort shards to ensure deterministic processing order
    input_shards.sort()
    
    buffer = ShardBuffer(shard_size, TempDir, output_repo, hf_token)
    
    for i, shard_file in enumerate(input_shards):
        progress(0.2 + (0.7 * i / len(input_shards)), desc=f"Processing {shard_file}")
        print(f"Downloading {shard_file}...")
        
        local_shard = hf_hub_download(repo_id=base_repo, filename=shard_file, token=hf_token, local_dir=TempDir)
        
        # Process tensors
        with MemoryEfficientSafeOpen(local_shard) as f:
            keys = f.keys()
            for k in keys:
                v = f.get_tensor(k)
                
                # MERGE LOGIC
                base_stem = get_key_stem(k)
                lora_keys = set(lora_pairs.keys())
                match = None
                
                if base_stem in lora_keys:
                    match = lora_pairs[base_stem]
                else:
                    if "to_q" in base_stem:
                        qkv_stem = base_stem.replace("to_q", "qkv")
                        if qkv_stem in lora_keys: match = lora_pairs[qkv_stem]
                    elif "to_k" in base_stem:
                        qkv_stem = base_stem.replace("to_k", "qkv")
                        if qkv_stem in lora_keys: match = lora_pairs[qkv_stem]
                    elif "to_v" in base_stem:
                        qkv_stem = base_stem.replace("to_v", "qkv")
                        if qkv_stem in lora_keys: match = lora_pairs[qkv_stem]

                if match and "down" in match and "up" in match:
                    down = match["down"]
                    up = match["up"]
                    alpha = match["alpha"]
                    rank = match["rank"]
                    scaling = scale * (alpha / rank)
                    
                    if len(v.shape) == 4 and len(down.shape) == 2:
                        down = down.unsqueeze(-1).unsqueeze(-1)
                        up = up.unsqueeze(-1).unsqueeze(-1)
                    
                    try:
                        if len(up.shape) == 4:
                            delta = (up.squeeze() @ down.squeeze()).reshape(up.shape[0], down.shape[1], 1, 1)
                        else:
                            delta = up @ down
                    except:
                        delta = up.T @ down 
                    
                    delta = delta * scaling
                    
                    # Slicing
                    valid_delta = True
                    if delta.shape == v.shape:
                        pass
                    elif delta.shape[0] == v.shape[0] * 3:
                        chunk = v.shape[0]
                        if "to_q" in k: delta = delta[0:chunk, ...]
                        elif "to_k" in k: delta = delta[chunk:2*chunk, ...]
                        elif "to_v" in k: delta = delta[2*chunk:, ...]
                        else: valid_delta = False
                    elif delta.numel() == v.numel():
                        delta = delta.reshape(v.shape)
                    else:
                        valid_delta = False
                        
                    if valid_delta:
                        v = v.to(dtype)
                        delta = delta.to(dtype)
                        v.add_(delta)
                        del delta

                # Add to buffer
                if v.dtype != dtype: v = v.to(dtype)
                buffer.add_tensor(k, v)
                del v
        
        # Cleanup Input Shard immediately
        os.remove(local_shard)
        gc.collect()

    # Final Flush
    buffer.flush()
    
    # Renaming logic (Retroactive):
    # Since we uploaded as model-00001.safetensors, but now we know total count...
    # Actually, Diffusers is fine with model-00001.safetensors format as long as index.json matches.
    # We just need to upload the index.
    
    print("Uploading Index...")
    index_data = {"metadata": {"total_size": 0}, "weight_map": buffer.index_map}
    with open(TempDir / "model.safetensors.index.json", "w") as f:
        json.dump(index_data, f, indent=4)
    api.upload_file(path_or_fileobj=TempDir / "model.safetensors.index.json", path_in_repo="model.safetensors.index.json", repo_id=output_repo, token=hf_token)
    
    cleanup_temp()
    return f"Done! Merged into {buffer.shard_count} shards at {output_repo}"

# =================================================================================
# TAB 2: EXTRACT LORA
# =================================================================================

def extract_lora_layer_by_layer(model_org, model_tuned, rank, clamp):
    org = MemoryEfficientSafeOpen(model_org)
    tuned = MemoryEfficientSafeOpen(model_tuned)
    lora_sd = {}
    print("Calculating diffs...")
    for key in tqdm(org.keys()):
        if key not in tuned.keys(): continue
        mat_org = org.get_tensor(key).float()
        mat_tuned = tuned.get_tensor(key).float()
        diff = mat_tuned - mat_org
        if torch.max(torch.abs(diff)) < 1e-4: continue
        
        out_dim, in_dim = diff.shape[:2]
        r = min(rank, in_dim, out_dim)
        is_conv = len(diff.shape) == 4
        if is_conv: diff = diff.flatten(start_dim=1)
            
        try:
            U, S, Vh = torch.linalg.svd(diff, full_matrices=False)
            U, S, Vh = U[:, :r], S[:r], Vh[:r, :]
            U = U @ torch.diag(S)
            dist = torch.cat([U.flatten(), Vh.flatten()])
            hi_val = torch.quantile(dist, clamp)
            U = U.clamp(-hi_val, hi_val)
            Vh = Vh.clamp(-hi_val, hi_val)
            if is_conv:
                U = U.reshape(out_dim, r, 1, 1)
                Vh = Vh.reshape(r, in_dim, mat_org.shape[2], mat_org.shape[3])
            else:
                U = U.reshape(out_dim, r)
                Vh = Vh.reshape(r, in_dim)
            stem = key.replace(".weight", "")
            lora_sd[f"{stem}.lora_up.weight"] = U
            lora_sd[f"{stem}.lora_down.weight"] = Vh
            lora_sd[f"{stem}.alpha"] = torch.tensor(r).float()
        except: pass
    out = TempDir / "extracted.safetensors"
    save_file(lora_sd, out)
    return str(out)

def task_extract(hf_token, org, tun, rank, out):
    cleanup_temp()
    login(hf_token)
    try:
        p1 = download_file(org, hf_token, filename="org.safetensors")
        p2 = download_file(tun, hf_token, filename="tun.safetensors")
        f = extract_lora_layer_by_layer(p1, p2, int(rank), 0.99)
        api.create_repo(repo_id=out, exist_ok=True, token=hf_token)
        api.upload_file(path_or_fileobj=f, path_in_repo="extracted.safetensors", repo_id=out, token=hf_token)
        return "Done"
    except Exception as e: return f"Error: {e}"

# =================================================================================
# TAB 3 & 4
# =================================================================================

def task_merge_adapters(hf_token, urls, beta, out_repo):
    cleanup_temp()
    login(hf_token)
    try:
        paths = [download_file(u.strip(), hf_token, filename=f"a_{i}.safetensors") for i,u in enumerate(urls.split(",")) if u.strip()]
        if not paths: return "No files"
        base = load_file(paths[0], device="cpu")
        for k in base:
            if base[k].dtype.is_floating_point: base[k] = base[k].float()
        for p in paths[1:]:
            c = load_file(p, device="cpu")
            for k in base:
                if k in c and "alpha" not in k:
                    base[k] = base[k] * beta + c[k].float() * (1-beta)
        out = TempDir / "merged_adapters.safetensors"
        save_file(base, out)
        api.create_repo(repo_id=out_repo, exist_ok=True, token=hf_token)
        api.upload_file(path_or_fileobj=out, path_in_repo="merged_adapters.safetensors", repo_id=out_repo, token=hf_token)
        return "Done"
    except Exception as e: return f"Error: {e}"

def task_resize(hf_token, lora, rank, out):
    return "See previous versions for full code."

# =================================================================================
# UI
# =================================================================================

css = ".container { max-width: 900px; margin: auto; }"

with gr.Blocks() as demo:
    gr.Markdown("# 🧰 Universal LoRA Toolkit V12 (Greedy Streaming)")
    
    with gr.Tabs():
        with gr.Tab("Merge + Reshard"):
            t1_token = gr.Textbox(label="Token", type="password")
            t1_base = gr.Textbox(label="Base Repo", value="ostris/Z-Image-De-Turbo")
            t1_sub = gr.Textbox(label="Subfolder", value="transformer")
            t1_lora = gr.Textbox(label="LoRA")
            with gr.Row():
                t1_scale = gr.Slider(label="Scale", value=1.0)
                t1_prec = gr.Radio(["bf16", "fp16", "float32"], value="bf16", label="Precision")
                t1_shard = gr.Slider(label="Shard Size (GB)", value=2.0, minimum=0.5, maximum=10.0, step=0.5)
            t1_out = gr.Textbox(label="Output")
            t1_struct = gr.Textbox(label="Structure Repo", value="Tongyi-MAI/Z-Image-Turbo")
            t1_priv = gr.Checkbox(label="Private", value=True)
            t1_btn = gr.Button("Merge & Reshard")
            t1_res = gr.Textbox(label="Result")
            t1_btn.click(task_merge, [t1_token, t1_base, t1_sub, t1_lora, t1_scale, t1_prec, t1_shard, t1_out, t1_struct, t1_priv], t1_res)

        with gr.Tab("Extract"):
            t2_token = gr.Textbox(label="Token", type="password")
            t2_org = gr.Textbox(label="Original")
            t2_tun = gr.Textbox(label="Tuned")
            t2_rank = gr.Number(label="Rank", value=32)
            t2_out = gr.Textbox(label="Output")
            t2_btn = gr.Button("Extract")
            t2_res = gr.Textbox(label="Result")
            t2_btn.click(task_extract, [t2_token, t2_org, t2_tun, t2_rank, t2_out], t2_res)

        with gr.Tab("Merge Adapters"):
            t3_token = gr.Textbox(label="Token", type="password")
            t3_urls = gr.Textbox(label="URLs")
            t3_beta = gr.Slider(label="Beta", value=0.9)
            t3_out = gr.Textbox(label="Output")
            t3_btn = gr.Button("Merge")
            t3_res = gr.Textbox(label="Result")
            t3_btn.click(task_merge_adapters, [t3_token, t3_urls, t3_beta, t3_out], t3_res)

if __name__ == "__main__":
    demo.queue().launch(css=css, ssr_mode=False)