# ---------- fine_tune_v7_auto_segment.py ----------
# (주식회사엠케이디 | MKD)
# MKDCore Fine-Tuning — Auto-Segments from PKL, Safe Stop/Resume, GPU Monitor, Checkpoints

import os
import sys
import time
import gc
import pickle
import signal
import random
import datetime
import subprocess
import threading

import torch
from torch import nn, optim
from tqdm import tqdm
import psutil

from mkd_core import MKDCore, MyTokenizer, SEQ_LEN


# ============================================================
# CONFIGURATION
# ============================================================
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
BASE_DIR = "/home/mkd-s003/AI/finetune"

TOKENIZER_JSON = os.path.join(BASE_DIR, "prod_tokenizer_v2.json")
TRAIN_PKL = os.path.join(BASE_DIR, "token_cache_train.pkl")
VAL_PKL = os.path.join(BASE_DIR, "token_cache_val.pkl")
MODEL_BASE = os.path.join(BASE_DIR, "mkdcore_v2.pt")

STATE_PATH = os.path.join(BASE_DIR, "mkdcore_state_finetune.pt")
CKPT_DIR = os.path.join(BASE_DIR, "checkpoints_ft")
LOG_REASON = os.path.join(BASE_DIR, "reason_of_interruption.log")
PROGRESS_FILE = os.path.join(BASE_DIR, "segment_progress.pkl")

os.makedirs(CKPT_DIR, exist_ok=True)
torch.cuda.empty_cache()

MAX_TRAIN_TOKENS_SEG = 20_000_000
MAX_VAL_TOKENS_SEG = 5_000_000

EPOCHS = 16
STEPS_PER_EPOCH = 20_000
LR = 2e-5
LR_DECAY = 0.9
GRAD_CLIP = 1.0
SAVE_EVERY = 5_000
PRINT_EVERY = 1_000
RESTART_DELAY = 60

CHUNK_PRINT_EVERY = 1000


# ============================================================
# MONITORING
# ============================================================
def system_status():
    cpu = psutil.cpu_percent(interval=0.1)
    mem = psutil.virtual_memory().percent
    if not torch.cuda.is_available():
        return f"CPU {cpu:.0f}% | RAM {mem:.1f}%"
    mem_alloc = torch.cuda.memory_allocated(0) / 1e9
    total_mem = torch.cuda.get_device_properties(0).total_memory / 1e9
    try:
        util = subprocess.check_output(
            ["nvidia-smi", "--query-gpu=utilization.gpu", "--format=csv,noheader,nounits"]
        ).decode().strip().split("\n")[0]
    except Exception:
        util = "?"
    return f"GPU {mem_alloc:.2f}/{total_mem:.0f} GB | Util {util}% | CPU {cpu:.0f}% | RAM {mem:.1f}%"


def log_reason(reason: str):
    with open(LOG_REASON, "a", encoding="utf-8") as f:
        f.write(f"[{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] {reason}\n")


# ============================================================
# SEGMENT PROGRESS
# ============================================================
def load_segment_progress():
    if not os.path.exists(PROGRESS_FILE):
        return {"train_pos": 0, "val_pos": 0}
    try:
        with open(PROGRESS_FILE, "rb") as f:
            data = pickle.load(f)
        return {
            "train_pos": int(data.get("train_pos", 0)),
            "val_pos": int(data.get("val_pos", 0)),
        }
    except Exception:
        return {"train_pos": 0, "val_pos": 0}


def save_segment_progress(train_pos: int, val_pos: int):
    with open(PROGRESS_FILE, "wb") as f:
        pickle.dump({"train_pos": int(train_pos), "val_pos": int(val_pos)}, f)


# ============================================================
# STREAM-LOAD TOKEN SEGMENT
# ============================================================
def load_tokens_segment(pkl_path: str, start_offset_tokens: int, max_tokens: int, label: str):
    tokens = []
    loaded_tokens = 0
    consumed = 0
    chunks = 0

    print(f"\nLoading {os.path.basename(pkl_path)} segment | start={start_offset_tokens:,} tokens | limit={max_tokens:,} tokens")
    with open(pkl_path, "rb") as f:
        while loaded_tokens < max_tokens:
            try:
                seq = pickle.load(f)
            except EOFError:
                break

            if not isinstance(seq, list) or len(seq) == 0:
                continue

            seq_len = len(seq)
            if consumed + seq_len <= start_offset_tokens:
                consumed += seq_len
                chunks += 1
                if chunks % CHUNK_PRINT_EVERY == 0:
                    print(f"  • {chunks:,} chunks scanned | skipped={consumed:,} | loaded={loaded_tokens:,} | {label}")
                continue

            in_seq_start = max(0, start_offset_tokens - consumed)
            take = seq_len - in_seq_start
            take = min(take, max_tokens - loaded_tokens)

            if take > 0:
                tokens.extend(seq[in_seq_start:in_seq_start + take])
                loaded_tokens += take

            consumed += seq_len
            chunks += 1

            if chunks % CHUNK_PRINT_EVERY == 0:
                print(f"  • {chunks:,} chunks | loaded={loaded_tokens:,} | RAM {psutil.virtual_memory().percent:.1f}% | {label}")

            if loaded_tokens >= max_tokens:
                break

    new_offset = start_offset_tokens + loaded_tokens
    print(f"Loaded {loaded_tokens:,} tokens from {os.path.basename(pkl_path)} | new_offset={new_offset:,} | {label}")
    if loaded_tokens == 0:
        print(f"Warning: no tokens loaded from {pkl_path} at start offset {start_offset_tokens:,}. Possibly at file end.")
    return torch.tensor(tokens, dtype=torch.long), new_offset


