import os
# T4 ke liye robust env
os.environ["NCCL_P2P_DISABLE"] = "1"
os.environ["NCCL_IB_DISABLE"] = "1"
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
# T4 optimization
os.environ["TOKENIZERS_PARALLELISM"] = "false"

import math
import time
import argparse
import logging
from pathlib import Path

import torch
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
# FIX 1: Naya amp API - T4 ke liye must (bf16 nahi, fp16)
from torch.amp import GradScaler, autocast
import numpy as np
from transformers import AutoTokenizer
from huggingface_hub import hf_hub_download, HfApi
from huggingface_hub.utils import EntryNotFoundError, HfHubHTTPError

try:
    from kaggle_secrets import UserSecretsClient
    HAS_KAGGLE_SECRETS = True
except ImportError:
    HAS_KAGGLE_SECRETS = False

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)

def get_hf_token():
    token = os.environ.get("HF_TOKEN")
    if token:
        token = token.strip()
        logger.info("Loaded HF_TOKEN from environment.")
        return token
    if HAS_KAGGLE_SECRETS:
        try:
            token = UserSecretsClient().get_secret("HF_TOKEN")
            if token:
                token = token.strip()
                os.environ["HF_TOKEN"] = token
                logger.info("Loaded HF_TOKEN from kaggle_secrets.")
                return token
        except Exception as e:
            logger.warning(f"Failed to get HF_TOKEN from kaggle_secrets: {e}")
    raise ValueError("HF_TOKEN not found in environment or kaggle_secrets.")

def setup_ddp():
    if "RANK" not in os.environ or "WORLD_SIZE" not in os.environ:
        return 0, 0, 1, torch.device("cuda" if torch.cuda.is_available() else "cpu"), False
    dist.init_process_group("nccl")
    local_rank = int(os.environ["LOCAL_RANK"])
    device = torch.device(f"cuda:{local_rank}")
    torch.cuda.set_device(device)
    return int(os.environ["RANK"]), local_rank, int(os.environ["WORLD_SIZE"]), device, True

def load_tokenizer_robust(hf_token: str, model_repo: str, subfolder: str, is_main: bool):
    if is_main: logger.info("Loading tokenizer...")
    local_candidates = [
        "/kaggle/working/ViuRec/tokenizer",
        "/kaggle/input/viuai-500m-tokenizer",
        "./tokenizer",
        "./viuai-500m-tokenizer"
    ]
    
    for fast in [True, False]:
        for p_str in local_candidates:
            p = Path(p_str)
            if (p / "tokenizer.json").exists() or (p / "vocab.json").exists() or (p / "tokenizer.model").exists():
                try:
                    tok = AutoTokenizer.from_pretrained(str(p), local_files_only=True, use_fast=fast, trust_remote_code=True)
                    if is_main: logger.info(f"Loaded tokenizer from {p} fast={fast}")
                    return tok
                except Exception as e:
                    if is_main: logger.warning(f"Local tokenizer {p} fast={fast} failed: {e}")
                    continue

        # HF se try - pehle ViuRec tokenizer subfolder, phir alag tokenizer repo
        hf_repos_to_try = [model_repo, "ViuAI/viuai-500m-tokenizer", "ViuAI/ViuRec"]
        for repo_try in hf_repos_to_try:
            for attempt_subfolder in [subfolder, None]:
                try:
                    if attempt_subfolder:
                        tok = AutoTokenizer.from_pretrained(repo_try, subfolder=attempt_subfolder, token=hf_token, use_fast=fast, trust_remote_code=True)
                    else:
                        tok = AutoTokenizer.from_pretrained(repo_try, token=hf_token, use_fast=fast, trust_remote_code=True)
                    if is_main: logger.info(f"Loaded tokenizer from HF {repo_try} subfolder={attempt_subfolder} fast={fast}")
                    return tok
                except Exception as e:
                    err = str(e).lower()
                    if "sentencepiece" in err or "tiktoken" in err or "protobuf" in err:
                        if is_main: logger.warning(f"Tokenizer needs sentencepiece/tiktoken/protobuf: {e}")
                        # fast=False pe dobara try karo
                        continue
                    if is_main: logger.debug(f"HF tokenizer try failed {repo_try} fast={fast} subfolder={attempt_subfolder}: {e}")
                    continue

    if is_main:
        logger.warning("Tokenizer failed, using dummy vocab 64000")
    return None

