from __future__ import annotations
import argparse
import json
import logging
import sys
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
import torch
from torch.utils.data import DataLoader
from transformers import AutoTokenizer
from transformers.optimization import Adafactor
from riemann.dataset import PackedTextDataset
from riemann.growth_ladder import get_stage, GrowthStage
from riemann.growth import build_model
from riemann.modeling_riemann import RiemannForCausalLM
from riemann.openpeer_weights import build_stage0_from_openpeer_pretrained
from riemann.device import (
    main_device,
    secondary_device,
    cpu_device,
    configure_cpu_threads,
    describe_devices,
)

logger = logging.getLogger("riemann.pretrain")
from riemann.paths import DATA_DIR, CHECKPOINT_ROOT

TOKENIZER_NAME = "gpt2"


def load_or_build_model(
    stage: GrowthStage, checkpoint_dir: Path, resume: bool, use_pretrained_init: bool
):
    if resume and (checkpoint_dir / "config.json").exists():
        logger.info("Resuming model weights from %s", checkpoint_dir)
        model = RiemannForCausalLM.from_pretrained(str(checkpoint_dir))
    elif use_pretrained_init and stage.name == "stage0_124m":
        logger.info("Building %s from OpenPeerLLM's own published weights", stage.name)
        model = build_stage0_from_openpeer_pretrained(stage.to_riemann_config())
    else:
        logger.info(
            "Building a fresh %s model (no checkpoint found / --resume not given)",
            stage.name,
        )
        model = build_model(stage)
    return model


def load_training_state(checkpoint_dir: Path) -> dict:
    state_path = checkpoint_dir / "training_state.json"
    if state_path.exists():
        return json.loads(state_path.read_text(encoding="utf-8"))
    return {"stage": None, "global_step": 0}


def save_checkpoint(
    model, tokenizer, checkpoint_dir: Path, stage: GrowthStage, global_step: int
) -> None:
    original_device = next(model.parameters()).device
    model.to("cpu")
    checkpoint_dir.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(str(checkpoint_dir))
    tokenizer.save_pretrained(str(checkpoint_dir))
    model.to(original_device)
    state = {"stage": stage.name, "global_step": global_step}
    (checkpoint_dir / "training_state.json").write_text(
        json.dumps(state, indent=2), encoding="utf-8"
    )
    logger.info("Saved checkpoint at step %d -> %s", global_step, checkpoint_dir)


def run_cpu_microstep(cpu_model, batch, grad_accum: int) -> float:
    outputs = cpu_model(input_ids=batch, labels=batch)
    loss = outputs.loss / grad_accum
    loss.backward()
    return outputs.loss.item()


def sync_cpu_mirror(cpu_model, model, cpu_dev) -> None:
    cpu_model.load_state_dict(
        {k: v.detach().to(cpu_dev) for k, v in model.state_dict().items()}
    )
    for p in cpu_model.parameters():
        p.grad = None


def build_multimodal_step(model, adapter, device, secondary, tokenizer, example):
    frames = example["frames"].to(secondary)
    per_frame = adapter.frame_encoder(frames)
    pooled = per_frame.mean(dim=0, keepdim=True).to(device)
    video_tok = adapter.to_lm(pooled)
    prompt_ids = tokenizer.encode(example["prompt"])[:256]
    prompt_ids_t = torch.tensor(prompt_ids, dtype=torch.long, device=device)
    text_embeds = model.word_embeddings(prompt_ids_t).unsqueeze(0)
    inputs_embeds = torch.cat([video_tok.unsqueeze(0), text_embeds], dim=1)
    labels = torch.cat(
        [torch.tensor([-100], dtype=torch.long, device=device), prompt_ids_t]
    ).unsqueeze(0)
    outputs = model(inputs_embeds=inputs_embeds, labels=labels)
    return outputs.loss