# ============================================================
# BATCHING
# ============================================================
def get_batch(source: torch.Tensor, seq_len: int):
    if len(source) <= seq_len + 1:
        raise ValueError("Dataset segment too small for the configured sequence length.")
    i = random.randint(0, len(source) - seq_len - 1)
    x = source[i:i + seq_len].unsqueeze(0).to(DEVICE)
    y = source[i + 1:i + seq_len + 1].unsqueeze(0).to(DEVICE)
    return x, y


@torch.no_grad()
def evaluate_loss(model, loss_fn, vocab_size, val_data):
    model.eval()
    total = 0.0
    iters = 25
    for _ in range(iters):
        x, y = get_batch(val_data, SEQ_LEN)
        with torch.amp.autocast("cuda", enabled=torch.cuda.is_available()):
            logits = model(x)
            loss = loss_fn(logits.view(-1, vocab_size), y.view(-1))
        total += loss.item()
    return total / iters


# ============================================================
# CTRL+C SAFE HANDLER
# ============================================================
def save_state_on_interrupt(model, opt, scaler):
    try:
        torch.save({
            "model": model.state_dict(),
            "optimizer": opt.state_dict(),
            "scaler": scaler.state_dict(),
            "timestamp": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        }, STATE_PATH)
        print(f"\nProgress saved → {STATE_PATH}")
    except Exception as e:
        print(f"Save failed: {e}")


def ctrl_c_handler_factory(model, opt, scaler):
    def handle_ctrl_c(sig, frame):
        print("\n[CTRL+C] Training interrupted. Saving state...")
        save_state_on_interrupt(model, opt, scaler)
        print("[1] Save & Stop")
        print("[2] Save & Restart in 60s")
        print(f"(No input → auto restart in {RESTART_DELAY}s)")

        choice = {"val": None}
        answered = threading.Event()

        def _read_choice():
            try:
                choice["val"] = input("Choose (1/2): ").strip()
            except Exception:
                choice["val"] = None
            finally:
                answered.set()

        t = threading.Thread(target=_read_choice, daemon=True)
        t.start()
        for sec in range(RESTART_DELAY, 0, -1):
            if answered.is_set():
                break
            print(f"\rRestarting in {sec:02d}s | {system_status()}", end="", flush=True)
            time.sleep(1)
        print()

        if not answered.is_set() or choice["val"] == "2":
            log_reason("Manual or auto restart")
            print("Restarting now...")
            os.execv(sys.executable, [sys.executable] + sys.argv)
        else:
            log_reason("Stopped safely by user")
            print("Training stopped safely.")
            sys.exit(0)

    return handle_ctrl_c


