#!/usr/bin/env python3
"""True streaming 4-bit NF4 quantization - one shard at a time."""

import argparse
import gc
import torch
import concurrent.futures
from pathlib import Path
from safetensors.torch import load_file, save_file
from bitsandbytes.functional import quantize_nf4
from transformers import AutoConfig


def quantize_weight_nf4(weight: torch.Tensor):
    """Quantize a single weight tensor to NF4 using bitsandbytes functional API."""
    if weight.dim() != 2:
        return weight, None

    # quantize_nf4 returns (quantized_tensor, quant_state)
    qweight, quant_state = quantize_nf4(
        weight,
        blocksize=64,
        compress_statistics=True,
    )
    return qweight, quant_state


def streaming_quantize(model_path: str, output_path: str):
    print(f"Streaming NF4 quantization: {model_path}")
    output_path = Path(output_path)
    output_path.mkdir(parents=True, exist_ok=True)

    import glob
    shards = sorted(glob.glob(f"{model_path}/*.safetensors"))
    
    # Rename existing 0-indexed files to 1-indexed
    for existing in output_path.glob("*.safetensors"):
        parts = existing.name.split("-")
        if len(parts) >= 3 and existing.name.startswith("model-00000-"):
            num = int(parts[1])
            new_name = f"model-{num+1:05d}-of-{len(shards):05d}.safetensors"
            new_path = output_path / new_name
            existing.rename(new_path)
            print(f"Renamed: {existing.name} -> {new_name}")
    
    print(f"Found {len(shards)} shards\n")
    
    config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)

    # Process continuously: max 4 shards per GPU, add to whichever GPU has room
    max_per_gpu = 4
    num_gpus = 2
    
    def process_shard(idx, shard_file, gpu_id):
        """Process a single shard (called per-GPU)."""
        shard_name = f"model-{idx+1:05d}-of-{len(shards):05d}.safetensors"
        if (output_path / shard_name).exists():
            return
        
        print(f"[{idx+1}/{len(shards)}] {Path(shard_file).name} (GPU {gpu_id})")
        
        # Load to assigned GPU
        state_dict = load_file(shard_file, device=f"cuda:{gpu_id}")
        
        weight_keys = [
            k for k, v in state_dict.items()
            if "weight" in k and isinstance(v, torch.Tensor) and v.dim() == 2
        ]
        
        print(f"   Quantizing {len(weight_keys)} tensors...")
        
        for key in weight_keys:
            try:
                weight = state_dict[key]
                qweight, qstate = quantize_weight_nf4(weight)
                state_dict[key] = qweight
                del weight, qweight, qstate
                gc.collect()
                torch.cuda.empty_cache()
            except Exception as e:
                print(f"   Warning: Failed on {key}: {e}")
                gc.collect()
                torch.cuda.empty_cache()
        
        # Move to CPU and save
        state_dict = {
            k: v.cpu() if isinstance(v, torch.Tensor) else v
            for k, v in state_dict.items()
        }
        
        save_file(state_dict, output_path / shard_name)
        print(f"   ✓ Saved {shard_name}")
        
        del state_dict
        gc.collect()
        torch.cuda.empty_cache()
    
    # Find first unsaved shard
    start_idx = 0
    for idx in range(len(shards)):
        shard_name = f"model-{idx+1:05d}-of-{len(shards):05d}.safetensors"
        if not (output_path / shard_name).exists():
            start_idx = idx
            break
    
    print(f"Starting from shard {start_idx+1}/16\n")
    
    # Process continuously: submit to GPU with fewer active tasks
    import threading
    gpu_locks = [threading.Lock() for _ in range(num_gpus)]
    gpu_counts = [0] * num_gpus
    
    def get_next_gpu():
        """Get the GPU with fewest active tasks (max 4 per GPU)."""
        with gpu_locks[0]:
            with gpu_locks[1]:
                if gpu_counts[0] <= gpu_counts[1] and gpu_counts[0] < max_per_gpu:
                    return 0
                elif gpu_counts[1] < max_per_gpu:
                    return 1
                else:
                    return None  # Both GPUs at max
    
    # Create a queue of unprocessed shards
    from collections import deque
    remaining = deque(range(start_idx, len(shards)))
    
    with concurrent.futures.ThreadPoolExecutor(max_workers=max_per_gpu * num_gpus) as executor:
        futures = []
        
        def submit_next():
            """Submit next shard to available GPU."""
            gpu_id = get_next_gpu()
            if gpu_id is None or not remaining:
                return
            
            idx = remaining.popleft()
            shard_file = shards[idx]
            
            with gpu_locks[gpu_id]:
                gpu_counts[gpu_id] += 1
            
            future = executor.submit(process_shard, idx, shard_file, gpu_id)
            futures.append(future)
            
            # When this future completes, release GPU slot and submit next
            def on_complete(f):
                with gpu_locks[gpu_id]:
                    gpu_counts[gpu_id] -= 1
                submit_next()  # Try to submit next shard
            
            future.add_done_callback(on_complete)
        
        # Start submitting
        while remaining:
            gpu_id = get_next_gpu()
            if gpu_id is None:
                break
            submit_next()
        
        # Wait for all to complete
        concurrent.futures.wait(futures)

    config.save_pretrained(output_path)
    print(f"\n✅ Done → {output_path}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-path", type=str, default="/data/models/Ornith-1.0-35B")
    parser.add_argument("--output-path", type=str, default="/data/models/Ornith-1.0-35B-nf4")
    args = parser.parse_args()

    streaming_quantize(args.model_path, args.output_path)