import json
import os
import time
from pathlib import Path

import torch
from peft import PeftModel
from transformers import AutoModelForImageTextToText, AutoProcessor
from video_inputs import build_video_item, prepare

CACHE = Path("/home/s56fasla/MOTIONX/cache")
RUNS = Path("/home/s56fasla/MOTIONX/runs")
BASE_MODEL = os.environ.get("BASE_MODEL", "Qwen/Qwen3-VL-8B-Instruct")
RUN_PREFIX = os.environ.get("RUN_PREFIX", "skivlm_Qwen3-VL-8B")
TAG = os.environ.get("TAG", "q3")
ARM = os.environ.get("ARM", "ski")
LIMIT = int(os.environ.get("LIMIT", 0))
OUT = Path("/home/s56fasla/MOTIONX/evals")
QUESTION = "Describe what the person is doing in this video."
VIDEO_FPS = 4.0
MAX_FRAMES = 32
MAX_PIXELS = 200704
MAX_NEW_TOKENS = 128

OUT.mkdir(parents=True, exist_ok=True)
device = "cuda"
run_dir = None if ARM == "base" else RUNS / f"{RUN_PREFIX}_{ARM}"

proc_src = run_dir / "processor" if run_dir else BASE_MODEL
processor = AutoProcessor.from_pretrained(proc_src)
model = AutoModelForImageTextToText.from_pretrained(
    BASE_MODEL, dtype=torch.bfloat16, device_map={"": 0})
if len(processor.tokenizer) != model.get_input_embeddings().weight.shape[0]:
    model.resize_token_embeddings(len(processor.tokenizer))
if run_dir:
    model = PeftModel.from_pretrained(model, run_dir / "lora")
model.eval()
print("arm", ARM, "lora", bool(run_dir), flush=True)

items = json.load(open(CACHE / "manifest_val_ok.json"))
if LIMIT:
    items = items[:LIMIT]
print("val clips", len(items), flush=True)

preds = []
t0 = time.time()
for k, m in enumerate(items):
    messages = [{"role": "user", "content": [
        build_video_item(m["video"], VIDEO_FPS, MAX_FRAMES, MAX_PIXELS),
        {"type": "text", "text": QUESTION}]}]
    try:
        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        videos, vkw = prepare(messages, m["video"])
        enc = processor(text=[text], videos=videos, return_tensors="pt", **vkw).to(device)
        with torch.no_grad():
            out = model.generate(**enc, max_new_tokens=MAX_NEW_TOKENS, do_sample=False)
        gen = processor.tokenizer.decode(
            out[0][enc["input_ids"].shape[1]:], skip_special_tokens=True).strip()
    except Exception as e:
        print("[skip]", m["stem"], type(e).__name__, flush=True)
        continue
    preds.append({"stem": m["stem"], "subset": m["subset"], "ref": m["text"], "pred": gen})
    if k % 25 == 0:
        el = time.time() - t0
        print(f"{k}/{len(items)} {el:.0f}s eta {el/max(1,k)*(len(items)-k)/60:.0f}min", flush=True)
        json.dump(preds, open(OUT / f"preds_{TAG}_{ARM}.json", "w"), indent=1)

json.dump(preds, open(OUT / f"preds_{TAG}_{ARM}.json", "w"), indent=1)
print("done", len(preds), "->", OUT / f"preds_{TAG}_{ARM}.json", flush=True)