# ============================================================
# MAIN
# ============================================================
def main():
    print(f"Device: {DEVICE}")

    tok = MyTokenizer(TOKENIZER_JSON)
    vocab_size = tok.vocab_size
    print(f"Vocab size: {vocab_size}")

    model = MKDCore(vocab_size).to(DEVICE)
    state = torch.load(MODEL_BASE, map_location=DEVICE)
    model.load_state_dict(state, strict=False)
    print(f"Starting from: {MODEL_BASE}")
    print("Model loaded successfully.")

    opt = optim.AdamW(model.parameters(), lr=LR, betas=(0.9, 0.95))
    loss_fn = nn.CrossEntropyLoss()
    scaler = torch.amp.GradScaler("cuda", enabled=torch.cuda.is_available())

    signal.signal(signal.SIGINT, ctrl_c_handler_factory(model, opt, scaler))

    prog = load_segment_progress()
    start_train = prog["train_pos"]
    start_val = prog["val_pos"]

    train_data, new_train_pos = load_tokens_segment(TRAIN_PKL, start_train, MAX_TRAIN_TOKENS_SEG, "train")
    val_data, new_val_pos = load_tokens_segment(VAL_PKL, start_val, MAX_VAL_TOKENS_SEG, "val")

    if len(train_data) == 0 or len(val_data) == 0:
        print("No tokens available for this segment. Consider increasing PKL caches or resetting segment progress.")
        print(f"Current offsets → train_pos={start_train:,}, val_pos={start_val:,}")
        return

    print(f"Dataset segment ready: {len(train_data):,} train | {len(val_data):,} val tokens")
    save_segment_progress(new_train_pos, new_val_pos)
    print(f"Advanced segment offsets saved → train_pos={new_train_pos:,} | val_pos={new_val_pos:,}")

    model_save = os.path.join(
        BASE_DIR,
        f"mkdcore_v2_segment_finetuned_{datetime.datetime.now().strftime('%Y%m%d_%H%M')}.pt"
    )
    print(f"\nFine-tuning started → {model_save}\n")

    train_losses, val_losses = [], []

    for epoch in range(EPOCHS):
        model.train()
        total_loss = 0.0
        start_time = time.time()

        with tqdm(total=STEPS_PER_EPOCH, desc=f"Epoch {epoch + 1}/{EPOCHS}", ncols=120) as pbar:
            for step in range(STEPS_PER_EPOCH):
                x, y = get_batch(train_data, SEQ_LEN)

                with torch.amp.autocast("cuda", enabled=torch.cuda.is_available()):
                    logits = model(x)
                    loss = loss_fn(logits.view(-1, vocab_size), y.view(-1))

                scaler.scale(loss).backward()
                scaler.unscale_(opt)
                torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP)
                scaler.step(opt)
                scaler.update()
                opt.zero_grad(set_to_none=True)

                total_loss += loss.item()
                avg_loss = total_loss / (step + 1)
                pbar.set_postfix({"Loss": f"{avg_loss:.4f}", "GPU": system_status()})
                pbar.update(1)

                if (step + 1) % PRINT_EVERY == 0:
                    elapsed = (time.time() - start_time) / 60
                    print(f"Step {step + 1} | Loss {avg_loss:.4f} | {system_status()} | Elapsed {elapsed:.1f}m")

                if (step + 1) % SAVE_EVERY == 0:
                    ckpt = os.path.join(CKPT_DIR, f"mkdcore_seg_ckpt_e{epoch + 1}_s{step + 1}.pt")
                    torch.save(model.state_dict(), ckpt)
                    print(f"Checkpoint saved → {ckpt}")

                if (step + 1) % 2000 == 0:
                    gc.collect()
                    if torch.cuda.is_available():
                        torch.cuda.empty_cache()

        val_loss = evaluate_loss(model, loss_fn, vocab_size, val_data)
        train_losses.append(total_loss / STEPS_PER_EPOCH)
        val_losses.append(val_loss)
        torch.save(model.state_dict(), model_save)
        print(f"Epoch {epoch + 1} | Train {train_losses[-1]:.4f} | Val {val_loss:.4f} | Saved → {model_save}")

        for g in opt.param_groups:
            g["lr"] *= LR_DECAY

    print("\nFine-tuning completed successfully.")
    print(f"Final model → {model_save}")
    print(f"Next run will start at offsets → train_pos={new_train_pos:,} | val_pos={new_val_pos:,}")
    print("You can re-run this same script to consume the next segment automatically.")


# ============================================================
# AUTO CONTINUOUS FINE-TUNING LOOP
# ============================================================
if __name__ == "__main__":
    rounds_completed = 0
    while True:
        try:
            main()
            rounds_completed += 1
            print(f"\nSegment {rounds_completed} finished successfully.")

            if os.path.exists(PROGRESS_FILE):
                with open(PROGRESS_FILE, "rb") as f:
                    prog = pickle.load(f)
                train_pos = prog.get("train_pos", 0)
                val_pos = prog.get("val_pos", 0)
            else:
                train_pos = val_pos = 0

            if train_pos == 0 and val_pos == 0:
                print("No progress file found — exiting.")
                break
            train_size = os.path.getsize(TRAIN_PKL)
            val_size = os.path.getsize(VAL_PKL)
            if train_pos >= train_size or val_pos >= val_size:
                print("All segments processed — training completed.")
                break

            print(f"Next run will continue from train_pos={train_pos:,}, val_pos={val_pos:,}")
            print(f"Waiting {RESTART_DELAY}s before starting next segment...")
            for t in range(RESTART_DELAY, 0, -1):
                print(f"\rRestarting in {t:02d}s | {system_status()}", end="", flush=True)
                time.sleep(1)
            print("\nStarting next segment now...\n")

        except KeyboardInterrupt:
            print("\nTraining stopped manually by user. Progress saved.")
            sys.exit(0)
        except Exception as e:
            print(f"\nError occurred: {e}")
            print(f"Restarting automatically in {RESTART_DELAY}s...")
            for t in range(RESTART_DELAY, 0, -1):
                print(f"\rRestart in {t:02d}s | {system_status()}", end="", flush=True)
                time.sleep(1)
            print("\nRestarting now...\n")