import os
import re
import shutil
import csv
import sys
import json
from pathlib import Path

import torch
import torch.distributed as dist
import math
from dotenv import load_dotenv

import hydra
from omegaconf import OmegaConf
from safetensors.torch import load_file

# Adjust imports to match your project structure
from lmr.config import initialize_config
from lmr.tokenizer import Tokenizer
from lmr.models import get_model 
from lmr.data import get_dataset_splits
from lmr.checkpointing import Checkpointing
from lmr.utils.seed import set_seed
from lmr.training import Bert_Trainer, Trainer
from lmr.generation import Generator
from lmr.benchmark import Benchmark
from lmr.ddp import unwrap_model

DATASET_DIR = Path("datasets")
CHECKPOINT_DIR = Path("/work/jf381/checkpoints")
BENCHMARK_DIR = Path("output")

# =============================================================================
# Helper Functions: Robust Loading
# =============================================================================

def load_weight_data(path_obj, device="cpu"):
    """核心：支持文件夹(sharded safetensors), 单文件(.safetensors), 或 (.pt)"""
    model_state = {}
    
    if path_obj.is_dir():
        index_file = path_obj / "model.safetensors.index.json"
        if index_file.exists():
            print(f"🔹 Detected sharded safetensors folder: {path_obj.name}")
            with open(index_file, 'r') as f:
                index_data = json.load(f)
            weight_map = index_data.get("weight_map", {})
            shards = set(weight_map.values())
            for shard_name in shards:
                shard_path = path_obj / shard_name
                model_state.update(load_file(str(shard_path), device=str(device)))
            return model_state
        else:
            # 如果是文件夹但没有 index，尝试找文件夹里的第一个权重文件
            possible = list(path_obj.glob("*.safetensors")) + list(path_obj.glob("*.pt"))
            if not possible: return None
            path_obj = possible[0]

    if path_obj.suffix == ".safetensors":
        print(f"🔹 Loading single safetensors: {path_obj.name}")
        return load_file(str(path_obj), device=str(device))
    else:
        print(f"🔹 Loading pickle (.pt): {path_obj.name}")
        ckpt = torch.load(path_obj, map_location=device)
        return ckpt.get("model", ckpt.get("state_dict", ckpt))

def load_state_dict_robust(model, checkpoint_path, strict=False):
    """
    真正的强力加载：自动处理前缀(module., model.)并匹配分片权重。
    """
    path_obj = Path(checkpoint_path)
    if not path_obj.exists():
        print(f"❌ Path not found: {checkpoint_path}")
        return False

    try:
        # 1. 获取权重数据
        model_state = load_weight_data(path_obj)
        if model_state is None: return False

        # 2. 规范化 Key (移除 DDP 或 HF 带来的前缀)
        ckpt_keys_map = {}
        for k in model_state.keys():
            clean_k = k.replace("module.", "").replace("_orig_mod.", "").replace("model.", "")
            ckpt_keys_map[clean_k] = k

        # 3. 准备目标模型 (Handle unwrapped model)
        target_model = unwrap_model(model)
        target_state = target_model.state_dict()
        
        filtered_state = {}
        matched_count = 0

        for k_target, v_target in target_state.items():
            # 同样规范化目标 key
            k_target_clean = k_target.replace("module.", "").replace("_orig_mod.", "").replace("model.", "")
            
            if k_target_clean in ckpt_keys_map:
                real_ckpt_key = ckpt_keys_map[k_target_clean]
                v_ckpt = model_state[real_ckpt_key]
                
                if v_ckpt.shape == v_target.shape:
                    filtered_state[k_target] = v_ckpt
                    matched_count += 1

        # 4. 执行加载
        msg = target_model.load_state_dict(filtered_state, strict=strict)
        print(f"✅ Success! Loaded {matched_count} parameters. Status: {msg}")
        return True
    except Exception as e:
        print(f"❌ Failed to load checkpoint: {e}")
        import traceback
        traceback.print_exc()
        return False

def setup_model_and_tokenizer(config):
    print(f"🔧 Initializing Tokenizer: {config.tokenizer_base}")
    tokenizer = Tokenizer(config.tokenizer_base)
    print(f"🔧 Initializing Model: {config.model}")
    model = get_model(config.model, tokenizer.vocab_size, tokenizer=tokenizer)
    return tokenizer, model

# =============================================================================
# Mode Logics
# =============================================================================

def run_generation_task(config, model, tokenizer, output_dir, ckpt_name_tag=""):
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model.to(device)
    model.eval()
    
    generator = Generator(config, model, tokenizer, device=device, output_dir=output_dir)
    results = generator.generate()
    
    metrics = {}
    if results and "verification" in results:
        metrics = {
            'acc': results['verification']['accuracy'],
            'correct': results['verification']['correct'],
            'total': results['verification']['total']
        }
        if results.get('token_accuracies'):
            import numpy as np
            metrics['token_acc'] = np.mean(results['token_accuracies'])

        if ckpt_name_tag:
            for fname in ["gsm8k_metrics.txt", "gsm8k_generations.txt"]:
                src = output_dir / fname
                if src.exists():
                    dst = output_dir / f"{src.stem}_{ckpt_name_tag}{src.suffix}"
                    shutil.move(src, dst)
    return metrics