def load_checkpoint(model, optimizer, scaler, api, token, repo_id, filename="checkpoints/latest_chk.pt", strict=False):
    # FIX 2: Path bug fix - hf_hub_download local_dir="." + filename="checkpoints/..." => ./checkpoints/...
    # Toh local_path bhi wahi hona chahiye
    local_path = Path(filename)  # checkpoints/latest_chk.pt
    local_path.parent.mkdir(parents=True, exist_ok=True)
    try:
        logger.info(f"Checking for existing checkpoint on HF: {filename}")
        # download return karta hai actual local path
        downloaded_path = hf_hub_download(repo_id=repo_id, filename=filename, local_dir=".", token=token)
        logger.info(f"Downloaded checkpoint to {downloaded_path}")
        # Security: trusted env me hi weights_only=False use karo
        ckpt = torch.load(downloaded_path, map_location="cpu", weights_only=False)
        missing, unexpected = model.load_state_dict(ckpt["model"], strict=strict)
        if missing or unexpected:
            logger.warning(f"Checkpoint loaded with Missing: {len(missing)} | Unexpected: {len(unexpected)}")
        else:
            logger.info("Model weights loaded successfully.")
            
        try:
            optimizer.load_state_dict(ckpt["optimizer"])
        except Exception as e:
            logger.warning(f"Optimizer state load failed, continuing: {e}")
        if "scaler" in ckpt and scaler is not None:
            try:
                scaler.load_state_dict(ckpt["scaler"])
            except Exception as e:
                logger.warning(f"Scaler state load failed: {e}")
            
        step = ckpt.get("step", 0)
        logger.info(f"Resuming from step {step}")
        return step
    except (EntryNotFoundError, HfHubHTTPError):
        logger.info("No checkpoint found on HF. Starting from scratch.")
        return 0
    except Exception as e:
        logger.error(f"Error loading checkpoint: {e}")
        return 0

class ShardPool:
    def __init__(self, api, repo_id, work_dir, token, vocab_size, max_pool=3):
        self.api = api
        self.repo_id = repo_id
        self.work_dir = Path(work_dir)
        self.work_dir.mkdir(parents=True, exist_ok=True)
        self.token = token
        self.max_pool = max_pool
        self.pool = []
        self.available_files = []
        self.vocab_size = vocab_size
        self._refresh_file_list()
        
    def _refresh_file_list(self):
        try:
            files = self.api.list_repo_files(repo_id=self.repo_id, repo_type="dataset", token=self.token)
        except Exception as e:
            logger.warning(f"Failed to list dataset repo {self.repo_id}, trying as model repo. Error: {e}")
            files = self.api.list_repo_files(repo_id=self.repo_id, repo_type="model", token=self.token)
        self.available_files = [f for f in files if f.startswith("shards/") and f.endswith(".bin") and not f.endswith("_val.bin")]
        np.random.shuffle(self.available_files)
        logger.info(f"Found {len(self.available_files)} training shards")
        
    def fill(self, ctx_len):
        max_attempts = self.max_pool * 5
        attempts = 0
        while len(self.pool) < self.max_pool:
            if attempts >= max_attempts:
                raise RuntimeError("Failed to fill ShardPool: too many invalid shards.")
            if not self.available_files:
                self._refresh_file_list()
                if not self.available_files:
                    raise RuntimeError("No shard files available")
            f = self.available_files.pop()
            local_p = hf_hub_download(repo_id=self.repo_id, repo_type="dataset", filename=f, local_dir=self.work_dir, token=self.token)
            
            dtype = np.uint16 if self.vocab_size <= 65535 else np.uint32
            mmap_arr = np.memmap(local_p, dtype=dtype, mode="r")
            
            if len(mmap_arr) < ctx_len + 1:
                logger.warning(f"Shard {f} too small ({len(mmap_arr)}). Discarding.")
                del mmap_arr
                Path(local_p).unlink(missing_ok=True)
                attempts += 1
                continue

            # FIX 3: Vocab validation - token vocab se bada toh shard corrupt
            if len(mmap_arr) > 1000:
                sample_max = int(mmap_arr[:1000].max())
                if sample_max >= self.vocab_size:
                    logger.warning(f"Shard {f} has token {sample_max} >= vocab {self.vocab_size}. Discarding.")
                    del mmap_arr
                    Path(local_p).unlink(missing_ok=True)
                    attempts += 1
                    continue
                
            self.pool.append((local_p, mmap_arr))
            logger.info(f"Loaded shard {f} ({len(mmap_arr)} tokens)")
            
    def sample_batch(self, batch_size, ctx_len, rng):
        xs, ys = [], []
        while len(xs) < batch_size:
            idx = rng.integers(0, len(self.pool))
            _, tokens = self.pool[idx]
            max_start = len(tokens) - ctx_len - 1
            if max_start <= 0:
                continue
            i = rng.integers(0, max_start)
            xs.append(np.array(tokens[i:i+ctx_len], dtype=np.int64))
            ys.append(np.array(tokens[i+1:i+1+ctx_len], dtype=np.int64))
        return np.stack(xs), np.stack(ys)
        
    def rotate_one(self):
        if not self.pool: return
        old_path, old_mmap = self.pool.pop(0)
        del old_mmap
        try:
            Path(old_path).unlink(missing_ok=True)
            logger.info(f"Rotated out shard {Path(old_path).name}")
        except Exception as e:
            logger.warning(f"Failed to delete {old_path}: {e}")