def main() -> None:
    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s"
    )
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument(
        "--stage", default="stage0_124m", help="growth ladder stage name"
    )
    parser.add_argument(
        "--resume",
        action="store_true",
        help="continue from this stage's saved checkpoint",
    )
    parser.add_argument(
        "--max-steps",
        type=int,
        default=None,
        help="stop after this many optimizer steps (for smoke tests)",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=None,
        help="override the stage's suggested batch size",
    )
    parser.add_argument(
        "--grad-accum",
        type=int,
        default=None,
        help="override the stage's suggested grad accumulation",
    )
    parser.add_argument(
        "--seq-len",
        type=int,
        default=None,
        help="override the stage's suggested sequence length",
    )
    parser.add_argument(
        "--num-workers",
        type=int,
        default=2,
        help="DataLoader worker processes for tokenization",
    )
    parser.add_argument("--log-every", type=int, default=10)
    parser.add_argument("--save-every", type=int, default=500)
    parser.add_argument(
        "--multimodal",
        action="store_true",
        help="also mix in CodeVideoBench video-conditioned steps",
    )
    parser.add_argument(
        "--multimodal-every",
        type=int,
        default=50,
        help="run one multimodal step every N optimizer steps",
    )
    parser.add_argument(
        "--no-pretrained-init",
        action="store_true",
        help="build stage0 from random weights instead of OpenPeerLLM's published checkpoint",
    )
    parser.add_argument(
        "--no-cpu-accel",
        action="store_true",
        help="disable routing a share of micro-batches to the CPU during training",
    )
    parser.add_argument(
        "--cpu-micro-step-every",
        type=int,
        default=4,
        help="1 out of every N micro-batches in a grad-accum cycle trains on CPU instead of the main GPU",
    )
    parser.add_argument(
        "--cpu-threads",
        type=int,
        default=None,
        help="override CPU thread count (defaults to all logical cores)",
    )
    parser.add_argument("--data-dir", default=str(DATA_DIR))
    parser.add_argument("--checkpoint-root", default=str(CHECKPOINT_ROOT))
    args = parser.parse_args()
    stage = get_stage(args.stage)
    batch_size = args.batch_size or stage.suggested_batch_size
    grad_accum = args.grad_accum or stage.suggested_grad_accum
    seq_len = args.seq_len or stage.suggested_seq_len
    checkpoint_dir = Path(args.checkpoint_root) / stage.name
    logger.info("Devices: %s", json.dumps(describe_devices()))
    if not stage.feasible_on_this_machine:
        logger.warning(
            "%s is marked NOT feasible on this machine. Proceeding anyway since you asked for it, but expect this to be slow or to run out of memory.",
            stage.name,
        )
    device = main_device()
    tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    model = load_or_build_model(
        stage, checkpoint_dir, args.resume, not args.no_pretrained_init
    )
    model.to(device)
    if stage.use_gradient_checkpointing:
        model.gradient_checkpointing_enable()
        model.config.use_cache = False
    model.train()
    cpu_dev = cpu_device()
    cpu_accel = (not args.no_cpu_accel) and device != cpu_dev
    cpu_model = None
    cpu_executor = None
    if cpu_accel:
        configure_cpu_threads(args.cpu_threads)
        cpu_model = type(model)(model.config)
        cpu_model.load_state_dict(
            {k: v.detach().to(cpu_dev) for k, v in model.state_dict().items()}
        )
        cpu_model.to(cpu_dev)
        cpu_model.train()
        cpu_executor = ThreadPoolExecutor(max_workers=1)
        logger.info(
            "CPU acceleration enabled: 1 of every %d micro-batches trains on %s",
            args.cpu_micro_step_every,
            cpu_dev,
        )
    state = load_training_state(checkpoint_dir) if args.resume else {"global_step": 0}
    global_step = state.get("global_step", 0)
    params = list(model.parameters())
    adapter = None
    multimodal_examples = []
    secondary = None
    if args.multimodal:
        from riemann.multimodal import (
            VideoPrefixAdapter,
            build_multimodal_training_examples,
        )

        secondary = secondary_device()
        adapter = VideoPrefixAdapter(lm_n_embd=model.config.hidden_size)
        adapter.frame_encoder.to(secondary)
        adapter.to_lm.to(device)
        params += list(adapter.parameters())
        multimodal_examples = build_multimodal_training_examples()
        logger.info(
            "Multimodal enabled: %d CodeVideoBench examples, frame encoder on %s",
            len(multimodal_examples),
            secondary,
        )
    optimizer = Adafactor(
        params, scale_parameter=True, relative_step=True, warmup_init=True, lr=None
    )
    corpus_path = Path(args.data_dir) / "pretrain_corpus.txt"
    if not corpus_path.exists():
        raise FileNotFoundError(
            f"{corpus_path} not found -- run scripts/prepare_datasets.py first"
        )
    dataset = PackedTextDataset(corpus_path, tokenizer, seq_len)
    loader = DataLoader(dataset, batch_size=batch_size, num_workers=args.num_workers)
    logger.info(
        "Training %s: batch_size=%d grad_accum=%d seq_len=%d device=%s (resuming from step %d)",
        stage.name,
        batch_size,
        grad_accum,
        seq_len,
        device,
        global_step,
    )
    mm_cursor = 0
    accum_loss = 0.0
    cpu_micro_steps = 0
    micro_step = 0
    pending_cpu_future = None
    start_time = time.time()
    optimizer.zero_grad()
    for batch in loader:
        micro_step += 1
        route_to_cpu = cpu_accel and micro_step % args.cpu_micro_step_every == 0
        if route_to_cpu:
            cpu_batch = batch.to(cpu_dev)
            if pending_cpu_future is not None:
                accum_loss += pending_cpu_future.result()
                cpu_micro_steps += 1
            pending_cpu_future = cpu_executor.submit(
                run_cpu_microstep, cpu_model, cpu_batch, grad_accum
            )
        else:
            batch = batch.to(device)
            outputs = model(input_ids=batch, labels=batch)
            loss = outputs.loss / grad_accum
            loss.backward()
            accum_loss += outputs.loss.item()
        if micro_step % grad_accum == 0:
            if pending_cpu_future is not None:
                accum_loss += pending_cpu_future.result()
                cpu_micro_steps += 1
                pending_cpu_future = None
                for p, cp in zip(model.parameters(), cpu_model.parameters()):
                    if cp.grad is None:
                        continue
                    cpu_grad = cp.grad.detach().to(device)
                    p.grad = cpu_grad if p.grad is None else p.grad + cpu_grad
            optimizer.step()
            optimizer.zero_grad()
            if cpu_accel:
                sync_cpu_mirror(cpu_model, model, cpu_dev)
            global_step += 1
            if (
                adapter is not None
                and multimodal_examples
                and (global_step % args.multimodal_every == 0)
            ):
                example = multimodal_examples[mm_cursor % len(multimodal_examples)]
                mm_cursor += 1
                mm_loss = build_multimodal_step(
                    model, adapter, device, secondary, tokenizer, example
                )
                (mm_loss / grad_accum).backward()
                optimizer.step()
                optimizer.zero_grad()
                logger.info(
                    "  [multimodal step on %s] loss=%.4f",
                    example["name"],
                    mm_loss.item(),
                )
            if global_step % args.log_every == 0:
                elapsed = time.time() - start_time
                avg_loss = accum_loss / (args.log_every * grad_accum)
                logger.info(
                    "step %d | loss %.4f | %.1fs elapsed | %d cpu micro-steps",
                    global_step,
                    avg_loss,
                    elapsed,
                    cpu_micro_steps,
                )
                accum_loss = 0.0
                cpu_micro_steps = 0
            if global_step % args.save_every == 0:
                save_checkpoint(model, tokenizer, checkpoint_dir, stage, global_step)
            if args.max_steps is not None and global_step >= args.max_steps:
                break
    if cpu_executor is not None:
        cpu_executor.shutdown(wait=True)
    save_checkpoint(model, tokenizer, checkpoint_dir, stage, global_step)
    logger.info("Done. Final step %d, checkpoint at %s", global_step, checkpoint_dir)


if __name__ == "__main__":
    main()