def generate(config):
    """Mode: generate (单次运行，支持加载最近的文件夹或.pt权重)"""
    tokenizer, model = setup_model_and_tokenizer(config)
    ckpt_dir = CHECKPOINT_DIR / config.checkpoint_name
    
    load_mode = getattr(config.benchmark, "checkpoint_mode", "recent")
    print(f"🔍 Looking for [{load_mode}] weights in {ckpt_dir}...")

    # 识别所有可能的 checkpoint (包括文件夹和 .pt 文件)
    all_items = list(ckpt_dir.iterdir())
    checkpoints = [f for f in all_items if ("optim" not in f.name and "sched" not in f.name and f.name != "metrics_summary.csv")]

    def sort_key(f):
        nums = re.findall(r'\d+', f.name)
        return int(nums[-1]) if nums else 0
    checkpoints.sort(key=sort_key)

    target_ckpt = None
    if load_mode == "best":
        best_candidates = [f for f in checkpoints if "best" in f.name.lower()]
        target_ckpt = best_candidates[0] if best_candidates else (checkpoints[-1] if checkpoints else None)
    else:
        target_ckpt = checkpoints[-1] if checkpoints else None

    if target_ckpt:
        print(f"🚀 Found target: {target_ckpt.name}")
        success = load_state_dict_robust(model, target_ckpt)
        if not success:
            print("⚠️ Load failed, check file integrity.")
    else:
        print(f"❌ No checkpoints found in {ckpt_dir}")

    run_generation_task(config, model, tokenizer, output_dir=ckpt_dir)

def generate_all(config):
    """Mode: generate_all (扫描目录并运行所有 checkpoitns)"""
    tokenizer, model = setup_model_and_tokenizer(config)
    ckpt_dir = CHECKPOINT_DIR / config.checkpoint_name
    
    # 查找所有子目录或 .pt 文件
    all_items = list(ckpt_dir.iterdir())
    checkpoints = [f for f in all_items if ("optim" not in f.name and "sched" not in f.name and f.name != "metrics_summary.csv")]
    
    def sort_key(f):
        nums = re.findall(r'\d+', f.name)
        return int(nums[-1]) if nums else 0
    checkpoints.sort(key=sort_key)
    
    print(f"\n🔎 Found {len(checkpoints)} checkpoints.")
    summary_path = ckpt_dir / "metrics_summary.csv"
    
    with open(summary_path, mode='w', newline='') as f:
        writer = csv.writer(f)
        writer.writerow(["checkpoint", "accuracy", "token_accuracy", "correct", "total"])

    for ckpt_path in checkpoints:
        print(f"\n{'-'*40}\nProcessing: {ckpt_path.name}\n{'-'*40}")
        if load_state_dict_robust(model, ckpt_path):
            metrics = run_generation_task(config, model, tokenizer, ckpt_dir, ckpt_name_tag=ckpt_path.name)
            if metrics:
                with open(summary_path, mode='a', newline='') as f:
                    writer = csv.writer(f)
                    writer.writerow([
                        ckpt_path.name,
                        f"{metrics.get('acc', 0):.4f}",
                        f"{metrics.get('token_acc', 0):.4f}",
                        metrics.get('correct', 0),
                        metrics.get('total', 0)
                    ])

def train_model(config):
    if torch.cuda.is_available() and torch.cuda.device_count() > 1:
        if not dist.is_initialized():
            dist.init_process_group(backend="nccl")
            torch.cuda.set_device(dist.get_rank() % torch.cuda.device_count())

    tokenizer, model = setup_model_and_tokenizer(config)
    tokenized_dataset_dir = DATASET_DIR / config.tokenizer_base
    splits = get_dataset_splits(config.dataset, 1024, tokenized_dataset_dir)
    checkpointing = Checkpointing(model, CHECKPOINT_DIR / config.checkpoint_name)
    
    if "bert" in str(config.model).lower():
        from lmr.training import Bert_Trainer
        trainer = Bert_Trainer(config.training, model, tokenizer, splits, checkpointing, None)
    else:
        trainer = Trainer(config.training, model, tokenizer, splits, checkpointing)
    trainer.train()

# =============================================================================
# Main Entry Point
# =============================================================================

@hydra.main(config_path="config", config_name="config", version_base="1.3")
def main(config):
    load_dotenv()
    set_seed(config)
    initialize_config(config)
    
    mode = config.mode
    if mode == "train":
        train_model(config)
    elif mode == "generate":
        generate(config)
    elif mode == "generate_all":
        generate_all(config)
    elif mode == "benchmark":
        tokenizer, model = setup_model_and_tokenizer(config)
        checkpointing = Checkpointing(model, CHECKPOINT_DIR / config.checkpoint_name)
        benchmarking = Benchmark(config.benchmark, model, tokenizer, checkpointing, BENCHMARK_DIR / config.checkpoint_name)
        benchmarking.run_benchmarks()
    else:
        print(f"❌ Unknown mode: {mode}")

if __name__ == "__main__":
    main()