def main():
    parser = argparse.ArgumentParser(description="ViuResonance Training - T4 Optimized")
    parser.add_argument("--max-steps", type=int, default=10000)
    parser.add_argument("--micro-batch", type=int, default=1, help="T4 16GB ke liye 1 rakho")
    parser.add_argument("--grad-accum", type=int, default=16)
    parser.add_argument("--eval-every", type=int, default=500)
    parser.add_argument("--save-every", type=int, default=500)
    parser.add_argument("--rotate-every", type=int, default=500)
    parser.add_argument("--lr", type=float, default=3e-4)
    parser.add_argument("--weight-decay", type=float, default=0.1)
    parser.add_argument("--ctx-len", type=int, default=1024, help="T4 pe 1024 safe, 2048 OOM de sakta hai")
    parser.add_argument("--vocab-size", type=int, default=None)
    args = parser.parse_args()

    # FIX 4: T4 TF32 optimization
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True
    torch.backends.cudnn.benchmark = True

    token = get_hf_token()
    rank, local_rank, world_size, device, is_ddp = setup_ddp()
    is_main = (rank == 0)
    
    # FIX 5: Seed properly for CUDA
    seed = 42 + rank
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    rng = np.random.default_rng(seed)

    # Load tokenizer
    tokenizer = load_tokenizer_robust(token, "ViuAI/ViuRec", "tokenizer", is_main)
    vocab_size = args.vocab_size or (len(tokenizer) if tokenizer is not None else 64000)
    if is_main:
        logger.info(f"Using vocab_size={vocab_size}, ctx_len={args.ctx_len}, device={device}, is_ddp={is_ddp}")

    # FIX: Dono naam support karo - model.py ya model_fixed_T4.py
    try:
        from model_fixed_T4 import ViuResonance100M, ModelConfig
    except ModuleNotFoundError:
        try:
            from model import ViuResonance100M, ModelConfig
        except ModuleNotFoundError:
            # Kaggle me file /kaggle/working/code/model.py ho sakti hai
            import importlib.util, sys
            for cand in ["./model_fixed_T4.py", "./model.py", "/kaggle/working/code/model_fixed_T4.py", "/kaggle/working/code/model.py"]:
                p = Path(cand)
                if p.exists():
                    spec = importlib.util.spec_from_file_location("model", str(p))
                    mod = importlib.util.module_from_spec(spec)
                    sys.modules["model"] = mod
                    spec.loader.exec_module(mod)
                    ViuResonance100M = mod.ViuResonance100M
                    ModelConfig = mod.ModelConfig
                    break
            else:
                raise ModuleNotFoundError("model.py / model_fixed_T4.py nahi mila - dono files same folder me rakho")

    config = ModelConfig(vocab_size=vocab_size, chunk_size=256) # T4 ke liye 256
    unwrapped_model = ViuResonance100M(config).to(device)
    
    # FIX 6: torch.compile pehle, DDP baad me - T4 par compile thoda heavy hai, fallback rakha hai
    try:
        unwrapped_model = torch.compile(unwrapped_model, mode="default") # T4 ke liye reduce-overhead bhi try kar sakte ho
        if is_main: logger.info("Compiled model successfully (before DDP).")
        compiled = True
    except Exception as e:
        if is_main: logger.warning(f"Failed to compile model: {e}. Falling back to eager.")
        compiled = False

    if is_ddp:
        model = DDP(unwrapped_model, device_ids=[local_rank], output_device=local_rank, broadcast_buffers=False)
    else:
        model = unwrapped_model

    # Optimizer - T4 Turing me fused AdamW support nahi hota, fallback rahega
    try:
        optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay, betas=(0.9, 0.95), fused=True)
        if is_main: logger.info("Using fused AdamW")
    except Exception as e:
        if is_main: logger.warning(f"fused=True failed ({e}). Using non-fused AdamW for T4.")
        optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay, betas=(0.9, 0.95))

    # FIX 7: Naya GradScaler API + T4 = fp16 (bf16_supported=False hoga T4 par)
    use_bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported()
    amp_dtype = torch.bfloat16 if use_bf16 else torch.float16
    scaler = GradScaler('cuda', enabled=not use_bf16)
    if is_main:
        logger.info(f"T4 Detected: bf16_supported={use_bf16}, using amp_dtype={amp_dtype}, scaler_enabled={scaler.is_enabled()}")

    api = HfApi()
    start_step = load_checkpoint(unwrapped_model, optimizer, scaler, api, token, "ViuAI/ViuRec", strict=False)

    work_dir = Path("/kaggle/working" if Path("/kaggle/working").exists() else ".")
    shard_pool = ShardPool(api, "ViuAI/viuai-500m-data", work_dir / "shards", token, vocab_size)
    shard_pool.fill(args.ctx_len)

    # Val Data (only main)
    val_tokens = None
    if is_main:
        logger.info("Loading validation data...")
        try:
            val_files = [f for f in api.list_repo_files(repo_id="ViuAI/viuai-500m-data", repo_type="dataset", token=token) if f.startswith("shards/") and f.endswith("_val.bin")]
        except Exception:
            val_files = [f for f in api.list_repo_files(repo_id="ViuAI/viuai-500m-data", repo_type="model", token=token) if f.startswith("shards/") and f.endswith("_val.bin")]
        try:
            if val_files:
                val_p = hf_hub_download(repo_id="ViuAI/viuai-500m-data", repo_type="dataset", filename=val_files[0], local_dir=work_dir, token=token)
                dtype = np.uint16 if vocab_size <= 65535 else np.uint32
                val_arr = np.memmap(val_p, dtype=dtype, mode="r")
                max_toks = 2_000_000 # T4 RAM ke liye 20M se 2M kiya
                val_tokens = np.array(val_arr[:max_toks])
                del val_arr
                # Val file ko delete mat karo taki dobara download na ho - T4 pe time bachega
                # Path(val_p).unlink(missing_ok=True)
        except Exception as e:
            logger.warning(f"Failed to load val data: {e}")

    @torch.no_grad()
    def evaluate():
        if val_tokens is None or len(val_tokens) < args.ctx_len + 1: return 0.0
        unwrapped_model.eval()
        val_loss = 0.0
        val_iters = 10 # T4 pe kam iter
        for _ in range(val_iters):
            i = rng.integers(0, len(val_tokens) - args.ctx_len - 1)
            x = torch.from_numpy(val_tokens[i:i+args.ctx_len].astype(np.int64)).unsqueeze(0).to(device, non_blocking=True)
            y = torch.from_numpy(val_tokens[i+1:i+1+args.ctx_len].astype(np.int64)).unsqueeze(0).to(device, non_blocking=True)
            with autocast('cuda', dtype=amp_dtype):
                _, loss = unwrapped_model(x, y, use_checkpoint=False)
            val_loss += loss.item()
        unwrapped_model.train()
        return val_loss / val_iters

    # Training Loop
    model.train()
    if is_main:
        logger.info(f"Starting training for {args.max_steps} steps from step {start_step}...")
        
    for step in range(start_step, args.max_steps):
        # LR Schedule
        warmup = 1000
        progress = min(1.0, step / max(1, warmup))
        if step < warmup:
            lr = args.lr * progress
        else:
            lr = args.lr * (0.5 * (1.0 + math.cos(math.pi * (step - warmup) / max(1, args.max_steps - warmup))))
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr

        optimizer.zero_grad(set_to_none=True)
        accum_loss = 0.0
        
        t0 = time.time()
        for micro_i in range(args.grad_accum):
            x_np, y_np = shard_pool.sample_batch(args.micro_batch, args.ctx_len, rng)
            x = torch.from_numpy(x_np).to(device, non_blocking=True)
            y = torch.from_numpy(y_np).to(device, non_blocking=True)
            
            with autocast('cuda', dtype=amp_dtype):
                _, loss = model(x, y)
                loss = loss / args.grad_accum
            
            scaler.scale(loss).backward()
            accum_loss += loss.item()
            del x, y, loss

        # Gradient Clipping
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        
        # NaN Check
        has_nan = not math.isfinite(accum_loss)
        if not has_nan:
            for p in model.parameters():
                if p.grad is not None and not torch.isfinite(p.grad).all():
                    has_nan = True
                    break
                
        if has_nan:
            if is_main: logger.warning(f"NaN/Inf detected at step {step}! Skipping.")
            optimizer.zero_grad(set_to_none=True)
            scaler.update() # FIX 8: NaN par bhi scaler update karna zaruri hai warna stuck hoga
            continue
            
        scaler.step(optimizer)
        scaler.update()

        dt = time.time() - t0
        
        if is_main and step % 10 == 0:
            tok_s = (args.micro_batch * args.grad_accum * args.ctx_len * world_size) / dt if dt > 0 else 0
            mem_alloc = torch.cuda.memory_allocated(device) / 1e9 if torch.cuda.is_available() else 0
            mem_res = torch.cuda.memory_reserved(device) / 1e9 if torch.cuda.is_available() else 0
            logger.info(f"step {step:5d}/{args.max_steps} | loss {accum_loss:.4f} | lr {lr:.2e} | {tok_s:.0f} tok/s | mem {mem_alloc:.2f}GB/{mem_res:.2f}GB")

        if step > start_step and step % args.rotate_every == 0:
            shard_pool.rotate_one()
            shard_pool.fill(args.ctx_len)
            if torch.cuda.is_available():
                torch.cuda.empty_cache() # T4 me fragmentation se bachne ke liye

        if step > start_step and step % args.eval_every == 0 and is_main:
            val_loss = evaluate()
            logger.info(f"--- Eval step {step}: val_loss {val_loss:.4f} ---")

        if step > start_step and step % args.save_every == 0:
            if is_ddp: dist.barrier()
            if is_main:
                logger.info(f"Saving checkpoint at step {step}...")
                ckpt_path = Path("checkpoints/latest_chk.pt")
                ckpt_path.parent.mkdir(parents=True, exist_ok=True)
                chk = {
                    "model": unwrapped_model.state_dict(),
                    "optimizer": optimizer.state_dict(),
                    "scaler": scaler.state_dict(),
                    "step": step,
                }
                torch.save(chk, ckpt_path)
                try:
                    api.upload_file(
                        path_or_fileobj=str(ckpt_path),
                        path_in_repo="checkpoints/latest_chk.pt",
                        repo_id="ViuAI/ViuRec",
                        token=token
                    )
                    logger.info("Upload successful.")
                except Exception as e:
                    logger.error(f"Upload failed: {e}")
            if is_ddp: dist.barrier()

    if is_ddp:
        dist.destroy_process_group()

if __name__ == "__main__":
    main()